test(bigquery): refactor request interceptor to resolve socket leaks in system tests

Author: chalmerlowe

packages/google-cloud-bigquery/google/cloud/bigquery/__init__.py

@@ -123,7 +123,7 @@
 except ImportError:
     bigquery_magics = None
 
-if sys.version_info < (3, 10):
+if sys.version_info < (3, 10):  # pragma: NO COVER
     warnings.warn(
         "The python-bigquery library no longer supports Python <= 3.9. "
         f"Your Python version is {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}. We "

packages/google-cloud-bigquery/google/cloud/bigquery/table.py

@@ -2353,7 +2353,9 @@ def to_arrow(
                 progress_bar.close()
         finally:
             if owns_bqstorage_client:
-                bqstorage_client._transport.close()
+                # mypy: bqstorage_client is guaranteed to be not None when owns_bqstorage_client is True,
+                # but mypy cannot infer this correlation. We ignore the union-attr error here.
+                bqstorage_client._transport.close()  # type: ignore[union-attr]
 
         if record_batches and bqstorage_client is not None:
             return pyarrow.Table.from_batches(record_batches)

packages/google-cloud-bigquery/tests/system/helpers.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
 import datetime
 import decimal
 import uuid
@@ -104,3 +105,30 @@ def _rate_limit_exceeded(forbidden):
     google.api_core.exceptions.Forbidden,
     error_predicate=_rate_limit_exceeded,
 )
+
+
+@contextlib.contextmanager
+def patch_tracked_requests():
+    """Context manager to patch google-auth requests and track/close their HTTP sessions.
+
+    This prevents socket leaks in system tests that use Workload Identity or metadata server auth.
+    """
+    import contextlib
+    import google.auth.transport.requests
+
+    original_init = google.auth.transport.requests.Request.__init__
+    tracked_requests = []
+
+    def patched_init(self, session=None):
+        original_init(self, session=session)
+        tracked_requests.append(self)
+
+    google.auth.transport.requests.Request.__init__ = patched_init
+    try:
+        yield tracked_requests
+    finally:
+        google.auth.transport.requests.Request.__init__ = original_init
+        for req in tracked_requests:
+            if hasattr(req, "session") and req.session is not None:
+                req.session.close()
+

packages/google-cloud-bigquery/tests/system/test_client.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import base64
+import contextlib
+import google.auth.transport.requests
 import copy
 import csv
 import datetime
@@ -155,6 +157,8 @@ def _load_json_schema(filename="schema.json"):
         return _parse_schema_resource(json.load(schema_file))
 
 
+
+
 class Config(object):
     """Run-time configuration to be modified at set-up.
 
@@ -234,23 +238,29 @@ def _create_bucket(self, bucket_name, location=None):
 
     def test_close_releases_open_sockets(self):
         current_process = psutil.Process()
-        conn_count_start = len(current_process.net_connections())
+        conn_start = current_process.net_connections()
+        conn_count_start = len(conn_start)
+
+        with helpers.patch_tracked_requests():
+            client = Config.CLIENT
+            client.query(
+                """
+                SELECT
+                    source_year AS year, COUNT(is_male) AS birth_count
+                FROM `bigquery-public-data.samples.natality`
+                GROUP BY year
+                ORDER BY year DESC
+                LIMIT 15
+                """
+            )
 
-        client = Config.CLIENT
-        client.query(
-            """
-            SELECT
-                source_year AS year, COUNT(is_male) AS birth_count
-            FROM `bigquery-public-data.samples.natality`
-            GROUP BY year
-            ORDER BY year DESC
-            LIMIT 15
-            """
-        )
+            client.close()
 
-        client.close()
+        import gc
 
-        conn_count_end = len(current_process.net_connections())
+        gc.collect()
+        conn_end = current_process.net_connections()
+        conn_count_end = len(conn_end)
         self.assertLessEqual(conn_count_end, conn_count_start)
 
     def test_create_dataset(self):
@@ -2174,25 +2184,31 @@ def test_dbapi_dry_run_query(self):
     def test_dbapi_connection_does_not_leak_sockets(self):
         pytest.importorskip("google.cloud.bigquery_storage")
         current_process = psutil.Process()
-        conn_count_start = len(current_process.net_connections())
-
-        # Provide no explicit clients, so that the connection will create and own them.
-        connection = dbapi.connect()
-        cursor = connection.cursor()
-
-        cursor.execute(
+        conn_start = current_process.net_connections()
+        conn_count_start = len(conn_start)
+
+        with helpers.patch_tracked_requests():
+            # Provide no explicit clients, so that the connection will create and own them.
+            connection = dbapi.connect()
+            cursor = connection.cursor()
+
+            cursor.execute(
+                """
+                SELECT id, `by`, timestamp
+                FROM `bigquery-public-data.hacker_news.full`
+                ORDER BY `id` ASC
+                LIMIT 100000
             """
-            SELECT id, `by`, timestamp
-            FROM `bigquery-public-data.hacker_news.full`
-            ORDER BY `id` ASC
-            LIMIT 100000
-        """
-        )
-        rows = cursor.fetchall()
-        self.assertEqual(len(rows), 100000)
+            )
+            rows = cursor.fetchall()
+            self.assertEqual(len(rows), 100000)
+
+            connection.close()
+        import gc
 
