Remove custom query job async timeout logic (#1109) (#1111)

* use dynamic schema in test_grant_access_to.py

* use dynamic schema in test_grant_access_to.py

* revert setup

* use dbt-common main to test against

* emove custom query job async timeout logic

* remove unneeded unit test

---------

Co-authored-by: Mike Alfare <[email protected]>
(cherry picked from commit 5501cd3)

.changes/unreleased/Fixes-20240219-103324.yaml

@@ -0,0 +1,6 @@
+kind: Fixes
+body: Remove custom query job async timeout logic as it has been fixed in bigquery-python
+time: 2024-02-19T10:33:24.3385-08:00
+custom:
+  Author: colin-rogers-dbt
+  Issue: "1081"

dbt/adapters/bigquery/connections.py

@@ -1,5 +1,3 @@
-import asyncio
-import functools
 import json
 import re
 from contextlib import contextmanager
@@ -740,25 +738,7 @@ def _query_and_results(
                 self._bq_job_link(query_job.location, query_job.project, query_job.job_id)
             )
 
-        # 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)
+        iterator = query_job.result(max_results=limit, timeout=job_execution_timeout)
         return query_job, iterator
 
     def _retry_and_handle(self, msg, conn, fn):

tests/functional/test_job_timeout.py

@@ -59,4 +59,5 @@ def profiles_config_update(self, dbt_profile_target):
 
     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
+        expected_error = f"Operation did not complete within the designated timeout of {_SHORT_TIMEOUT} seconds."
+        assert expected_error in result[0].message

tests/unit/test_bigquery_connection_manager.py

@@ -1,6 +1,4 @@
-import time
 import json
-import pytest
 import unittest
 from contextlib import contextmanager
 from requests.exceptions import ConnectionError
@@ -122,26 +120,6 @@ def test_query_and_results(self, mock_bq):
             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)
         args, kwargs = self.mock_client.copy_table.call_args