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

Author: AbgarSim

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

@@ -70,6 +70,9 @@
 # Protect against environments where bigquery library is not available.
 try:
   import regex
+  import threading
+  from cachetools import TTLCache, cachedmethod
+  from cachetools.keys import hashkey
   from apitools.base.py.exceptions import HttpError
   from apitools.base.py.exceptions import HttpForbiddenError
   from apitools.base.py.transfer import Upload
@@ -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
@@ -789,10 +803,22 @@ def _insert_all_rows(
     return not errors, errors
 
   @retry.with_exponential_backoff(
-      num_retries=MAX_RETRIES,
-      retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter)
+    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

@@ -25,6 +25,9 @@
 import math
 import re
 import unittest
+
+from unittest.mock import Mock
+from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
 from typing import Optional
 from typing import Sequence
 
@@ -292,6 +295,35 @@ 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):
+    client = Mock()
+    client.tables = Mock()
+
+    returned_table = Mock(name="BigQueryTable")
+    client.tables.Get = 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,), _kwargs = 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(