-        connection.close()
-        conn_count_end = len(current_process.net_connections())
+        gc.collect()
+        conn_end = current_process.net_connections()
+        conn_count_end = len(conn_end)
         self.assertLessEqual(conn_count_end, conn_count_start)
 
     def _load_table_for_dml(self, rows, dataset_id, table_id):

packages/google-cloud-bigquery/tests/system/test_magics.py

@@ -14,11 +14,15 @@
 
 """System tests for Jupyter/IPython connector."""
 
+import contextlib
 import re
 
+import google.auth.transport.requests
+
 import pytest
 import psutil
 
+from . import helpers
 
 IPython = pytest.importorskip("IPython")
 io = pytest.importorskip("IPython.utils.io")
@@ -45,30 +49,35 @@ def ipython_interactive(ipython):
         yield ipython
 
 
+
+
 def test_bigquery_magic(ipython_interactive):
     ip = IPython.get_ipython()
     current_process = psutil.Process()
-    conn_count_start = len(current_process.net_connections())
-
-    # Deprecated, but should still work in google-cloud-bigquery 3.x.
-    with pytest.warns(FutureWarning, match="bigquery_magics"):
-        ip.extension_manager.load_extension("google.cloud.bigquery")
-
-    sql = """
-        SELECT
-            CONCAT(
-            'https://stackoverflow.com/questions/',
-            CAST(id as STRING)) as url,
-            view_count
-        FROM `bigquery-public-data.stackoverflow.posts_questions`
-        WHERE tags like '%google-bigquery%'
-        ORDER BY view_count DESC
-        LIMIT 10
-    """
-    with io.capture_output() as captured:
-        result = ip.run_cell_magic("bigquery", "--use_rest_api", sql)
-
-    conn_count_end = len(current_process.net_connections())
+    conn_start = current_process.net_connections()
+    conn_count_start = len(conn_start)
+
+    with helpers.patch_tracked_requests():
+        # Deprecated, but should still work in google-cloud-bigquery 3.x.
+        with pytest.warns(FutureWarning, match="bigquery_magics"):
+            ip.extension_manager.load_extension("google.cloud.bigquery")
+
+        sql = """
+            SELECT
+                CONCAT(
+                'https://stackoverflow.com/questions/',
+                CAST(id as STRING)) as url,
+                view_count
+            FROM `bigquery-public-data.stackoverflow.posts_questions`
+            WHERE tags like '%google-bigquery%'
+            ORDER BY view_count DESC
+            LIMIT 10
+        """
+        with io.capture_output() as captured:
+            result = ip.run_cell_magic("bigquery", "--use_rest_api", sql)
+
+    conn_end = current_process.net_connections()
+    conn_count_end = len(conn_end)
 
     lines = re.split("\n|\r", captured.stdout)
     # Removes blanks & terminal code (result of display clearing)
@@ -83,3 +92,4 @@ def test_bigquery_magic(ipython_interactive):
     # than expected when running system tests on Kokoro, thus using the <= assertion.
     # That's still fine, however, since the sockets are apparently not leaked.
     assert conn_count_end <= conn_count_start  # system resources are released
+

packages/google-cloud-bigquery/tests/unit/test_magics.py

@@ -45,7 +45,7 @@
 
 @pytest.fixture()
 def use_local_magics_context(monkeypatch):
