Repository: dbt-labs/dbt-bigquery (python)

Issue:
[CT-881] refactor retries
We currently have a non-trivial amount of code to support our homegrown retry strategy via `_retry_and_handle`. This in inferior to using Google's provided and configurable `Retry()` decorator.

This is an extension to the work started in #230.

Goals:
- Make sure all network calls have the `Retry()` decorator.
- If the user provides configuration for the job_* retry/timeout options, plumb them into the Retry() decorator in the appropriate ways.
- Delete homegrown retry code

Context:
It is reasonable that a user would want a very quick timeout if a job fails to create/submit, and a much much _much_ longer timeout for a job that's actually running (e.g. models that take >1 hr to run).

Pull request description:
Timeout BQ queries 
Time out and cancel a model execution if it exceeds the provided number of seconds. 

Currently the bigquery python client does not provide us with a mechanism to specify a timeout on the database side. As a workaround until they do we can do the next best thing: timeout the execution, cancel the query and raise an exception. 

Note we will only use this new logic if a user specifies the `job_execution_timeout` parameter in their profile. If it is not set we continue using the same behavior as today.

resolves #231
[docs](https://github.com/dbt-labs/docs.getdbt.com/issues/new/choose) dbt-labs/docs.getdbt.com/#

<!---
  Include the number of the issue addressed by this PR above if applicable.
  PRs for code changes without an associated issue *will not be merged*.
  See CONTRIBUTING.md for more information.

  Include the number of the docs issue that was opened for this PR. If
  this change has no user-facing implications, "N/A" suffices instead. New
  docs tickets can be created by clicking the link above or by going to
  https://github.com/dbt-labs/docs.getdbt.com/issues/new/choose.
-->

### Problem

<!---
  Describe the problem this PR is solving. What is the application state
  before this PR is merged?
-->

### Solution

<!---
  Describe the way this PR solves the above problem. Add as much detail as you
  can to help reviewers understand your changes. Include any alternatives and
  tradeoffs you considered.
-->

### Checklist

- [ ] I have read [the contributing guide](https://github.com/dbt-labs/dbt-core/blob/main/CONTRIBUTING.md) and understand what's expected of me
- [ ] I have run this code in development and it appears to resolve the stated issue
- [ ] This PR includes tests, or tests are not required/relevant for this PR
- [ ] This PR has no interface changes (e.g. macros, cli, logs, json artifacts, config files, adapter interface, etc) or this PR has already received feedback and approval from Product or DX


Patch:
diff --git a/.changes/unreleased/Fixes-20230829-162111.yaml b/.changes/unreleased/Fixes-20230829-162111.yaml
new file mode 100644
index 0000000..5d34acd
--- /dev/null
+++ b/.changes/unreleased/Fixes-20230829-162111.yaml
@@ -0,0 +1,6 @@
+kind: Fixes
+body: Time out queries if user supplies `job_execution_timeout`
+time: 2023-08-29T16:21:11.69291-07:00
+custom:
+  Author: colin-rogers-dbt McKnight-42
+  Issue: "231"
diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py
index 0d3c52c..c136042 100644
--- a/dbt/adapters/bigquery/connections.py
+++ b/dbt/adapters/bigquery/connections.py
@@ -1,7 +1,10 @@
+import asyncio
+import functools
 import json
 import re
 from contextlib import contextmanager
 from dataclasses import dataclass, field
+
 from mashumaro.helper import pass_through
 
 from functools import lru_cache
@@ -710,7 +713,6 @@ class BigQueryConnectionManager(BaseConnectionManager):
         # Cannot reuse job_config if destination is set and ddl is used
         job_config = google.cloud.bigquery.QueryJobConfig(**job_params)
         query_job = client.query(query=sql, job_config=job_config, timeout=job_creation_timeout)
-
         if (
             query_job.location is not None
             and query_job.job_id is not None
@@ -720,8 +722,25 @@ class BigQueryConnectionManager(BaseConnectionManager):
                 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)
+        # only use async logic if user specifies a timeout
+        if job_execution_timeout:
+            loop = asyncio.new_event_loop()
+            future_iterator = asyncio.wait_for(
+                loop.run_in_executor(None, functools.partial(query_job.result, max_results=limit)),
+                timeout=job_execution_timeout,
+            )
 
+            try:
+                iterator = loop.run_until_complete(future_iterator)
+            except asyncio.TimeoutError:
+                query_job.cancel()
+                raise DbtRuntimeError(
+                    f"Query exceeded configured timeout of {job_execution_timeout}s"
+                )
+            finally:
+                loop.close()
+        else:
+            iterator = query_job.result(max_results=limit)
         return query_job, iterator
 
     def _retry_and_handle(self, msg, conn, fn):


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 0ba0091..78f3d82 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -12,7 +12,7 @@ def pytest_addoption(parser):
     parser.addoption("--profile", action="store", default="oauth", type=str)
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="class")
 def dbt_profile_target(request):
     profile_type = request.config.getoption("--profile")
     if profile_type == "oauth":
diff --git a/tests/functional/test_job_timeout.py b/tests/functional/test_job_timeout.py
new file mode 100644
index 0000000..be559e8
--- /dev/null
+++ b/tests/functional/test_job_timeout.py
@@ -0,0 +1,62 @@
+import pytest
+
+from dbt.tests.util import run_dbt
+
+_REASONABLE_TIMEOUT = 300
+_SHORT_TIMEOUT = 1
+
+_LONG_RUNNING_MODEL_SQL = """
+    {{ config(materialized='table') }}
+    with array_1 as (
+    select generated_ids from UNNEST(GENERATE_ARRAY(1, 200000)) AS generated_ids
+    ),
+    array_2 as (
+    select generated_ids from UNNEST(GENERATE_ARRAY(2, 200000)) AS generated_ids
+    )
+
+    SELECT array_1.generated_ids
+    FROM array_1
+    LEFT JOIN array_1 as jnd on 1=1
+    LEFT JOIN array_2 as jnd2 on 1=1
+    LEFT JOIN array_1 as jnd3 on jnd3.generated_ids >= jnd2.generated_ids
+"""
+
+_SHORT_RUNNING_QUERY = """
+    SELECT 1 as id
+    """
+
+
+class TestSuccessfulJobRun:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "model.sql": _SHORT_RUNNING_QUERY,
+        }
+
+    @pytest.fixture(scope="class")
+    def profiles_config_update(self, dbt_profile_target):
+        outputs = {"default": dbt_profile_target}
+        outputs["default"]["job_execution_timeout_seconds"] = _REASONABLE_TIMEOUT
+        return {"test": {"outputs": outputs, "target": "default"}}
+
+    def test_bigquery_job_run_succeeds_within_timeout(self, project):
+        result = run_dbt()
+        assert len(result) == 1
+
+
+class TestJobTimeout:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "model.sql": _LONG_RUNNING_MODEL_SQL,
+        }
+
+    @pytest.fixture(scope="class")
+    def profiles_config_update(self, dbt_profile_target):
+        outputs = {"default": dbt_profile_target}
+        outputs["default"]["job_execution_timeout_seconds"] = _SHORT_TIMEOUT
+        return {"test": {"outputs": outputs, "target": "default"}}
+
+    def test_job_timeout(self, project):
+        result = run_dbt(["run"], expect_pass=False)  # project setup will fail
+        assert f"Query exceeded configured timeout of {_SHORT_TIMEOUT}s" in result[0].message
diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py
index bb98db8..10cb3f5 100644
--- a/tests/unit/test_bigquery_adapter.py
+++ b/tests/unit/test_bigquery_adapter.py
@@ -1,3 +1,5 @@
+import time
+
 import agate
 import decimal
 import json
