feat: bigquery_fetcher entrypoint

.pre-commit-config.yaml

@@ -34,6 +34,7 @@ repos:
         additional_dependencies: [
           pytest==7.1.2,
           sentry-arroyo==2.22.0,
+          google-cloud-bigquery==3.41.0,
         ]
         files: ^py/.+
 default_language_version:

Dockerfile

@@ -1,7 +1,9 @@
 FROM us-docker.pkg.dev/sentryio/dhi-mirror/python:3.13-debian13-dev AS build
 
 COPY py/ /app/py/
-RUN pip install --no-cache-dir /app/py
+
+# Install `bigquery` extra to enable the `bigquery_fetcher` entrypoint
+RUN pip install --no-cache-dir "/app/py[bigquery]"
 
 FROM us-docker.pkg.dev/sentryio/dhi-mirror/python:3.13-debian13
 

py/pyproject.toml

@@ -16,6 +16,10 @@ dependencies = [
     'sentry-arroyo>=2.22.0'
 ]
 
+[project.optional-dependencies]
+# Only needed for the `bigquery_fetcher` entrypoint
+bigquery = ['google-cloud-bigquery>=3.0.0']
+
 [tool.setuptools.package-data]
 usageaccountant = ["py.typed"]
 

py/requirements-test.txt

@@ -1,5 +1,6 @@
 sentry-arroyo>=2.22.0
 pytest>=7.2.1
 mypy>=1.4.1
+google-cloud-bigquery>=3.0.0
 black>=24.10.0
 flake8>=6

py/tests/test_bigquery_fetcher.py