-    if magics is not None:
+    if magics is not None:  # pragma: NO COVER
         local_context = magics.Context()
         local_context._project = "unit-test-project"
         mock_credentials = mock.create_autospec(
@@ -2195,13 +2195,10 @@ def test_bigquery_magic_create_dataset_fails(monkeypatch):
 
 
 @pytest.mark.usefixtures("ipython_interactive")
-def test_bigquery_magic_with_location(monkeypatch):
+def test_bigquery_magic_with_location(monkeypatch, use_local_magics_context):
     ip = IPython.get_ipython()
     monkeypatch.setattr(bigquery, "bigquery_magics", None)
     bigquery.load_ipython_extension(ip)
-    magics.context.credentials = mock.create_autospec(
-        google.auth.credentials.Credentials, instance=True
-    )
 
     run_query_patch = mock.patch(
         "google.cloud.bigquery.magics.magics._run_query", autospec=True

packages/google-cloud-bigquery/tests/unit/test_table.py

@@ -79,6 +79,82 @@ def test_ctor_with_key(self):
         self.assertEqual(encryption_config.kms_key_name, self.KMS_KEY_NAME)
 
 
+class TestPropertyGraphReference(unittest.TestCase):
+    PROJECT = "my-project"
+    DATASET_ID = "my_dataset"
+    PROPERTY_GRAPH_ID = "my_pg"
+
+    def _get_target_class(self):
+        from google.cloud.bigquery.table import PropertyGraphReference
+
+        return PropertyGraphReference
+
+    def _make_one(self, *args, **kw):
+        return self._get_target_class()(*args, **kw)
+
+    def test_ctor(self):
+        dataset_ref = DatasetReference(self.PROJECT, self.DATASET_ID)
+        ref = self._make_one(dataset_ref, self.PROPERTY_GRAPH_ID)
+        self.assertEqual(ref.project, self.PROJECT)
+        self.assertEqual(ref.dataset_id, self.DATASET_ID)
+        self.assertEqual(ref.property_graph_id, self.PROPERTY_GRAPH_ID)
+
+    def test_from_api_repr(self):
+        resource = {
+            "projectId": self.PROJECT,
+            "datasetId": self.DATASET_ID,
+            "propertyGraphId": self.PROPERTY_GRAPH_ID,
+        }
+        ref = self._get_target_class().from_api_repr(resource)
+        self.assertEqual(ref.project, self.PROJECT)
+        self.assertEqual(ref.dataset_id, self.DATASET_ID)
+        self.assertEqual(ref.property_graph_id, self.PROPERTY_GRAPH_ID)
+
+    def test_to_api_repr(self):
+        dataset_ref = DatasetReference(self.PROJECT, self.DATASET_ID)
+        ref = self._make_one(dataset_ref, self.PROPERTY_GRAPH_ID)
+        resource = ref.to_api_repr()
+        expected = {
+            "projectId": self.PROJECT,
+            "datasetId": self.DATASET_ID,
+            "propertyGraphId": self.PROPERTY_GRAPH_ID,
+        }
+        self.assertEqual(resource, expected)
+
+    def test___str__(self):
+        dataset_ref = DatasetReference(self.PROJECT, self.DATASET_ID)
+        ref = self._make_one(dataset_ref, self.PROPERTY_GRAPH_ID)
+        self.assertEqual(
+            str(ref), f"{self.PROJECT}.{self.DATASET_ID}.{self.PROPERTY_GRAPH_ID}"
+        )
+
+    def test___repr__(self):
+        dataset_ref = DatasetReference(self.PROJECT, self.DATASET_ID)
+        ref = self._make_one(dataset_ref, self.PROPERTY_GRAPH_ID)
+        expected = (
+            f"PropertyGraphReference({dataset_ref!r}, '{self.PROPERTY_GRAPH_ID}')"
+        )
+        self.assertEqual(repr(ref), expected)
+
+    def test___eq__(self):
+        dataset_ref1 = DatasetReference(self.PROJECT, self.DATASET_ID)
+        ref1 = self._make_one(dataset_ref1, self.PROPERTY_GRAPH_ID)
+        dataset_ref2 = DatasetReference(self.PROJECT, self.DATASET_ID)
+        ref2 = self._make_one(dataset_ref2, self.PROPERTY_GRAPH_ID)
+        self.assertEqual(ref1, ref2)
+
+        ref3 = self._make_one(dataset_ref1, "other_pg")
+        self.assertNotEqual(ref1, ref3)
+        self.assertNotEqual(ref1, object())
+
+    def test___hash__(self):
+        dataset_ref1 = DatasetReference(self.PROJECT, self.DATASET_ID)
+        ref1 = self._make_one(dataset_ref1, self.PROPERTY_GRAPH_ID)
+        dataset_ref2 = DatasetReference(self.PROJECT, self.DATASET_ID)
+        ref2 = self._make_one(dataset_ref2, self.PROPERTY_GRAPH_ID)
+        self.assertEqual(hash(ref1), hash(ref2))
+
+
 class TestTableBase:
     @staticmethod
     def _get_target_class():