[bid-manager-api] feat: Add query cache and report reuse

Author: Vasist10

libs/community/google/bid-manager/garf_bid_manager/api_clients.py

@@ -15,6 +15,7 @@
 
 import csv
 import io
+import json
 import logging
 import os
 import pathlib
@@ -34,6 +35,8 @@
 _DEFAULT_API_SCOPES = ['https://www.googleapis.com/auth/doubleclickbidmanager']
 
 _SERVICE_ACCOUNT_CREDENTIALS_FILE = str(pathlib.Path.home() / 'dbm.json')
+_QUERY_CACHE_ENV = 'GARF_BID_MANAGER_QUERY_CACHE_DIR'
+_DEFAULT_QUERY_CACHE_DIR = pathlib.Path.home() / '.garf/bid_manager'
 
 
 class BidManagerApiClientError(exceptions.BidManagerApiError):
@@ -50,6 +53,7 @@ def __init__(
       'GARF_BID_MANAGER_CREDENTIALS_FILE', _SERVICE_ACCOUNT_CREDENTIALS_FILE
     ),
     auth_mode: Literal['oauth', 'service_account'] = 'oauth',
+    query_cache_dir: str | pathlib.Path | None = None,
     **kwargs: str,
   ) -> None:
     """Initializes BidManagerApiClient."""