@@ -0,0 +1,212 @@
+import unittest
+from datetime import date, datetime, timezone
+from io import StringIO
+from json import loads
+from typing import Any, Iterable, List, Mapping, cast
+from unittest.mock import Mock
+
+from arroyo.backends.kafka.consumer import KafkaPayload
+from arroyo.backends.local.backend import LocalBroker
+from arroyo.backends.local.storages.memory import MemoryMessageStorage
+from arroyo.types import Partition, Topic
+from arroyo.utils.clock import MockedClock
+from google.cloud import bigquery
+
+from usageaccountant import accumulator
+from usageaccountant import bigquery_fetcher as bqf
+from usageaccountant.accumulator import UsageUnit
+from usageaccountant.fetcher_utils import UsageAccumulatorRecord
+
+# A day-aligned epoch (2026-06-11 00:00:00 GMT).
+INVENTORY_TS = 1781136000
+
+
+class MockBQ:
+    def __init__(self, rows: Iterable[Mapping[str, Any]]):
+        self.rows = rows
+
+    def query(self, query: str) -> Mock:
+        result_fn = Mock(return_value=self.rows)
+        return Mock(result=result_fn)
+
+
+class TestBigQueryFetcher(unittest.TestCase):
+    query_str = (
+        '[{"query": "SELECT app_feature, amount FROM t", '
+        '"unit": "bytes", "shared_resource_id": "gcs_objectstore"}]'
+    )
+
+    def setUp(self) -> None:
+        storage: MemoryMessageStorage[KafkaPayload] = MemoryMessageStorage()
+        self.broker = LocalBroker(storage, MockedClock())
+        self.topic = Topic("test_bq_fetcher")
+        self.broker.create_topic(self.topic, 1)
+
+        producer = self.broker.get_producer()
+        # granularity of 1s so an overridden timestamp is preserved verbatim.
+        self.usage_accumulator = accumulator.UsageAccumulator(
+            1, topic_name="test_bq_fetcher", producer=producer
+        )
+
+    def test_parse_and_assert_query_file(self) -> None:
+        assert bqf.parse_and_assert_query_file(StringIO(self.query_str))
+
+    def test_parse_and_assert_query_file_empty(self) -> None:
+        with self.assertRaises(AssertionError):
+            bqf.parse_and_assert_query_file(StringIO("[]"))
+
+    def test_parse_and_assert_query_file_not_list(self) -> None:
+        with self.assertRaises(AssertionError):
+            bqf.parse_and_assert_query_file(StringIO('{"query": "x"}'))
+
+    def test_parse_and_assert_query_file_missing_query(self) -> None:
+        with self.assertRaises(AssertionError):
+            bqf.parse_and_assert_query_file(
+                StringIO('[{"unit": "bytes", "shared_resource_id": "x"}]')
+            )
+
+    def test_parse_and_assert_query_file_missing_resource(self) -> None:
+        with self.assertRaises(AssertionError):
+            bqf.parse_and_assert_query_file(
+                StringIO('[{"query": "x", "unit": "bytes"}]')
+            )
+
+    def test_parse_and_assert_query_file_missing_unit(self) -> None:
+        with self.assertRaises(AssertionError):
+            bqf.parse_and_assert_query_file(
+                StringIO('[{"query": "x", "shared_resource_id": "y"}]')
+            )
+
+    def test_normalize_timestamp_none(self) -> None:
+        assert bqf.normalize_timestamp(None) is None
+
+    def test_normalize_timestamp_int(self) -> None:
+        assert bqf.normalize_timestamp(INVENTORY_TS) == INVENTORY_TS
+
+    def test_normalize_timestamp_float(self) -> None:
+        assert bqf.normalize_timestamp(INVENTORY_TS + 0.9) == INVENTORY_TS
+
+    def test_normalize_timestamp_datetime(self) -> None:
+        dt = datetime(2026, 6, 11, tzinfo=timezone.utc)
+        assert bqf.normalize_timestamp(dt) == INVENTORY_TS
+
+    def test_normalize_timestamp_date(self) -> None:
+        assert bqf.normalize_timestamp(date(2026, 6, 11)) == INVENTORY_TS
+
+    def test_normalize_timestamp_bool_rejected(self) -> None:
+        with self.assertRaises(TypeError):
+            bqf.normalize_timestamp(True)
+
+    def test_normalize_timestamp_unsupported_type(self) -> None:
+        with self.assertRaises(TypeError):
+            bqf.normalize_timestamp("2026-06-11")
+
+    def test_process_rows_with_timestamp(self) -> None:
+        rows: List[Mapping[str, Any]] = [
+            {
+                "app_feature": "attachments",
+                "amount": 100,
+                "timestamp": INVENTORY_TS,
+            },
+            {
+                "app_feature": "preprod",
+                "amount": 200,
+                "timestamp": INVENTORY_TS,
+            },
+        ]
+        records = bqf.process_rows(rows, UsageUnit.BYTES, "gcs_objectstore")
+        assert records == [
+            UsageAccumulatorRecord(
+                "gcs_objectstore",
+                "attachments",
+                100,
+                UsageUnit.BYTES,
+                INVENTORY_TS,
+            ),
+            UsageAccumulatorRecord(
+                "gcs_objectstore",
+                "preprod",
+                200,
+                UsageUnit.BYTES,
+                INVENTORY_TS,
+            ),
+        ]
+
+    def test_process_rows_without_timestamp(self) -> None:
+        rows: List[Mapping[str, Any]] = [
+            {"app_feature": "attachments", "amount": 100}
+        ]
+        records = bqf.process_rows(
+            rows, UsageUnit.BYTES, "bigtable_objectstore"
+        )
+        assert records == [
+            UsageAccumulatorRecord(
+                "bigtable_objectstore",
+                "attachments",
+                100,
+                UsageUnit.BYTES,
+                None,
+            )
+        ]
+
+    def test_process_rows_skips_nulls(self) -> None:
+        rows: List[Mapping[str, Any]] = [
+            {"app_feature": None, "amount": 100},
+            {"app_feature": "attachments", "amount": None},
+            {"app_feature": "preprod", "amount": 50},
+        ]
+        records = bqf.process_rows(rows, UsageUnit.BYTES, "gcs_objectstore")
+        assert records == [
+            UsageAccumulatorRecord(
+                "gcs_objectstore", "preprod", 50, UsageUnit.BYTES, None
+            )
+        ]
+
+    def test_main_produces_records(self) -> None:
+        client = MockBQ(
+            [
+                {
+                    "app_feature": "attachments",
+                    "amount": "123",
+                    "timestamp": INVENTORY_TS,
+                }
+            ]
+        )
+
+        bqf.main(
+            query_file=StringIO(self.query_str),
+            usage_accumulator=self.usage_accumulator,
+            bq_client=cast(bigquery.Client, client),
+            dry_run=False,
+        )
+
+        msg = self.broker.consume(Partition(self.topic, 0), 0)
+        assert msg is not None
+        payload = loads(msg.payload.value.decode("utf-8"))
+        assert payload == {
+            "timestamp": INVENTORY_TS,
+            "shared_resource_id": "gcs_objectstore",
+            "app_feature": "attachments",
+            "usage_unit": "bytes",
+            "amount": 123,
+        }
+        assert self.broker.consume(Partition(self.topic, 0), 1) is None
+
+    def test_main_dry_run(self) -> None:
+        client = MockBQ(
+            [
+                {
+                    "app_feature": "attachments",
+                    "amount": "123",
+                    "timestamp": INVENTORY_TS,
+                }
+            ]
+        )
+        bqf.main(
+            query_file=StringIO(self.query_str),
+            usage_accumulator=self.usage_accumulator,
+            bq_client=cast(bigquery.Client, client),
+            dry_run=True,
+        )
+        # Nothing is produced in dry-run mode.
+        assert self.broker.consume(Partition(self.topic, 0), 0) is None

