feat(bigquery): support fetching bigtable shard boundaries

py/mypy.ini

@@ -1,2 +1,11 @@
 [mypy]
 python_version = 3.13
+
+# google-cloud-bigtable is only imported lazily by the bigtable_sharding
+# row-key sampler and is not installed in the type-check env, so let mypy treat
+# it as untyped rather than failing to resolve the import.
+[mypy-google.cloud.bigtable]
+ignore_missing_imports = True
+
+[mypy-google.cloud.bigtable.*]
+ignore_missing_imports = True

py/pyproject.toml

@@ -17,8 +17,12 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-# Only needed for the `bigquery_fetcher` entrypoint
-bigquery = ['google-cloud-bigquery>=3.0.0']
+# Only needed for the `bigquery_fetcher` entrypoint. google-cloud-bigtable
+# backs the SampleRowKeys row-key sharding for Bigtable-backed queries.
+bigquery = [
+    'google-cloud-bigquery>=3.0.0',
+    'google-cloud-bigtable>=2.0.0',
+]
 
 [tool.setuptools.package-data]
 usageaccountant = ["py.typed"]

py/requirements-test.txt

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

py/tests/test_bigquery_fetcher.py

@@ -1,8 +1,8 @@
 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 json import dumps, loads
+from typing import Any, Iterable, List, Mapping, Sequence, cast
 from unittest.mock import Mock
 
 from arroyo.backends.kafka.consumer import KafkaPayload
@@ -16,6 +16,7 @@
 from usageaccountant import bigquery_fetcher as bqf
 from usageaccountant.accumulator import UsageUnit
 from usageaccountant.fetcher_utils import UsageAccumulatorRecord
+from usageaccountant.query_sharding import QuerySharder, ShardedQuery
 
 # A day-aligned epoch (2026-06-11 00:00:00 GMT).
 INVENTORY_TS = 1781136000
@@ -30,6 +31,43 @@ def query(self, query: str) -> Mock:
         return Mock(result=result_fn)
 
 
+class ShardedMockBQ:
+    """
+    Returns one preset batch of rows per successive ``query`` call and records
+    the (query, parameters) of each call for assertions.
+    """
+
+    def __init__(self, batches: Sequence[List[Mapping[str, Any]]]):
+        self.batches = list(batches)
+        self.calls: List[Any] = []
+
+    def query(self, query: str, job_config: Any = None) -> Mock:
+        params = list(job_config.query_parameters) if job_config else []
+        self.calls.append((query, params))
+        rows = self.batches.pop(0)
+        return Mock(result=Mock(return_value=rows))
+
+
+class FakeSharder(QuerySharder):
+    """
+    A backend-agnostic sharder returning preset shard queries, so the fetcher's
+    orchestration can be tested without any real sharding backend.
+    """
+
+    def __init__(self, shard_queries: List[ShardedQuery]) -> None:
+        self.shard_queries = shard_queries
+
+    def validate_config(
+        self, shard_config: Mapping[str, Any], query: str
+    ) -> None:
+        pass
+
+    def build_shard_queries(
+        self, query: str, shard_config: Mapping[str, Any]
+    ) -> List[ShardedQuery]:
+        return self.shard_queries
+
+
 class TestBigQueryFetcher(unittest.TestCase):
     query_str = (
         '[{"query": "SELECT app_feature, amount FROM t", '
@@ -210,3 +248,147 @@ def test_main_dry_run(self) -> None:
         )
         # Nothing is produced in dry-run mode.
         assert self.broker.consume(Partition(self.topic, 0), 0) is None