@@ -59,6 +63,10 @@ def __init__(
     self.kwargs = kwargs
     self._client = None
     self._credentials = None
+    cache_dir = query_cache_dir or os.getenv(_QUERY_CACHE_ENV)
+    self.query_cache_dir = (
+      pathlib.Path(cache_dir) if cache_dir else _DEFAULT_QUERY_CACHE_DIR
+    )
 
   @property
   def credentials(self):
@@ -87,36 +95,36 @@ def client(self):
   def get_response(
     self, request: query_editor.BidManagerApiQuery, **kwargs: str
   ) -> api_clients.GarfApiResponse:
-    query = _build_request(request)
-    query_response = self.client.queries().create(body=query).execute()
-    report_response = (
-      self.client.queries()
-      .run(queryId=query_response['queryId'], synchronous=False)
-      .execute()
-    )
-    query_id = report_response['key']['queryId']
-    report_id = report_response['key']['reportId']
-    logging.info(
-      'Query %s is running, report %s has been created and is '
-      'currently being generated.',
-      query_id,
-      report_id,
-    )
+    query_hash = request.hash
+    query_id = None
+    report_id = None
+    status = None
 
-    get_request = (
-      self.client.queries()
-      .reports()
-      .get(
-        queryId=report_response['key']['queryId'],
-        reportId=report_response['key']['reportId'],
+    if cached_ids := self._load_cached_query_reference(query_hash):
+      cached_query_id, cached_report_id = cached_ids
+      logging.info(
+        'Attempting to reuse DV360 report %s for query hash %s.',
+        cached_report_id,
+        query_hash,
       )
-    )
+      try:
+        status = self._get_report_status(cached_query_id, cached_report_id)
+        query_id, report_id = cached_query_id, cached_report_id
+      except Exception as exc:  # pylint: disable=broad-except
+        logging.warning(
+          'Unable to reuse DV360 report %s (hash %s), regenerating. Reason: %s',
+          cached_report_id,
+          query_hash,
+          exc,
+        )
+        status = None
 
-    status = _check_if_report_is_done(get_request)
+    if status is None:
+      query_id, report_id = self._run_query(request)
+      self._save_cached_query_reference(query_hash, query_id, report_id)
+      status = self._get_report_status(query_id, report_id)
 
-    logging.info(
-      'Report %s generated successfully. Now downloading.', report_id
-    )
+    logging.info('Report %s generated successfully. Now downloading.', report_id)
     with smart_open.open(
       status['metadata']['googleCloudStoragePath'], 'r', encoding='utf-8'
     ) as f:
@@ -144,6 +152,64 @@ def _get_oauth_credentials(self):
     )
 
 
+  def _run_query(
+    self, request: query_editor.BidManagerApiQuery
+  ) -> tuple[str, str]:
+    query = _build_request(request)
+    query_response = self.client.queries().create(body=query).execute()
+    report_response = (
+      self.client.queries()
+      .run(queryId=query_response['queryId'], synchronous=False)
+      .execute()
+    )
+    query_id = report_response['key']['queryId']
+    report_id = report_response['key']['reportId']
+    logging.info(
+      'Query %s is running, report %s has been created and is currently '
+      'being generated.',
+      query_id,
+      report_id,
+    )
+    return query_id, report_id
+
+  def _get_report_status(self, query_id: str, report_id: str):
+    get_request = (
+      self.client.queries()
+      .reports()
+      .get(
+        queryId=query_id,
+        reportId=report_id,
+      )
+    )
+    return _check_if_report_is_done(get_request)
+
+  def _load_cached_query_reference(
+    self, query_hash: str
+  ) -> tuple[str, str] | None:
+    cache_path = self.query_cache_dir / f'{query_hash}.txt'
+    if not cache_path.is_file():
+      return None
+    try:
+      with open(cache_path, 'r', encoding='utf-8') as cache_file:
+        data = json.load(cache_file)
+      return data['query_id'], data['report_id']
+    except (OSError, ValueError, KeyError) as exc:
+      logging.warning(
+        'Failed to load DV360 cache file %s, ignoring. Reason: %s',
+        cache_path,
+        exc,
+      )
+      return None
+
+  def _save_cached_query_reference(
+    self, query_hash: str, query_id: str, report_id: str
+  ) -> None:
+    self.query_cache_dir.mkdir(parents=True, exist_ok=True)
+    cache_path = self.query_cache_dir / f'{query_hash}.txt'
+    with open(cache_path, 'w', encoding='utf-8') as cache_file:
+      json.dump({'query_id': query_id, 'report_id': report_id}, cache_file)
+
+
 def _build_request(request: query_editor.BidManagerApiQuery):
   """Builds Bid Manager API query object from BidManagerApiQuery."""
   query = {

libs/community/google/bid-manager/tests/unit/test_api_clients.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
+import io
+import json
+from unittest import mock
+
 import pytest
 from garf_bid_manager import api_clients, query_editor
 
@@ -85,3 +90,110 @@ def test_process_api_response():
   ]
 
   assert result == expected_result
+
+
+def _build_query_spec():
+  query = """
+    SELECT
+      advertiser AS advertiser,
+      metric_impressions AS impressions
+    FROM standard
+    WHERE advertiser = 1
+    """
+  return query_editor.BidManagerApiQuery(text=query, title='test').generate()
+
+
+def _mock_report_status():
+  return {
+    'key': {'queryId': '123', 'reportId': '456'},
+    'metadata': {
+      'status': {'state': 'DONE'},
+      'googleCloudStoragePath': 'gs://bucket/report.csv',
+    },
+  }
+
+
+def _patch_open(monkeypatch, data: str):
+  @contextlib.contextmanager
+  def fake_open(*_, **__):
+    yield io.StringIO(data)
+
+  monkeypatch.setattr(api_clients.smart_open, 'open', fake_open)
+
+
+def test_get_response_reuses_cached_report(tmp_path, monkeypatch):
+  spec = _build_query_spec()
+  cache_file = tmp_path / f'{spec.hash}.txt'
+  cache_file.write_text(
+    json.dumps({'query_id': '123', 'report_id': '456'}),
+    encoding='utf-8',
+  )
+
+  client = api_clients.BidManagerApiClient(query_cache_dir=tmp_path)
+
+  mock_client = mock.Mock()
+  mock_queries = mock.Mock()
+  mock_reports = mock.Mock()
+  mock_get_request = mock.Mock()
+  mock_get_request.execute.return_value = _mock_report_status()
+  mock_reports.get.return_value = mock_get_request
+  mock_queries.reports.return_value = mock_reports
+  mock_client.queries.return_value = mock_queries
+  client._client = mock_client
+
+  _patch_open(monkeypatch, 'col1,col2\nvalue1,value2\n')
+
+  response = client.get_response(spec)
+
+  assert response.results == [
+    {'FILTER_ADVERTISER': 'value1', 'METRIC_IMPRESSIONS': 'value2'}
+  ]
+  mock_queries.create.assert_not_called()
+  mock_queries.run.assert_not_called()
+
+
+def test_get_response_creates_and_caches_report(tmp_path, monkeypatch):
+  spec = _build_query_spec()
+  client = api_clients.BidManagerApiClient(query_cache_dir=tmp_path)
+
+  mock_client = mock.Mock()
+  mock_queries = mock.Mock()
+  mock_reports = mock.Mock()
+
+  mock_create = mock.Mock()
+  mock_create.execute.return_value = {'queryId': '999'}
+  mock_run = mock.Mock()
+  mock_run.execute.return_value = {
+    'key': {'queryId': '999', 'reportId': '555'}
+  }
+  mock_get_request = mock.Mock()
+  mock_get_request.execute.return_value = {
+    'key': {'queryId': '999', 'reportId': '555'},
+    'metadata': {
+      'status': {'state': 'DONE'},
+      'googleCloudStoragePath': 'gs://bucket/report.csv',
+    },
+  }
+
+  mock_reports.get.return_value = mock_get_request
+  mock_queries.create.return_value = mock_create
+  mock_queries.run.return_value = mock_run
+  mock_queries.reports.return_value = mock_reports
+  mock_client.queries.return_value = mock_queries
+  client._client = mock_client
+
+  _patch_open(monkeypatch, 'col1,col2\nvalue1,value2\n')
+
+  response = client.get_response(spec)
+
+  assert response.results == [
+    {'FILTER_ADVERTISER': 'value1', 'METRIC_IMPRESSIONS': 'value2'}
+  ]
+  mock_queries.create.assert_called_once()
+  mock_queries.run.assert_called_once()
+  cache_file = tmp_path / f'{spec.hash}.txt'
+  assert cache_file.is_file()
+  assert json.loads(cache_file.read_text(encoding='utf-8')) == {
+    'query_id': '999',
+    'report_id': '555',
+  }