Billing unit tests (#656)

* Billing api extra labels (#619)

* Added compute_category, cromwell_sub_workflow_name, cromwell_workflow_id, goog_pipelines_worker and wdl_task_name to extended view and created relevant filters and API points.

* Added labels to all BQ queries, refactoring billing layer.

* Added examples to billing-total-cost API regarding the new filters.

* Billing - fixing styling issues after the first Billing release (#624)

* Temporarily disable seqr and hail from /topics API.

* Autoselect 1st topic / 1st project value from the DDL.

* Merging Billing.css into index.css

* Small fix - reusing extRecords in FieldSelector component.

* Refactoring duplicated code in FieldSelector.

* Added Stages to the Group by DDL.

* Billing API IsBillingEnabled (#626)

* Added API point to check if billing is enabled.

* Added simple Total Cost By Batch Page. (#627)

* Added simple Total Cost By Batch Page.

* Billing cost by category (#629)

* Added simple Total Cost By Batch Page.

* Fixed autoselect day format.

* Fixing day format for autoselect (missing leading 0)

* Added first draft of billing page to show detail SKU per selected cost category over selected time periods (day, week, month or invoice month)

* Small fix for BillingCostByBatch page, disable search if searchBy is empty or < 6 chars.

* New: Billing API GET namespaces, added namespace to allowed fields for total cost.

* Implemented HorizontalStackedBarChart, updated Billing By Invoice Month page to enable toggle between chart and table view.

* Stacked Bars Chart with option to accumulate data. (#634)

* Implemented Stacked bars with option to accumulate data.

* Added budget bar to billing horizontal bar chart, added background color for the billing table to reflect the chart colours.

* Added simple prediction of billing stacked bar chart.

* Billing hail batch layout (#633)

* Added simple Total Cost By Batch Page.

* Removing debug prints.

* Fixed autoselect day format.

* Fixing day format for autoselect (missing leading 0)

* Added first draft of billing page to show detail SKU per selected cost category over selected time periods (day, week, month or invoice month)

* Small fix for BillingCostByBatch page, disable search if searchBy is empty or < 6 chars.

* New: Billing API GET namespaces, added namespace to allowed fields for total cost.

* Implemented HorizontalStackedBarChart, updated Billing By Invoice Month page to enable toggle between chart and table view.

* ADD: Cost by Analysis page

* ADD: add start of Analysis grid

* ADD: add start of Analysis grid

* FIX: table fixes for the HailBatchGrid

* API: api changes to enable query of the raw table

* API: fixed and working with updated get_total_cost endpoint

* API: fix typing of get_total_cost (default return is now a list[dict] and can be converted in the layer/route to a specific output type

* API: add endpoint to get costs by batch_id

* API: done

* IN PROGRESS: modifying Cost By Analysis to use new endpoints

* IN PROGRESS: changes to Cost By Analysis, linking with backend API.

* IN PROGRESS: changes to Cost By Analysis, grid grouping by ar/batch/job.

* NEW: finalising Cost By Analysis page

* ADD: durations to Cost By Analysis page

---------

Co-authored-by: Milo Hyben <[email protected]>

* FIX: Billing - fixing time_column condition.

* Removing draft billing page.

* Remove unused API point & cleanup, changes as per code review.

* Small Frontend refactoring, reflecting PR review.

* Updating billing style for dark mode.

* Optimised Frontend, replacing reduce with forEach where possible.

* Refactoring Billing DB structures.

* Cleaning up unused dependencies.

* FIX: replaced button 'color=red' with 'negative' property.

* FIX: replace HEX color for pattern with CSS var.

* FIX: replace async call with sync for a simple function.

* FIX: dark mode for Horizontal Stacked Bar.

* FIX: billing cost by analysis page, esp. search control resizing and functionality.

* FIX: duplicated keys in the grid on Billing Cost By Analysis page.

* FIX: refactoring BQ tables, small fixes for billing pages.

* FIX: BillingCostPageAnalysis, keeping the old record until loading of data finishes.

* FIX: Billing StackedChart various issues.

* Linting

* FIX: missing filters checks, updating charts when loading.

* FIX: silenece linting no attribute msg for Middleware.

* Refactoring filters, implemented first Billing GraphQL integration.

* Fixing linting.

* Added unit tests for BQ filters.

* Fixing linting.

* Added tests for billing routes.

* Removing billing GraphQL, there will be another PR for this.

* Adding doc strings to tests.

* Changing to staticmethod where relevant, added more unitests.

* Linting

* Refactoring string constant so both pylint and unittest are happy.

* Unittests for BQ Function filter.

* Unittests for BillingArBatchTable.

* Added pytz dependency to dev so we can write unit test with timezone aware datetime.

* Fixing timezone aware unit test for billing function filter.

* Refactoring BQ Unittests, added more for BillingBaseTable class.

* Linting

* Added more unit tests for BillingBaseTable.

* Add last missing unittest for BilingBaseTable.

* Added unittests for BillingLayer.

* Linting.

* Merge dev try 2.

* More unit tests for BillingLayer.

* More unit tests for Billing routes functions.

* More unit tests for Billing APi routes.

* Billing unit tests - use mockup author.

* Added BillingLayer muckup to unit tests.

* Billing unit tests refactoring, implementing feedback from PR.

* Fixing billing unit tests.

* More fixes for billing unit tests.

* Removing tests related to old getLabelValue functions.

* Update test/test_api_billing.py

Co-authored-by: Michael Franklin <[email protected]>

* Setting up mock of is_billing_enabled in test class setUp function.

---------

Co-authored-by: Sabrina Yan <[email protected]>
Co-authored-by: Michael Franklin <[email protected]>

db/python/layers/billing.py

@@ -18,7 +18,7 @@ class BillingLayer(BqBaseLayer):
 
     def table_factory(
         self,
-        source: BillingSource,
+        source: BillingSource | None = None,
         fields: list[BillingColumn] | None = None,
         filters: dict[BillingColumn, str | list | dict] | None = None,
     ) -> (

db/python/tables/bq/billing_base.py

@@ -117,8 +117,9 @@ def _execute_query(
         # otherwise return as BQ iterator
         return self._connection.connection.query(query, job_config=job_config)
 
+    @staticmethod
     def _query_to_partitioned_filter(
-        self, query: BillingTotalCostQueryModel
+        query: BillingTotalCostQueryModel,
     ) -> BillingFilter:
         """
         By default views are partitioned by 'day',
@@ -137,15 +138,18 @@ def _query_to_partitioned_filter(
         )
         return billing_filter
 
-    def _filter_to_optimise_query(self) -> str:
+    @staticmethod
+    def _filter_to_optimise_query() -> str:
         """Filter string to optimise BQ query"""
         return 'day >= TIMESTAMP(@start_day) AND day <= TIMESTAMP(@last_day)'
 
-    def _last_loaded_day_filter(self) -> str:
+    @staticmethod
+    def _last_loaded_day_filter() -> str:
         """Last Loaded day filter string"""
         return 'day = TIMESTAMP(@last_loaded_day)'
 
-    def _convert_output(self, query_job_result):
+    @staticmethod
+    def _convert_output(query_job_result):
         """Convert query result to json"""
         if not query_job_result or query_job_result.result().total_rows == 0:
             # return empty list if no record found
@@ -325,8 +329,8 @@ async def _execute_running_cost_query(
             self._execute_query(_query, query_params),
         )
 
+    @staticmethod
     async def _append_total_running_cost(
-        self,
         field: BillingColumn,
         is_current_month: bool,
         last_loaded_day: str | None,
@@ -437,9 +441,8 @@ async def _append_running_cost_records(
 
         return results
 
-    def _prepare_order_by_string(
-        self, order_by: dict[BillingColumn, bool] | None
-    ) -> str:
+    @staticmethod
+    def _prepare_order_by_string(order_by: dict[BillingColumn, bool] | None) -> str:
         """Prepare order by string"""
         if not order_by:
             return ''
@@ -452,9 +455,8 @@ def _prepare_order_by_string(
 
         return f'ORDER BY {",".join(order_by_cols)}' if order_by_cols else ''
 
-    def _prepare_aggregation(
-        self, query: BillingTotalCostQueryModel
-    ) -> tuple[str, str]:
+    @staticmethod
+    def _prepare_aggregation(query: BillingTotalCostQueryModel) -> tuple[str, str]:
         """Prepare both fields for aggregation and group by string"""
         # Get columns to group by
 
@@ -479,7 +481,8 @@ def _prepare_aggregation(
 
         return fields_selected, group_by
 
-    def _prepare_labels_function(self, query: BillingTotalCostQueryModel):
+    @staticmethod
+    def _prepare_labels_function(query: BillingTotalCostQueryModel):
         if not query.filters:
             return None
 
@@ -558,7 +561,7 @@ async def get_total_cost(
             where_str = f'WHERE {where_str}'
 
         _query = f"""
-        {func_filter.fun_implementation if func_filter else ''}
+        {func_filter.func_implementation if func_filter else ''}
 
         WITH t AS (
             SELECT {time_group.field}{time_group.separator} {fields_selected},

db/python/tables/bq/billing_filter.py

@@ -2,6 +2,7 @@
 
 import dataclasses
 import datetime
+from typing import Any
 
 from db.python.tables.bq.generic_bq_filter import GenericBQFilter
 from db.python.tables.bq.generic_bq_filter_model import GenericBQFilterModel
@@ -46,3 +47,17 @@ class BillingFilter(GenericBQFilterModel):
     goog_pipelines_worker: GenericBQFilter[str] = None
     wdl_task_name: GenericBQFilter[str] = None
     namespace: GenericBQFilter[str] = None
+
+    def __eq__(self, other: Any) -> bool:
+        """Equality operator"""
+        result = super().__eq__(other)
+        if not result or not isinstance(other, BillingFilter):
+            return False
+
+        # compare all attributes
+        for att in self.__dict__:
+            if getattr(self, att) != getattr(other, att):
+                return False
+
+        # all attributes are equal
+        return True

db/python/tables/bq/billing_gcp_daily.py

@@ -1,4 +1,5 @@
 from datetime import datetime, timedelta
+from typing import Any
 
 from google.cloud import bigquery
 
@@ -9,7 +10,7 @@
 )
 from db.python.tables.bq.billing_filter import BillingFilter
 from db.python.tables.bq.generic_bq_filter import GenericBQFilter
-from models.models import BillingTotalCostQueryModel
+from models.models import BillingColumn, BillingTotalCostQueryModel
 
 
 class BillingGcpDailyTable(BillingBaseTable):
@@ -21,8 +22,9 @@ def get_table_name(self):
         """Get table name"""
         return self.table_name
 
+    @staticmethod
     def _query_to_partitioned_filter(
-        self, query: BillingTotalCostQueryModel
+        query: BillingTotalCostQueryModel,
     ) -> BillingFilter:
         """
         add extra filter to limit materialized view partition
@@ -77,7 +79,9 @@ async def _last_loaded_day(self):
 
         return None
 
-    def _prepare_daily_cost_subquery(self, field, query_params, last_loaded_day):
+    def _prepare_daily_cost_subquery(
+        self, field: BillingColumn, query_params: list[Any], last_loaded_day: str
+    ):
         """prepare daily cost subquery"""
 
         # add extra filter to limit materialized view partition

db/python/tables/bq/billing_raw.py

@@ -16,8 +16,9 @@ def get_table_name(self):
         """Get table name"""
         return self.table_name
 
+    @staticmethod
     def _query_to_partitioned_filter(
-        self, query: BillingTotalCostQueryModel
+        query: BillingTotalCostQueryModel,
     ) -> BillingFilter:
         """
         Raw BQ billing table is partitioned by usage_end_time

db/python/tables/bq/function_bq_filter.py

@@ -27,14 +27,14 @@ class FunctionBQFilter:
 
     def __init__(self, name: str, implementation: str):
         self.func_name = name
-        self.fun_implementation = implementation
+        self.func_implementation = implementation
         # param_id is a counter for parameterised values
         self._param_id = 0
 
     def to_sql(
         self,
         column_name: BillingColumn,
-        func_params: str | list[Any] | dict[Any, Any],
+        func_params: str | list[Any] | dict[Any, Any] | None = None,
         func_operator: str = None,
     ) -> tuple[str, list[bigquery.ScalarQueryParameter | bigquery.ArrayQueryParameter]]:
         """
@@ -103,7 +103,9 @@ def _sql_value_prep(key: str, value: Any) -> bigquery.ScalarQueryParameter:
         if isinstance(value, float):
             return bigquery.ScalarQueryParameter(key, 'FLOAT64', value)
         if isinstance(value, datetime):
-            return bigquery.ScalarQueryParameter(key, 'STRING', value)
+            return bigquery.ScalarQueryParameter(
+                key, 'STRING', value.isoformat(timespec='seconds')
+            )
 
         # otherwise as string parameter
         return bigquery.ScalarQueryParameter(key, 'STRING', value)

db/python/tables/bq/generic_bq_filter.py

@@ -12,6 +12,19 @@ class GenericBQFilter(GenericFilter[T]):
     Generic BigQuery filter is BQ specific filter class, based on GenericFilter
     """
 
+    def __eq__(self, other):
+        """Equality operator"""
+        if not isinstance(other, GenericBQFilter):
+            return False
+
+        keys = ['eq', 'in_', 'nin', 'gt', 'gte', 'lt', 'lte']
+        for att in keys:
+            if getattr(self, att) != getattr(other, att):
+                return False
+
+        # all attributes are equal
+        return True
+
     def to_sql(
         self, column: str, column_name: str = None
     ) -> tuple[str, dict[str, T | list[T] | Any | list[Any]]]:
@@ -38,13 +51,17 @@ def to_sql(
                 values[k] = self._sql_value_prep(k, self.in_[0])
             else:
                 k = self.generate_field_name(_column_name + '_in')
-                conditionals.append(f'{column} IN ({self._sql_cond_prep(k, self.in_)})')
+                conditionals.append(
+                    f'{column} IN UNNEST({self._sql_cond_prep(k, self.in_)})'
+                )
                 values[k] = self._sql_value_prep(k, self.in_)
         if self.nin is not None:
             if not isinstance(self.nin, list):
                 raise ValueError('NIN filter must be a list')
             k = self.generate_field_name(column + '_nin')
-            conditionals.append(f'{column} NOT IN ({self._sql_cond_prep(k, self.nin)})')
+            conditionals.append(
+                f'{column} NOT IN UNNEST({self._sql_cond_prep(k, self.nin)})'
+            )
             values[k] = self._sql_value_prep(k, self.nin)
         if self.gt is not None:
             k = self.generate_field_name(column + '_gt')
@@ -83,9 +100,14 @@ def _sql_value_prep(key, value):
         Overrides the default _sql_value_prep to handle BQ parameters
         """
         if isinstance(value, list):
-            return bigquery.ArrayQueryParameter(
-                key, 'STRING', ','.join([str(v) for v in value])
-            )
+            if value and isinstance(value[0], int):
+                return bigquery.ArrayQueryParameter(key, 'INT64', value)
+            if value and isinstance(value[0], float):
+                return bigquery.ArrayQueryParameter(key, 'FLOAT64', value)
+
+            # otherwise all list records as string
+            return bigquery.ArrayQueryParameter(key, 'STRING', [str(v) for v in value])
+
         if isinstance(value, Enum):
             return GenericBQFilter._sql_value_prep(key, value.value)
         if isinstance(value, int):

requirements-dev.txt

@@ -15,3 +15,5 @@ strawberry-graphql[debug-server]==0.206.0
 functions_framework
 google-cloud-bigquery
 google-cloud-pubsub
+# following required to unit test some billing functions
+pytz