@@ -634,18 +636,39 @@ class TestBigQueryConnectionManager(unittest.TestCase):
 
     @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery")
     def test_query_and_results(self, mock_bq):
+        self.mock_client.query = Mock(return_value=Mock(state="DONE"))
         self.connections._query_and_results(
             self.mock_client,
             "sql",
             {"job_param_1": "blah"},
             job_creation_timeout=15,
-            job_execution_timeout=100,
+            job_execution_timeout=3,
+        )
+
+        mock_bq.QueryJobConfig.assert_called_once()
+        self.mock_client.query.assert_called_once_with(
+            query="sql", job_config=mock_bq.QueryJobConfig(), timeout=15
         )
 
+    @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery")
+    def test_query_and_results_timeout(self, mock_bq):
+        self.mock_client.query = Mock(
+            return_value=Mock(result=lambda *args, **kwargs: time.sleep(4))
+        )
+        with pytest.raises(dbt.exceptions.DbtRuntimeError) as exc:
+            self.connections._query_and_results(
+                self.mock_client,
+                "sql",
+                {"job_param_1": "blah"},
+                job_creation_timeout=15,
+                job_execution_timeout=1,
+            )
+
         mock_bq.QueryJobConfig.assert_called_once()
         self.mock_client.query.assert_called_once_with(
             query="sql", job_config=mock_bq.QueryJobConfig(), timeout=15
         )
+        assert "Query exceeded configured timeout of 1s" in str(exc.value)
 
     def test_copy_bq_table_appends(self):
         self._copy_table(write_disposition=dbt.adapters.bigquery.impl.WRITE_APPEND)