From ef3e147fa29289839a056b064b8ca4007acf7962 Mon Sep 17 00:00:00 2001 From: Matt Hammerly Date: Wed, 14 Oct 2026 16:40:24 -0700 Subject: [PATCH] Revert "feat: bigquery_fetcher entrypoint (#67)" This reverts commit ba67dd429efb97b21817ecd6fdd2f7118c1a3209. --- .pre-commit-config.yaml | 1 - Dockerfile | 3 +- py/pyproject.toml | 4 - py/requirements-test.txt | 1 - py/tests/test_bigquery_fetcher.py | 212 ------------------------- py/tests/test_fetcher.py | 57 +++++-- py/tests/test_fetcher_utils.py | 57 ------- py/usageaccountant/accumulator.py | 16 +- py/usageaccountant/bigquery_fetcher.py | 187 ---------------------- py/usageaccountant/datadog_fetcher.py | 85 ++++++++-- py/usageaccountant/fetcher_utils.py | 85 ---------- 11 files changed, 128 insertions(+), 580 deletions(-) delete mode 100644 py/tests/test_bigquery_fetcher.py delete mode 100644 py/tests/test_fetcher_utils.py delete mode 100644 py/usageaccountant/bigquery_fetcher.py delete mode 100644 py/usageaccountant/fetcher_utils.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3266981..a77d68e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -34,7 +34,6 @@ repos: additional_dependencies: [ pytest==7.1.2, sentry-arroyo==2.22.0, - google-cloud-bigquery==3.41.0, ] files: ^py/.+ default_language_version: diff --git a/Dockerfile b/Dockerfile index 33fb385..b78b708 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,8 +2,7 @@ FROM us-docker.pkg.dev/sentryio/dhi-mirror/python:3.13-debian13-dev AS build COPY py/ /app/py/ -# Install `bigquery` extra to enable the `bigquery_fetcher` entrypoint -RUN python3 -m venv /opt/venv && /opt/venv/bin/pip install --no-cache-dir "/app/py[bigquery]" +RUN python3 -m venv /opt/venv && /opt/venv/bin/pip install --no-cache-dir /app/py FROM us-docker.pkg.dev/sentryio/dhi-mirror/python:3.13-debian13 diff --git a/py/pyproject.toml b/py/pyproject.toml index fb01ee9..08e8b02 100644 --- a/py/pyproject.toml +++ b/py/pyproject.toml @@ -16,10 +16,6 @@ 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"] diff --git a/py/requirements-test.txt b/py/requirements-test.txt index b114940..eaf4ac3 100644 --- a/py/requirements-test.txt +++ b/py/requirements-test.txt @@ -1,6 +1,5 @@ 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 diff --git a/py/tests/test_bigquery_fetcher.py b/py/tests/test_bigquery_fetcher.py deleted file mode 100644 index fbcbed5..0000000 --- a/py/tests/test_bigquery_fetcher.py +++ /dev/null @@ -1,212 +0,0 @@ -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 diff --git a/py/tests/test_fetcher.py b/py/tests/test_fetcher.py index 812a1d5..73564c3 100644 --- a/py/tests/test_fetcher.py +++ b/py/tests/test_fetcher.py @@ -1,4 +1,5 @@ import copy +import os import unittest from io import StringIO from unittest.mock import Mock, patch @@ -13,7 +14,6 @@ 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,6 +22,12 @@ 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) @@ -56,6 +62,40 @@ 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( @@ -148,19 +188,17 @@ def test_parse_and_assert_response_scope_no_app_feature(self) -> None: def test_process_series_data(self) -> None: expected_record_list = [ - UsageAccumulatorRecord( + ddf.UsageAccumulatorRecord( resource_id="rc_long_redis", app_feature="shared", amount=2, usage_type=accumulator.UsageUnit.BYTES, - timestamp=None, ), - UsageAccumulatorRecord( + ddf.UsageAccumulatorRecord( resource_id="rc_long_redis", app_feature="shared", amount=3, usage_type=accumulator.UsageUnit.BYTES, - timestamp=None, ), ] @@ -178,12 +216,11 @@ def test_process_series_data(self) -> None: def test_process_series_data_with_none(self) -> None: expected_record_list = [ - UsageAccumulatorRecord( + ddf.UsageAccumulatorRecord( resource_id="rc_long_redis", app_feature="shared", amount=2, usage_type=accumulator.UsageUnit.BYTES, - timestamp=None, ) ] @@ -201,19 +238,17 @@ def test_process_series_data_with_none(self) -> None: def test_process_series_data_feature_tag(self) -> None: expected_record_list = [ - UsageAccumulatorRecord( + ddf.UsageAccumulatorRecord( resource_id="rc_long_redis", app_feature="shared", amount=2, usage_type=accumulator.UsageUnit.BYTES, - timestamp=None, ), - UsageAccumulatorRecord( + ddf.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( diff --git a/py/tests/test_fetcher_utils.py b/py/tests/test_fetcher_utils.py deleted file mode 100644 index bb7cc12..0000000 --- a/py/tests/test_fetcher_utils.py +++ /dev/null @@ -1,57 +0,0 @@ -import os -import unittest -from io import StringIO - -from usageaccountant import accumulator -from usageaccountant import fetcher_utils as fu - - -class TestFetcherUtils(unittest.TestCase): - kafka_config_str = ( - '{"bootstrap_servers": ["kafka.service.host:1234"],' - '"config_params": ' - '{"compression.type": "lz4",' - '"message.max.bytes": 10}}' - ) - - def test_parse_and_assert_kafka_config(self) -> None: - kafka_io = StringIO(self.kafka_config_str) - fu.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, fu.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( - fu.parse_and_assert_kafka_config(kafka_io), expected_kafka_config - ) - - def test_assert_valid_unit(self) -> None: - for unit in accumulator.UsageUnit: - fu.assert_valid_unit(unit.value.upper()) - - def test_assert_valid_unit_invalid(self) -> None: - with self.assertRaises(ValueError): - fu.assert_valid_unit("parsecs") diff --git a/py/usageaccountant/accumulator.py b/py/usageaccountant/accumulator.py index ca41f45..8f18ac6 100644 --- a/py/usageaccountant/accumulator.py +++ b/py/usageaccountant/accumulator.py @@ -128,7 +128,6 @@ def record( app_feature: str, amount: int, usage_type: UsageUnit, - timestamp: Optional[int] = None, ) -> None: """ Record a chunk of usage of a resource for a feature. @@ -143,11 +142,6 @@ def record( `app_feature` identifies the product feature. `amount` is the amount of resource used. `usage_type` is the unit of measure for `amount`. - `timestamp` optionally overrides the time the usage is associated - with (as unixtime seconds). This is useful for larger batch - rollups (like with the `bigquery_fetcher` entrypoint) that may - be handling a lot of usage data from a day earlier than the - current time. """ now = time.time() if ( @@ -157,12 +151,12 @@ def record( self.flush() self.__first_timestamp = now - bucket_time = now if timestamp is None else timestamp - bucket_time = ( - floor(bucket_time / self.__granularity_sec) - * self.__granularity_sec + key = UsageKey( + floor(now / self.__granularity_sec) * self.__granularity_sec, + resource_id, + app_feature, + usage_type, ) - key = UsageKey(bucket_time, resource_id, app_feature, usage_type) if ( key not in self.__usage_batch diff --git a/py/usageaccountant/bigquery_fetcher.py b/py/usageaccountant/bigquery_fetcher.py deleted file mode 100644 index 11f984e..0000000 --- a/py/usageaccountant/bigquery_fetcher.py +++ /dev/null @@ -1,187 +0,0 @@ -import argparse -import logging -from datetime import date, datetime, timezone -from typing import Any, Iterable, List, Mapping, Optional, Sequence, TextIO - -from google.cloud import bigquery - -from usageaccountant.accumulator import UsageAccumulator, UsageUnit -from usageaccountant.fetcher_utils import ( - UsageAccumulatorRecord, - assert_valid_unit, - log_records, - parse_and_assert_kafka_config, - post_to_usage_accumulator, -) - -logger = logging.getLogger("bigquery_fetcher") -logging.basicConfig( - level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s" -) - -APP_FEATURE_COLUMN = "app_feature" -AMOUNT_COLUMN = "amount" -# Optional column identifying the time the usage applies to (e.g. the Storage -# Insights inventory date). Prefer emitting Unix epoch seconds as an ``INT64`` -# (e.g. ``UNIX_SECONDS(...)``); ``DATE`` and ``TIMESTAMP`` columns are also -# accepted and normalized (see ``normalize_timestamp``). When absent, the -# accumulator stamps the record with the run time. -TIMESTAMP_COLUMN = "timestamp" - - -def normalize_timestamp(value: Any) -> Optional[int]: - """ - Normalizes a BigQuery ``timestamp`` column value to integer Unix epoch - seconds (or ``None`` when the value is SQL ``NULL``). - - The expected and preferred type is an ``INT64`` of Unix epoch seconds - (e.g. from ``UNIX_SECONDS(...)``), which the BigQuery client already - returns as a Python ``int`` -- so no conversion is needed. For robustness - we additionally accept: - * ``FLOAT64`` -> truncated to ``int`` - * ``TIMESTAMP`` -> ``datetime.datetime`` (tz-aware) via ``.timestamp()`` - * ``DATE`` -> ``datetime.date``, interpreted as midnight UTC - - Any other type (including ``bool``, which is a subclass of ``int``) raises - ``TypeError`` rather than being silently coerced to a nonsensical epoch. - """ - if value is None: - return None - # bool is a subclass of int; reject explicitly so True/False do not become - # epoch 1/0. - if isinstance(value, bool): - raise TypeError(f"invalid timestamp value: {value!r}") - # datetime is a subclass of date, so check it first. - if isinstance(value, datetime): - return int(value.timestamp()) - if isinstance(value, date): - return int( - datetime( - value.year, value.month, value.day, tzinfo=timezone.utc - ).timestamp() - ) - if isinstance(value, (int, float)): - return int(value) - raise TypeError( - f"unsupported timestamp type: {type(value).__name__} ({value!r})" - ) - - -def parse_and_assert_query_file( - query_file: TextIO, -) -> Sequence[Mapping[str, str]]: - """ - Validates that the query file is a non-empty list and that each entry - contains a ``query``, a ``shared_resource_id`` and a ``unit``. - """ - import json - - query_list = json.loads(query_file.read()) - assert isinstance(query_list, list) - assert query_list, "query file must contain at least one query" - - for query_dict in query_list: - assert "query" in query_dict - assert "shared_resource_id" in query_dict - assert "unit" in query_dict - - return query_list - - -def process_rows( - rows: Iterable[Mapping[str, Any]], unit: UsageUnit, shared_resource_id: str -) -> Sequence[UsageAccumulatorRecord]: - """ - Maps the rows returned by a rollup query into UsageAccumulatorRecords. - - Each row must expose an ``app_feature`` and an ``amount`` column, and may - optionally expose a ``timestamp`` column (see ``normalize_timestamp`` for - the accepted types). - """ - record_list = [] - for row in rows: - app_feature = row[APP_FEATURE_COLUMN] - amount = row[AMOUNT_COLUMN] - # Skip rows with no usage recorded for this feature/day. - if app_feature is None or amount is None: - continue - - timestamp = normalize_timestamp(row.get(TIMESTAMP_COLUMN, None)) - - record_list.append( - UsageAccumulatorRecord( - resource_id=shared_resource_id, - app_feature=app_feature, - amount=int(amount), - usage_type=unit, - timestamp=timestamp, - ) - ) - - return record_list - - -def main( - query_file: TextIO, - usage_accumulator: UsageAccumulator, - bq_client: bigquery.Client, - dry_run: bool, -) -> None: - """ - query_file: File with a list of dictionaries, each containing a BigQuery - ``query``, a ``shared_resource_id`` and a ``unit``. - usage_accumulator: UsageAccumulator object. - bq_client: BigQuery client used to run the queries. - dry_run: When True, log the records instead of producing them to Kafka. - """ - query_list = parse_and_assert_query_file(query_file) - record_list: List[UsageAccumulatorRecord] = [] - for query_dict in query_list: - query = query_dict["query"] - unit = query_dict["unit"] - shared_resource_id = query_dict["shared_resource_id"] - - assert_valid_unit(unit) - usage_unit = UsageUnit(unit.lower()) - - rows = bq_client.query(query).result() - record_list.extend(process_rows(rows, usage_unit, shared_resource_id)) - - if dry_run: - log_records(logger, record_list) - else: - post_to_usage_accumulator(record_list, usage_accumulator) - usage_accumulator.flush() - usage_accumulator.close() - - -if __name__ == "__main__": # pragma: no cover - parser = argparse.ArgumentParser(description="bigquery_fetcher") - parser.add_argument( - "--query_file", - type=argparse.FileType("r"), - help="JSON file containing BigQuery rollup queries, " - "shared_resource_ids and units", - ) - parser.add_argument( - "--kafka_config_file", - type=argparse.FileType("r"), - help="File containing kafka_config for initializing UsageAccumulator", - ) - parser.add_argument( - "--dry_run", - action="store_true", - help="Log data instead of sending it to Kafka", - ) - args = parser.parse_args() - - kafka_config = parse_and_assert_kafka_config(args.kafka_config_file) - - bq_client = bigquery.Client() - - main( - args.query_file, - UsageAccumulator(kafka_config=kafka_config), - bq_client, - args.dry_run, - ) diff --git a/py/usageaccountant/datadog_fetcher.py b/py/usageaccountant/datadog_fetcher.py index ee62d9f..6400f64 100644 --- a/py/usageaccountant/datadog_fetcher.py +++ b/py/usageaccountant/datadog_fetcher.py @@ -8,6 +8,7 @@ Dict, List, Mapping, + NamedTuple, Optional, Sequence, TextIO, @@ -17,13 +18,10 @@ ) from urllib.request import Request, urlopen -from usageaccountant.accumulator import UsageAccumulator, UsageUnit -from usageaccountant.fetcher_utils import ( - UsageAccumulatorRecord, - assert_valid_unit, - log_records, - parse_and_assert_kafka_config, - post_to_usage_accumulator, +from usageaccountant.accumulator import ( + KafkaConfig, + UsageAccumulator, + UsageUnit, ) headers = { @@ -58,6 +56,13 @@ class DatadogResponse(TypedDict): series: Sequence[DatadogResponseSeries] +class UsageAccumulatorRecord(NamedTuple): + resource_id: str + app_feature: str + amount: int + usage_type: UsageUnit + + def parse_and_assert_query_file( query_file: TextIO, ) -> Sequence[Mapping[str, str]]: @@ -75,6 +80,33 @@ def parse_and_assert_query_file( return query_list +def parse_and_assert_kafka_config(kafka_config_file: TextIO) -> KafkaConfig: + """ + Validates if the kafka config file contains + bootstrap_servers and config_params. + """ + kafka_config = json.loads(kafka_config_file.read()) + + assert isinstance(kafka_config, Mapping) + assert "bootstrap_servers" in kafka_config + assert "config_params" in kafka_config + + bootstrap_servers = kafka_config["bootstrap_servers"] + config_params = kafka_config["config_params"] + + if "sasl.password" in config_params: + config_params["sasl.password"] = os.path.expandvars( + config_params["sasl.password"] + ) + + assert isinstance(bootstrap_servers, Sequence) + assert isinstance(config_params, Mapping) + + return KafkaConfig( + bootstrap_servers=bootstrap_servers, config_params=config_params + ) + + def assert_valid_query(query: str) -> None: """ Validates if query string contains @@ -91,6 +123,13 @@ def assert_valid_query(query: str) -> None: assert "app_feature" in query or "feature" in query +def assert_valid_unit(unit: str) -> None: + """ + Validates if unit provided is supported by UsageUnit enum. + """ + UsageUnit(unit.lower()) + + def parse_and_assert_unit(series_list: Sequence[DatadogResponseSeries]) -> Any: """ Extracts unit from a DatadogResponseSeries object. @@ -286,13 +325,41 @@ def process_series_data( # point[0] is the timestamp amount=int(point[1]), usage_type=parsed_unit, - timestamp=None, ) ) return record_list +def post_to_usage_accumulator( + record_list: Sequence[UsageAccumulatorRecord], + usage_accumulator: UsageAccumulator, +) -> None: + """ + Posts UsageAccumulatorRecords to UsageAccumulator. + """ + for record in record_list: + usage_accumulator.record( + resource_id=record.resource_id, + app_feature=record.app_feature, + amount=record.amount, + usage_type=record.usage_type, + ) + + +def log_records(record_list: Sequence[UsageAccumulatorRecord]) -> None: + """ + Logs UsageAccumulatorRecords. + """ + for record in record_list: + logger.info( + f"resource_id: {record.resource_id}, " + f"app_feature: {record.app_feature}, " + f"amount: {record.amount}, " + f"usage_type: {record.usage_type}." + ) + + def main( query_file: TextIO, start_time: int, @@ -344,7 +411,7 @@ def main( ) if dry_run: - log_records(logger, record_list) + log_records(record_list) else: post_to_usage_accumulator(record_list, usage_accumulator) usage_accumulator.flush() diff --git a/py/usageaccountant/fetcher_utils.py b/py/usageaccountant/fetcher_utils.py deleted file mode 100644 index ed60b1b..0000000 --- a/py/usageaccountant/fetcher_utils.py +++ /dev/null @@ -1,85 +0,0 @@ -import json -import os -from logging import Logger -from typing import Mapping, NamedTuple, Optional, Sequence, TextIO - -from usageaccountant.accumulator import ( - KafkaConfig, - UsageAccumulator, - UsageUnit, -) - - -class UsageAccumulatorRecord(NamedTuple): - resource_id: str - app_feature: str - amount: int - usage_type: UsageUnit - timestamp: Optional[int] - - -def parse_and_assert_kafka_config(kafka_config_file: TextIO) -> KafkaConfig: - """ - Validates if the kafka config file contains - bootstrap_servers and config_params. - """ - kafka_config = json.loads(kafka_config_file.read()) - - assert isinstance(kafka_config, Mapping) - assert "bootstrap_servers" in kafka_config - assert "config_params" in kafka_config - - bootstrap_servers = kafka_config["bootstrap_servers"] - config_params = kafka_config["config_params"] - - if "sasl.password" in config_params: - config_params["sasl.password"] = os.path.expandvars( - config_params["sasl.password"] - ) - - assert isinstance(bootstrap_servers, Sequence) - assert isinstance(config_params, Mapping) - - return KafkaConfig( - bootstrap_servers=bootstrap_servers, config_params=config_params - ) - - -def assert_valid_unit(unit: str) -> None: - """ - Validates if unit provided is supported by UsageUnit enum. - """ - UsageUnit(unit.lower()) - - -def post_to_usage_accumulator( - record_list: Sequence[UsageAccumulatorRecord], - usage_accumulator: UsageAccumulator, -) -> None: - """ - Posts UsageAccumulatorRecords to the UsageAccumulator. - """ - for record in record_list: - usage_accumulator.record( - resource_id=record.resource_id, - app_feature=record.app_feature, - amount=record.amount, - usage_type=record.usage_type, - timestamp=record.timestamp, - ) - - -def log_records( - logger: Logger, record_list: Sequence[UsageAccumulatorRecord] -) -> None: - """ - Logs UsageAccumulatorRecords instead of producing them. - """ - for record in record_list: - logger.info( - f"resource_id: {record.resource_id}, " - f"app_feature: {record.app_feature}, " - f"amount: {record.amount}, " - f"usage_type: {record.usage_type}, " - f"timestamp: {record.timestamp}." - )