[BEAM-34076] Added TTL caching for BigQuery table definitions

Author: AbgarSim

sdks/python/apache_beam/io/gcp/bigquery_tools.py

@@ -37,6 +37,7 @@
 import sys
 import time
 import uuid
+import threading
 from json.decoder import JSONDecodeError
 from typing import Optional
 from typing import Sequence
@@ -66,6 +67,8 @@
 from apache_beam.typehints.typehints import Any
 from apache_beam.utils import retry
 from apache_beam.utils.histogram import LinearBucket
+from cachetools import TTLCache, cachedmethod
+from cachetools.keys import hashkey
 
 # Protect against environments where bigquery library is not available.
 try:
@@ -359,6 +362,9 @@ class BigQueryWrapper(object):
 
   HISTOGRAM_METRIC_LOGGER = MetricLogger()
 
+  _TABLE_CACHE_MAX_SIZE = 1024
+  _TABLE_CACHE_TTL_SECONDS = 300
+
   def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None):
     self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions())
     self.gcp_bq_client = client or gcp_bigquery.Client(
@@ -393,6 +399,14 @@ def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None):
       self._temporary_table_suffix = uuid.uuid4().hex
       self.temp_dataset_id = temp_dataset_id or self._get_temp_dataset()
 
+    self._table_cache_ttl_secs = self._TABLE_CACHE_TTL_SECONDS
+    self._table_cache_maxsize = self._TABLE_CACHE_MAX_SIZE
+    self._table_cache = TTLCache(
+        maxsize=self._table_cache_maxsize,
+        ttl=self._table_cache_ttl_secs,
+    )
+    self._table_cache_lock = threading.RLock()
+
     self.created_temp_dataset = False
 
   @property
@@ -791,8 +805,20 @@ def _insert_all_rows(
   @retry.with_exponential_backoff(
       num_retries=MAX_RETRIES,
       retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter)
+  def _get_table_uncached(self, project_id, dataset_id, table_id):
+    """Uncached tables.get call with retry/backoff."""
+    request = bigquery.BigqueryTablesGetRequest(
+        projectId=project_id, datasetId=dataset_id, tableId=table_id)
+    return self.client.tables.Get(request)
+
+  @cachedmethod(
+      cache=lambda self: self._table_cache,
+      lock=lambda self: self._table_cache_lock,
+      key=lambda self, project_id, dataset_id, table_id: hashkey(
+          project_id, dataset_id, table_id),
+  )
   def get_table(self, project_id, dataset_id, table_id):
-    """Lookup a table's metadata object.
+    """Lookup a table's metadata object (cached with TTL, thread-safe).
 
     Args:
       client: bigquery.BigqueryV2 instance
@@ -805,10 +831,7 @@ def get_table(self, project_id, dataset_id, table_id):
     Raises:
       HttpError: if lookup failed.
     """
-    request = bigquery.BigqueryTablesGetRequest(
-        projectId=project_id, datasetId=dataset_id, tableId=table_id)
-    response = self.client.tables.Get(request)
-    return response
+    return self._get_table_uncached(project_id, dataset_id, table_id)
 
   def _create_table(
       self,

sdks/python/apache_beam/io/gcp/bigquery_tools_test.py

@@ -292,6 +292,38 @@ def test_temporary_dataset_is_unique(self, patched_time_sleep):
       wrapper.create_temporary_dataset('project-id', 'location')
     self.assertTrue(client.datasets.Get.called)
 
+  def test_get_table_invokes_tables_get_and_caches_result(self):
+
+    from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
+
+    client = mock.Mock()
+    client.tables = mock.Mock()
+
+    returned_table = mock.Mock(name="BigQueryTable")
+    client.tables.Get = mock.Mock(return_value=returned_table)
+
+    wrapper = BigQueryWrapper(client=client)
+
+    project_id = "my-project"
+    dataset_id = "my_dataset"
+    table_id = "my_table"
+
+    table1 = wrapper.get_table(project_id, dataset_id, table_id)
+
+    assert table1 is returned_table
+    assert client.tables.Get.call_count == 1
+
+    (request,), _ = client.tables.Get.call_args
+    assert isinstance(request, bigquery.BigqueryTablesGetRequest)
+    assert request.projectId == project_id
+    assert request.datasetId == dataset_id
+    assert request.tableId == table_id
+
+    table2 = wrapper.get_table(project_id, dataset_id, table_id)
+
+    assert table2 is returned_table
+    assert client.tables.Get.call_count == 1  # still 1 => cached
+
   def test_get_or_create_dataset_created(self):
     client = mock.Mock()
     client.datasets.Get.side_effect = HttpError(