Added WIP Billing API point to query Big table billing aggregate.

Author: milo-hyben

api/routes/__init__.py

@@ -8,3 +8,4 @@
 from api.routes.web import router as web_router
 from api.routes.enum import router as enum_router
 from api.routes.sequencing_groups import router as sequencing_groups_router
+from api.routes.billing import router as billing_router

api/routes/billing.py

@@ -0,0 +1,34 @@
+"""
+Billing routes
+"""
+
+import json
+from fastapi import APIRouter
+from api.utils.db import (
+    BqConnection,
+    get_projectless_bq_connection,
+)
+from db.python.layers.billing import BillingLayer
+from models.models.billing import BillingRecord
+
+router = APIRouter(prefix='/billing', tags=['billing'])
+
+
+@router.post(
+    '/search',
+    response_model=list[BillingRecord],
+    operation_id='postBillingSearch',
+)
+async def post_search(
+    filter: str = None,
+    connection: BqConnection = get_projectless_bq_connection,
+) -> list[BillingRecord]:
+    """Get Billing records"""
+
+    filter_dict = {}
+    if filter:
+        filter_dict = json.loads(filter)
+
+    billing_layer = BillingLayer(connection)
+    records = await billing_layer.search_records(filter_dict)
+    return records

api/utils/db.py

@@ -10,6 +10,7 @@
 from api.settings import get_default_user
 from api.utils.gcp import email_from_id_token
 from db.python.connect import SMConnections, Connection
+from db.python.gcp_connect import BqConnection, PubSubConnection
 
 
 EXPECTED_AUDIENCE = getenv('SM_OAUTHAUDIENCE')
@@ -76,6 +77,16 @@ async def dependable_get_connection(author: str = Depends(authenticate)):
     return await SMConnections.get_connection_no_project(author)
 
 