py/tests/test_fetcher.py

@@ -1,5 +1,4 @@
 import copy
-import os
 import unittest
 from io import StringIO
 from unittest.mock import Mock, patch
@@ -14,6 +13,7 @@
 from tests.test_data import datadog_response
 from usageaccountant import accumulator
 from usageaccountant import datadog_fetcher as ddf
+from usageaccountant.fetcher_utils import UsageAccumulatorRecord
 
 
 class TestDatadogFetcher(unittest.TestCase):
@@ -22,12 +22,6 @@ class TestDatadogFetcher(unittest.TestCase):
         'by {app_feature,shared_resource_id}","unit": "bytes",'
         '"shared_resource_id": "rc_long_redis"}]'
     )
-    kafka_config_str = (
-        '{"bootstrap_servers": ["kafka.service.host:1234"],'
-        '"config_params": '
-        '{"compression.type": "lz4",'
-        '"message.max.bytes": 10}}'
-    )
 
     def setUp(self) -> None:
         self.good_response = copy.deepcopy(datadog_response.good_response)
@@ -62,40 +56,6 @@ def test_parse_and_assert_query_file_no_query(self) -> None:
             AssertionError, ddf.parse_and_assert_query_file, query_io
         )
 
-    def test_parse_and_assert_kafka_config(self) -> None:
-        kafka_io = StringIO(self.kafka_config_str)
-        ddf.parse_and_assert_kafka_config(kafka_io)
-
-    def test_parse_and_assert_kafka_config_no_config_params(self) -> None:
-        kafka_io = StringIO(
-            '{"bootstrap_servers": ["kafka.service.host:1234"]}'
-        )
-        self.assertRaises(
-            AssertionError, ddf.parse_and_assert_kafka_config, kafka_io
-        )
-
-    def test_parse_and_assert_expandvars(self) -> None:
-        os.environ["SASL_PASSWORD"] = "supersecret"
-
-        expected_kafka_config = accumulator.KafkaConfig(
-            bootstrap_servers=["kafka.service.host:1234"],
-            config_params={
-                "sasl.username": "alice",
-                "sasl.password": "supersecret",
-            },
-        )
-
-        kafka_io = StringIO(
-            '{"bootstrap_servers": ["kafka.service.host:1234"],'
-            '"config_params": '
-            '{"sasl.username": "alice",'
-            '"sasl.password": "${SASL_PASSWORD}"}}'
-        )
-
-        self.assertEqual(
-            ddf.parse_and_assert_kafka_config(kafka_io), expected_kafka_config
-        )
-
     def test_parse_and_assert_unit(self) -> None:
         parsed_unit = ddf.parse_and_assert_unit(
             cast(
@@ -188,17 +148,19 @@ def test_parse_and_assert_response_scope_no_app_feature(self) -> None:
 
     def test_process_series_data(self) -> None:
         expected_record_list = [
-            ddf.UsageAccumulatorRecord(
+            UsageAccumulatorRecord(
                 resource_id="rc_long_redis",
                 app_feature="shared",
                 amount=2,
                 usage_type=accumulator.UsageUnit.BYTES,
+                timestamp=None,
             ),
-            ddf.UsageAccumulatorRecord(
+            UsageAccumulatorRecord(
                 resource_id="rc_long_redis",
                 app_feature="shared",
                 amount=3,
                 usage_type=accumulator.UsageUnit.BYTES,
+                timestamp=None,
             ),
         ]
 
@@ -216,11 +178,12 @@ def test_process_series_data(self) -> None:
 
     def test_process_series_data_with_none(self) -> None:
         expected_record_list = [
-            ddf.UsageAccumulatorRecord(
+            UsageAccumulatorRecord(
                 resource_id="rc_long_redis",
                 app_feature="shared",
                 amount=2,
                 usage_type=accumulator.UsageUnit.BYTES,
+                timestamp=None,
             )
         ]
 
@@ -238,17 +201,19 @@ def test_process_series_data_with_none(self) -> None:
 
     def test_process_series_data_feature_tag(self) -> None:
         expected_record_list = [
-            ddf.UsageAccumulatorRecord(
+            UsageAccumulatorRecord(
                 resource_id="rc_long_redis",
                 app_feature="shared",
                 amount=2,
                 usage_type=accumulator.UsageUnit.BYTES,
+                timestamp=None,
             ),
-            ddf.UsageAccumulatorRecord(
+            UsageAccumulatorRecord(
                 resource_id="rc_long_redis",
                 app_feature="shared",
                 amount=3,
                 usage_type=accumulator.UsageUnit.BYTES,
+                timestamp=None,
             ),
         ]
         parsed = ddf.parse_and_assert_response_series(