Add endpoint to query known cells

Author: ThiagoAVicente

main.py

@@ -5,8 +5,7 @@
 from fastapi import FastAPI
 
 from src.routers.v1 import v1_router
-from src.routers.v1.latency_router import ClickHouse
-from src.routers.v1.raw_router import Influx
+from src.services.databases import ClickHouse, Influx
 
 from src.sink import KafkaSinkManager
 

src/models/raw.py

@@ -3,7 +3,7 @@
 from influxdb_client.client.write.point import Point
 from typing import Optional
 
-RAW:str = "raw"
+RAW_MEASUREMENT:str = "raw"
 
 
 
@@ -38,10 +38,10 @@ class Raw(BaseModel):
 
     def to_point(self) -> Point:
         """Convert to InfluxDB Point for writing"""
-        point = Point(RAW).time(self.timestamp)
+        point = Point(RAW_MEASUREMENT).time(self.timestamp)
 
         # Tags (indexed, used for filtering/grouping)
-        tags = ["cell_index", "network"]
+        tags = {"cell_index", "network"}
 
         # Add tags and fields
         for key, value in self.model_dump(exclude={"timestamp"}).items():

src/routers/v1/__init__.py

@@ -1,7 +1,9 @@
 from fastapi import APIRouter
 from src.routers.v1.latency_router import router as latencyR
 from src.routers.v1.raw_router import router as rawR
+from src.routers.v1.cells_router import router as cellR
 
 v1_router = APIRouter()
 v1_router.include_router(latencyR, prefix="/processed", tags=["v1", "latency"])
 v1_router.include_router(rawR, prefix="/raw", tags=["v1", "latency"])
+v1_router.include_router(cellR, prefix="/cell", tags=["v1", "cell"])

src/routers/v1/cells_router.py

@@ -0,0 +1,19 @@
+"""
+Endpoints for query data
+"""
+from fastapi import APIRouter, HTTPException
+from src.services.databases import Influx
+
+router = APIRouter()
+
+@router.get("/", response_model=list[int])
+def get_processed_latency():
+    """
+    Return a list of know cell indexes
+    """
+    try:
+        results = Influx.service.get_known_cells()
+        return results
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error querying processed latency: {str(e)}")

src/routers/v1/latency_router.py

@@ -2,16 +2,11 @@
 Endpoints for query data
 """
 from fastapi import APIRouter, HTTPException, Query
-from src.configs.clickhouse_conf import ClickhouseConf
-from src.services.clickhouse import ClickHouseService
+from src.services.databases import ClickHouse
 from src.models.processed_latency import ProcessedLatency
 
 router = APIRouter()
 
-class ClickHouse():
-    conf = ClickhouseConf()
-    service = ClickHouseService()
-
 
 @router.get("/latency/example", response_model=list[ProcessedLatency])
 def get_latency_example():

src/routers/v1/raw_router.py

@@ -1,15 +1,10 @@
 from fastapi import APIRouter, HTTPException, Query
-from src.configs.influx_conf import InfluxConf
-from src.services.influx import InfluxService
 from src.models.raw import RawResponse
 from datetime import datetime
-
+from src.services.databases import Influx
 
 router = APIRouter()
 
-class Influx():
-    conf = InfluxConf()
-    service = InfluxService()
 
 
 @router.get("/", response_model=RawResponse)

src/services/clickhouse.py

@@ -3,7 +3,6 @@
 from clickhouse_connect.driver.client import Client
 from src.configs.clickhouse_conf import ClickhouseConf
 from src.models.processed_latency import ProcessedLatency
-from src.services.db_service import DBService
 from src.services.clickhouse_query import QueryCH
 
 
@@ -78,7 +77,7 @@ def transform_processor_output(data: dict) -> dict:
 
     return transformed
 
-class ClickHouseService(DBService):
+class ClickHouseService:
     def __init__(self) -> None:
         self.conf =   ClickhouseConf()
         self.client:  Client = None
@@ -92,9 +91,6 @@ def connect(self):
             password = self.conf.password,
         )
 
-    def get_data(self, batch_number: int = 1, batch_size: int = 50) -> list:
-        return []
-
     def write_data(self, data: dict) -> None:
         """Write a single processed latency record to ClickHouse"""
         try:

src/services/databases.py

@@ -0,0 +1,14 @@
+from src.configs.clickhouse_conf import ClickhouseConf
+from src.services.clickhouse import ClickHouseService
+from src.configs.influx_conf import InfluxConf
+from src.services.influx import InfluxService
+
+
+class ClickHouse:
+    conf = ClickhouseConf()
+    service = ClickHouseService()
+
+
+class Influx:
+    conf = InfluxConf()
+    service = InfluxService()

src/services/db_service.py

@@ -1,12 +1,10 @@
 from influxdb_client.client.influxdb_client import InfluxDBClient
-from influxdb_client.client.write_api import ASYNCHRONOUS, WriteApi
-from influxdb_client.client.query_api import QueryApi
+from influxdb_client.client.write_api import ASYNCHRONOUS
 from src.configs.influx_conf import InfluxConf
-from src.models.raw import Raw
-from src.services.db_service import DBService
+from src.models.raw import Raw, RAW_MEASUREMENT
 from src.services.influx_query import QueryIF
 
-class InfluxService(DBService):
+class InfluxService:
     def __init__(self) -> None:
         self.conf=      InfluxConf()
 
@@ -22,10 +20,6 @@ def connect(self):
         self.write_api = self.client.write_api(write_options=ASYNCHRONOUS)
         self.query_api = self.client.query_api()
 
-    def get_data(self, batch_number:int = 1 ,batch_size:int = 50) -> list[Raw]:
-        #TODO: add query
-        return []
-
     def query_raw_data(self, start_time: int, end_time: int, cell_index: int, batch_number: int):
         _LIMIT = 50
         offset = (batch_number-1) * _LIMIT
@@ -61,3 +55,23 @@ def write_batch(self, data_list:list[dict]) -> None :
         raw_list = [Raw(**d) for d in data_list]
         points = [r.to_point() for r in raw_list]
         self.write_api.write(bucket=self.conf.bucket, org = self.conf.org ,record = points)
+
+    def get_known_cells(self) -> list[int]:
+        """Returns a list of known cell indexes"""
+
+        query = QueryIF.get_known_cells.format(
+            bucket=self.conf.bucket,
+            measurement=RAW_MEASUREMENT
+        )
+
+        tables = self.query_api.query(query)
+
+        cells = []
+        for table in tables:
+            for record in table.records:
+                try:
+                    cells.append(int(record["_value"]))
+                except (TypeError, ValueError):
+                    pass
+
+        return set(cells)