+async def dependable_get_bq_connection(author: str = Depends(authenticate)):
+    """FastAPI handler for getting connection withOUT project"""
+    return await BqConnection.get_connection_no_project(author)
+
+
+async def dependable_get_pubsub_connection(author: str = Depends(authenticate)):
+    """FastAPI handler for getting connection withOUT project"""
+    return await PubSubConnection.get_connection_no_project(author)
+
+
 def validate_iap_jwt_and_get_email(iap_jwt, audience):
     """
     Validate an IAP JWT and return email
@@ -102,3 +113,5 @@ def validate_iap_jwt_and_get_email(iap_jwt, audience):
 get_project_readonly_connection = Depends(dependable_get_readonly_project_connection)
 get_project_write_connection = Depends(dependable_get_write_project_connection)
 get_projectless_db_connection = Depends(dependable_get_connection)
+get_projectless_bq_connection = Depends(dependable_get_bq_connection)
+get_projectless_pubsub_connection = Depends(dependable_get_pubsub_connection)

db/python/gcp_connect.py

@@ -0,0 +1,76 @@
+"""
+Code for connecting to Big Query database
+"""
+import logging
+import os
+
+import google.cloud.bigquery as bq
+from google.cloud import pubsub_v1
+from db.python.utils import InternalError
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+
+class BqConnection:
+    """Stores a Big Query DB connection, project and author"""
+
+    def __init__(
+        self,
+        author: str,
+    ):
+        self.gcp_project = os.getenv('METAMIST_GCP_PROJECT')
+        self.connection: bq.Client = bq.Client(project=self.gcp_project)
+        self.author: str = author
+
+    @staticmethod
+    async def get_connection_no_project(author: str):
+        """Get a db connection from a project and user"""
+        # maybe it makes sense to perform permission checks here too
+        logger.debug(f'Authenticate no-project connection with {author!r}')
+
+        # we don't authenticate project-less connection, but rely on the
+        # the endpoint to validate the resources
+
+        return BqConnection(author=author)
+
+
+class BqDbBase:
+    """Base class for big query database subclasses"""
+
+    def __init__(self, connection: BqConnection):
+        if connection is None:
+            raise InternalError(
+                f'No connection was provided to the table {self.__class__.__name__!r}'
+            )
+        if not isinstance(connection, BqConnection):
+            raise InternalError(
+                f'Expected connection type Connection, received {type(connection)}, '
+                f'did you mean to call self._connection?'
+            )
+
+        self._connection = connection
+
+
+class PubSubConnection:
+    """Stores a PubSub connection, project and author"""
+
+    def __init__(
+        self,
+        author: str,
+        topic: str,
+    ):
+        self.client: pubsub_v1.PublisherClient = pubsub_v1.PublisherClient()
+        self.author: str = author
+        self.topic: str = os.getenv('METAMIST_GCP_PROJECT') + topic
+
+    @staticmethod
+    async def get_connection_no_project(author: str):
+        """Get a pubsub connection from a project and user"""
+        # maybe it makes sense to perform permission checks here too
+        logger.debug(f'Authenticate no-project connection with {author!r}')
+
+        # we don't authenticate project-less connection, but rely on the
+        # the endpoint to validate the resources
+
+        return PubSubConnection(author=author)

db/python/layers/billing.py

@@ -0,0 +1,56 @@
+from typing import Any
+
+from models.models import BillingRecord
+
+from db.python.gcp_connect import BqDbBase
+from db.python.layers.bq_base import BqBaseLayer
+
+
+class BillingLayer(BqBaseLayer):
+    """Billing layer"""
+
+    async def search_records(
+        self,
+        filter: dict[str, Any] | None = None,
+    ) -> list[BillingRecord] | None:
+        """
+        Get ETL record for the given request_id
+        """
+        billing_db = BillingDb(self.connection)
+        return await billing_db.search_records(filter)
+
+
+class BillingDb(BqDbBase):
+    """Db layer for billing related routes"""
+
+    async def search_records(
+        self,
+        filter: dict[str, Any] | None = None,
+    ) -> list[BillingRecord] | None:
+        """Get Billing record from BQ"""
+
+        filter_parts = []
+
+        for k, v in filter.items():
+            # value is a tupple of (operator, value)
+            op_value = list(v)
+            filter_parts.append(f'{k} {op_value[0]} {op_value[1]}')
+
+        filter = ' AND '.join(filter_parts)
+
+        if filter_parts:
+            filter = f'WHERE {filter}'
+
+        _query = f"""
+        SELECT * 
+        FROM `{self._connection.gcp_project}.billing_aggregate.aggregate` 
+        {filter}
+        LIMIT 10
+        """
+
+        query_job_result = list(self._connection.connection.query(_query).result())
+
+        if query_job_result:
+            return [BillingRecord.from_json(dict(row)) for row in query_job_result]
+
+        raise ValueError('No record found')

db/python/layers/bq_base.py

@@ -0,0 +1,13 @@
+from db.python.gcp_connect import BqConnection
+
+
+class BqBaseLayer:
+    """Base of all Big Query DB layers"""
+
+    def __init__(self, connection: BqConnection):
+        self.connection = connection
+
+    @property
+    def author(self):
+        """Get author from connection"""
+        return self.connection.author

models/models/__init__.py

@@ -59,3 +59,4 @@
     ProjectSummaryInternal,
     WebProject,
 )
+from models.models.billing import BillingRecord

models/models/billing.py

@@ -0,0 +1,69 @@
+import json
+from models.base import SMBase
+
+
+class BillingRecord(SMBase):
+    """Return class for the Billing record"""
+
+    id: str
+    topic: str
+    service_id: str
+    service_description: str
+
+    gcp_project_id: str
+    gcp_project_name: str
+
+    dataset: str
+    batch_id: str
+    job_id: str
+    batch_name: str
+
+    cost: str
+    currency: str
+    invoice_month: str
+    cost_type: str
+
+    class Config:
+        """Config for BillingRecord Response"""
+
+        orm_mode = True
+
+    @staticmethod
+    def from_json(record):
+        """Create BillingRecord from json"""
+
+        print('\n\n ========================= \n\n')
+
+        print(type(record))
+        print(record)
+
+        print('\n\n ========================= \n\n')
+
+        record['service'] = record['service'] if record['service'] else {}
+        record['project'] = record['project'] if record['project'] else {}
+        record['invoice'] = record['invoice'] if record['invoice'] else {}
+
+        labels = {}
+
+        if record['labels']:
+            for lbl in record['labels']:
+                labels[lbl['key']] = lbl['value']
+
+        record['labels'] = labels
+
+        return BillingRecord(
+            id=record['id'],
+            topic=record['topic'],
+            service_id=record['service'].get('id', ''),
+            service_description=record['service'].get('description', ''),
+            gcp_project_id=record['project'].get('id', ''),
+            gcp_project_name=record['project'].get('name', ''),
+            dataset=record['labels'].get('dataset', ''),
+            batch_id=record['labels'].get('batch_id', ''),
+            job_id=record['labels'].get('job_id', ''),
+            batch_name=record['labels'].get('batch_name', ''),
+            cost=record['cost'],
+            currency=record['currency'],
+            invoice_month=record['invoice'].get('month', ''),
+            cost_type=record['cost_type'],
+        )