From b62ce7469a47de4e1f526fc551817b3e0919ddb2 Mon Sep 17 00:00:00 2001 From: Colin Rogers <111200756+colin-rogers-dbt@users.noreply.github.com> Date: Fri, 19 Apr 2024 09:32:20 -0700 Subject: [PATCH] Preserve nicely formatted timeout exception (#1187) * use dynamic schema in test_grant_access_to.py * use dynamic schema in test_grant_access_to.py * revert setup * Add descriptive error message on query timeout * Add descriptive error message on query timeout --- dbt/adapters/bigquery/connections.py | 10 +++++++--- tests/unit/test_bigquery_adapter.py | 2 +- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py index 09558a14a..f96bc1381 100644 --- a/dbt/adapters/bigquery/connections.py +++ b/dbt/adapters/bigquery/connections.py @@ -1,3 +1,4 @@ +from concurrent.futures import TimeoutError import json import re from contextlib import contextmanager @@ -736,9 +737,12 @@ def _query_and_results( logger.debug( self._bq_job_link(query_job.location, query_job.project, query_job.job_id) ) - - iterator = query_job.result(max_results=limit, timeout=job_execution_timeout) - return query_job, iterator + try: + iterator = query_job.result(max_results=limit, timeout=job_execution_timeout) + return query_job, iterator + except TimeoutError: + exc = f"Operation did not complete within the designated timeout of {job_execution_timeout} seconds." + raise TimeoutError(exc) def _retry_and_handle(self, msg, conn, fn): """retry a function call within the context of exception_handler.""" diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py index 34abd0caf..248e8e2be 100644 --- a/tests/unit/test_bigquery_adapter.py +++ b/tests/unit/test_bigquery_adapter.py @@ -20,7 +20,7 @@ from dbt.adapters.bigquery.connections import _sanitize_label, _VALIDATE_LABEL_LENGTH_LIMIT from dbt_common.clients import agate_helper import dbt_common.exceptions -from dbt.context.manifest import generate_query_header_context +from dbt.context.query_header import generate_query_header_context from dbt.contracts.files import FileHash from dbt.contracts.graph.manifest import ManifestStateCheck from dbt.logger import GLOBAL_LOGGER as logger # noqa