+
+    # --- Section: Bigtable row-key sharding ---
+
+    shard_query = (
+        "SELECT app_feature, amount FROM t "
+        "WHERE TRUE {shard_predicate} GROUP BY app_feature"
+    )
+
+    def shard_query_file(self) -> str:
+        return dumps(
+            [
+                {
+                    "query": self.shard_query,
+                    "unit": "bytes",
+                    "shared_resource_id": "bigtable_objectstore",
+                    "bigtable_shard": {
+                        "instance": "objectstore",
+                        "table": "objectstore",
+                    },
+                }
+            ]
+        )
+
+    def test_parse_shard_ok(self) -> None:
+        assert bqf.parse_and_assert_query_file(
+            StringIO(self.shard_query_file())
+        )
+
+    def test_parse_shard_invalid_delegates(self) -> None:
+        # bigtable_shard validation is delegated to bigtable_sharding; a
+        # config missing a required key must still fail the parse.
+        query_file = dumps(
+            [
+                {
+                    "query": self.shard_query,
+                    "unit": "bytes",
+                    "shared_resource_id": "x",
+                    "bigtable_shard": {"table": "objectstore"},
+                }
+            ]
+        )
+        with self.assertRaises(AssertionError):
+            bqf.parse_and_assert_query_file(StringIO(query_file))
+
+    def test_aggregate_records_sums_across_shards(self) -> None:
+        records = [
+            UsageAccumulatorRecord(
+                "bigtable_objectstore", "attachments", 10, UsageUnit.BYTES, 7
+            ),
+            UsageAccumulatorRecord(
+                "bigtable_objectstore", "attachments", 20, UsageUnit.BYTES, 7
+            ),
+            UsageAccumulatorRecord(
+                "bigtable_objectstore", "profiles", 5, UsageUnit.BYTES, 7
+            ),
+        ]
+        aggregated = bqf.aggregate_records(records)
+        by_feature = {r.app_feature: r.amount for r in aggregated}
+        assert by_feature == {"attachments": 30, "profiles": 5}
+
+    def test_run_sharded_query_runs_each_shard(self) -> None:
+        # run_sharded_query is backend-agnostic: it runs each ShardedQuery the
+        # sharder hands it (sql + parameters) and concatenates the rows.
+        shards = [
+            ShardedQuery("SELECT shard_a", [], "[None, 'g')"),
+            ShardedQuery(
+                "SELECT shard_b",
+                [
+                    bigquery.ScalarQueryParameter(
+                        "shard_start", "STRING", "g"
+                    ),
+                    bigquery.ScalarQueryParameter("shard_end", "STRING", "m"),
+                ],
+                "['g', 'm')",
+            ),
+            ShardedQuery("SELECT shard_c", [], "['m', None)"),
+        ]
+        client = ShardedMockBQ(
+            [
+                [{"app_feature": "a", "amount": 1}],
+                [{"app_feature": "a", "amount": 2}],
+                [{"app_feature": "b", "amount": 3}],
+            ]
+        )
+        rows = list(
+            bqf.run_sharded_query(
+                cast(bigquery.Client, client),
+                "unused",
+                {"any": "config"},
+                FakeSharder(shards),
+            )
+        )
+        assert rows == [
+            {"app_feature": "a", "amount": 1},
+            {"app_feature": "a", "amount": 2},
+            {"app_feature": "b", "amount": 3},
+        ]
+        # Each shard's sql and parameters reach the client, in order.
+        assert [query for query, _ in client.calls] == [
+            "SELECT shard_a",
+            "SELECT shard_b",
+            "SELECT shard_c",
+        ]
+        assert [(p.name, p.value) for p in client.calls[1][1]] == [
+            ("shard_start", "g"),
+            ("shard_end", "m"),
+        ]
+
+    def test_main_sharded_aggregates_and_produces(self) -> None:
+        client = ShardedMockBQ(
+            [
+                [{"app_feature": "attachments", "amount": "10"}],
+                # attachments straddles a shard boundary.
+                [{"app_feature": "attachments", "amount": "20"}],
+                [{"app_feature": "profiles", "amount": "5"}],
+            ]
+        )
+        shards = [
+            ShardedQuery(f"SELECT shard_{i}", [], f"shard {i}")
+            for i in range(3)
+        ]
+        bqf.main(
+            query_file=StringIO(self.shard_query_file()),
+            usage_accumulator=self.usage_accumulator,
+            bq_client=cast(bigquery.Client, client),
+            dry_run=False,
+            sharders={"bigtable_shard": FakeSharder(shards)},
+        )
+
+        by_feature = {}
+        index = 0
+        while True:
+            msg = self.broker.consume(Partition(self.topic, 0), index)
+            if msg is None:
+                break
+            payload = loads(msg.payload.value.decode("utf-8"))
+            by_feature[payload["app_feature"]] = payload
+            index += 1
+
+        assert by_feature["attachments"]["amount"] == 30
+        assert by_feature["attachments"]["shared_resource_id"] == (
+            "bigtable_objectstore"
+        )
+        assert by_feature["profiles"]["amount"] == 5

py/tests/test_bigtable_sharding.py

