diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a77d68e..3266981 100644 --- a/.pre-commit-config.yaml +++ b/.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: diff --git a/Dockerfile b/Dockerfile index a373cad..17c1d23 100644 --- a/Dockerfile +++ b/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 diff --git a/py/pyproject.toml b/py/pyproject.toml index 08e8b02..fb01ee9 100644 --- a/py/pyproject.toml +++ b/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"] diff --git a/py/requirements-test.txt b/py/requirements-test.txt index eaf4ac3..b114940 100644 --- a/py/requirements-test.txt +++ b/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 diff --git a/py/tests/test_bigquery_fetcher.py b/py/tests/test_bigquery_fetcher.py new file mode 100644 index 0000000..fbcbed5 --- /dev/null +++ b/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 diff --git a/py/tests/test_fetcher.py b/py/tests/test_fetcher.py index 73564c3..812a1d5 100644 --- a/py/tests/test_fetcher.py +++ b/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( diff --git a/py/tests/test_fetcher_utils.py b/py/tests/test_fetcher_utils.py new file mode 100644 index 0000000..bb7cc12 --- /dev/null +++ b/py/tests/test_fetcher_utils.py @@ -0,0 +1,57 @@ +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 8f18ac6..ca41f45 100644 --- a/py/usageaccountant/accumulator.py +++ b/py/usageaccountant/accumulator.py @@ -128,6 +128,7 @@ 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. @@ -142,6 +143,11 @@ 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 ( @@ -151,12 +157,12 @@ def record( self.flush() self.__first_timestamp = now - key = UsageKey( - floor(now / self.__granularity_sec) * self.__granularity_sec, - resource_id, - app_feature, - usage_type, + bucket_time = now if timestamp is None else timestamp + bucket_time = ( + floor(bucket_time / self.__granularity_sec) + * self.__granularity_sec ) + 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 new file mode 100644 index 0000000..11f984e --- /dev/null +++ b/py/usageaccountant/bigquery_fetcher.py @@ -0,0 +1,187 @@ +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 6400f64..ee62d9f 100644 --- a/py/usageaccountant/datadog_fetcher.py +++ b/py/usageaccountant/datadog_fetcher.py @@ -8,7 +8,6 @@ Dict, List, Mapping, - NamedTuple, Optional, Sequence, TextIO, @@ -18,10 +17,13 @@ ) from urllib.request import Request, urlopen -from usageaccountant.accumulator import ( - KafkaConfig, - UsageAccumulator, - UsageUnit, +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, ) headers = { @@ -56,13 +58,6 @@ 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]]: @@ -80,33 +75,6 @@ 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 @@ -123,13 +91,6 @@ 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. @@ -325,41 +286,13 @@ 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, @@ -411,7 +344,7 @@ def main( ) if dry_run: - log_records(record_list) + log_records(logger, 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 new file mode 100644 index 0000000..ed60b1b --- /dev/null +++ b/py/usageaccountant/fetcher_utils.py @@ -0,0 +1,85 @@ +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}." + )