@@ -0,0 +1,142 @@
+import unittest
+from typing import Any, Mapping, Sequence
+
+from usageaccountant import bigtable_sharding as bts
+from usageaccountant.bigtable_sharding import BigtableQuerySharder
+
+# A shardable query: the placeholder sits in a ``WHERE TRUE`` context so the
+# single-shard (empty predicate) case is still valid SQL.
+SHARD_QUERY = (
+    "SELECT app_feature, amount FROM t "
+    "WHERE TRUE {shard_predicate} GROUP BY app_feature"
+)
+
+
+def sampler_gm(_config: Mapping[str, Any]) -> Sequence[bytes]:
+    """Two interior boundaries -> three shards."""
+    return [b"g", b"m"]
+
+
+class TestValidateConfig(unittest.TestCase):
+    def test_ok(self) -> None:
+        BigtableQuerySharder().validate_config(
+            {"instance": "objectstore", "table": "objectstore"}, SHARD_QUERY
+        )
+
+    def test_missing_instance(self) -> None:
+        with self.assertRaises(AssertionError):
+            BigtableQuerySharder().validate_config(
+                {"table": "objectstore"}, SHARD_QUERY
+            )
+
+    def test_missing_table(self) -> None:
+        with self.assertRaises(AssertionError):
+            BigtableQuerySharder().validate_config(
+                {"instance": "objectstore"}, SHARD_QUERY
+            )
+
+    def test_missing_placeholder(self) -> None:
+        with self.assertRaises(AssertionError):
+            BigtableQuerySharder().validate_config(
+                {"instance": "objectstore", "table": "objectstore"},
+                "SELECT app_feature, amount FROM t",
+            )
+
+
+class TestShardRangesFromSamples(unittest.TestCase):
+    def test_empty(self) -> None:
+        assert bts.shard_ranges_from_samples([]) == [(None, None)]
+
+    def test_only_end_sentinel(self) -> None:
+        # The trailing empty key marks end-of-table, not a usable boundary.
+        assert bts.shard_ranges_from_samples([b""]) == [(None, None)]
+
+    def test_basic_sorted_and_bounded(self) -> None:
+        # Unsorted input, with the end sentinel, yields sorted half-open
+        # ranges spanning (None -> ... -> None).
+        assert bts.shard_ranges_from_samples([b"m", b"g", b""]) == [
+            (None, "g"),
+            ("g", "m"),
+            ("m", None),
+        ]
+
+    def test_dedup(self) -> None:
+        assert bts.shard_ranges_from_samples([b"g", b"g", b"m"]) == [
+            (None, "g"),
+            ("g", "m"),
+            ("m", None),
+        ]
+
+    def test_skips_non_utf8(self) -> None:
+        # A non-UTF-8 boundary is dropped (merging its two shards).
+        assert bts.shard_ranges_from_samples([b"\xff\xfe", b"m"]) == [
+            (None, "m"),
+            ("m", None),
+        ]
+
+
+class TestBuildShardClause(unittest.TestCase):
+    def test_both_bounds(self) -> None:
+        clause, params = bts.build_shard_clause("g", "m", "rowkey")
+        assert clause == (
+            " AND rowkey >= @shard_start AND rowkey < @shard_end"
+        )
+        assert [(p.name, p.value) for p in params] == [
+            ("shard_start", "g"),
+            ("shard_end", "m"),
+        ]
+
+    def test_start_only(self) -> None:
+        clause, params = bts.build_shard_clause("m", None, "rowkey")
+        assert clause == " AND rowkey >= @shard_start"
+        assert [(p.name, p.value) for p in params] == [("shard_start", "m")]
+
+    def test_end_only(self) -> None:
+        clause, params = bts.build_shard_clause(None, "g", "rowkey")
+        assert clause == " AND rowkey < @shard_end"
+        assert [(p.name, p.value) for p in params] == [("shard_end", "g")]
+
+    def test_unbounded(self) -> None:
+        clause, params = bts.build_shard_clause(None, None, "rowkey")
+        assert clause == ""
+        assert params == []
+
+    def test_invalid_column(self) -> None:
+        with self.assertRaises(ValueError):
+            bts.build_shard_clause("g", "m", "rowkey; DROP TABLE t")
+
+
+class TestBuildShardQueries(unittest.TestCase):
+    def test_defaults_to_live_sampler(self) -> None:
+        assert BigtableQuerySharder().sampler is bts.default_sampler
+
+    def test_predicate_per_shard(self) -> None:
+        sharder = BigtableQuerySharder(sampler=sampler_gm)
+        shard_queries = sharder.build_shard_queries(
+            SHARD_QUERY, {"instance": "i", "table": "t"}
+        )
+        assert len(shard_queries) == 3
+
+        first = shard_queries[0]
+        assert "{shard_predicate}" not in first.sql
+        assert "rowkey < @shard_end" in first.sql
+        assert "rowkey >= @shard_start" not in first.sql
+        assert first.description == "[None, 'g')"
+
+        mid = shard_queries[1]
+        assert "rowkey >= @shard_start" in mid.sql
+        assert "rowkey < @shard_end" in mid.sql
+        assert [(p.name, p.value) for p in mid.parameters] == [
+            ("shard_start", "g"),
+            ("shard_end", "m"),
+        ]
+        assert mid.description == "['g', 'm')"
+
+        last = shard_queries[2]
+        assert "rowkey >= @shard_start" in last.sql
+        assert "rowkey < @shard_end" not in last.sql
+        assert last.description == "['m', None)"
+
+
+if __name__ == "__main__":
+    unittest.main()