feat: add location query

Author: NejcKle

rest_api/main.py

@@ -2,18 +2,20 @@
 from contextlib import asynccontextmanager
 from functools import partial
 
-from fastapi import Depends, FastAPI
+from fastapi import Body, Depends, FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import ORJSONResponse, PlainTextResponse
 from models.return_models import (
     ApplicationListResponse,
+    LocationListResponse,
     MeasurementListResponse,
     SensorListResponse,
 )
 from prometheus_fastapi_instrumentator import Instrumentator
 from psycopg2 import pool
 from src.handlers import (
     fetch_applications,
+    fetch_locations_for_device_list,
     fetch_measurements_for_sensor,
     fetch_sensors,
     fetch_sensors_for_application,
@@ -140,6 +142,26 @@ def get_sensors_for_application(
     return {"sensors": sensors}
 
 
+@app.post(
+    "/api/locations",
+    response_model=LocationListResponse,
+    tags=["measurements"],
+    responses={
+        401: {"description": "Forbidden - Invalid API Key"},
+        403: {"description": "Forbidden - Not authenticated"},
+        404: {"description": "Not Found - Sensor not found"},
+        500: {"description": "Internal Server Error"},
+    },
+)
+def get_locations(
+    conn=Depends(partial(get_db_connection, connection_pool)),
+    sensor_list=Body(...),
+):
+    """Fetch locations for a given list of sensors."""
+    ret = fetch_locations_for_device_list(conn, sensor_list)
+    return {"locations": ret}
+
+
 @app.get(
     "/api/measurements",
     response_model=MeasurementListResponse,

rest_api/models/return_models.py

@@ -31,6 +31,12 @@ class MeasurementListResponse(BaseModel):
     page_info: PageInfo  # Forward reference to PageInfo
 
 
+class LocationListResponse(BaseModel):
+    """Data model for a list of locations."""
+
+    locations: list[Measurement]
+
+
 class Sensor(BaseModel):
     """Data model for sensor information."""
 

rest_api/src/handlers.py

@@ -155,28 +155,95 @@ def fetch_measurements_for_sensor(
                         )
             except Exception as e:
                 print(f"Error executing query on table {table_name}: {e}")
-            # query = sql.SQL("SELECT * FROM {} WHERE id = ANY(%s)").format(
-            #     sql.Identifier(table_name)
-            # )
-            # cur.execute(query, (ids,))
-            # columns = [desc[0] for desc in cur.description]
-            # for row in cur.fetchall():
-            #     data = {col: val for col, val in zip(columns[3:], row[3:])}
-            #     filtered_data = {k: v for k, v in data.items() if v is not None}
-            #     measurements.append(
-            #         Measurement(
-            #             # sensor_id=row[1],
-            #             table_name=table_name,
-            #             timestamp=str(row[2]),
-            #             data=filtered_data,
-            #         )
-            #     )
 
         # Now sort all measurements by timestamp descending
         measurements.sort(key=lambda x: x.timestamp, reverse=True)
         return measurements
 
 
+def fetch_locations_for_device_list(conn, device_list: list[str]) -> list:
+    """Fetch the latest location for each device in the device_list."""
+    with conn.cursor() as cur:
+        query = """
+            SELECT
+            m.measurement_id,
+            r.measurement_table_name,
+            s.sensor_name
+            FROM measurements m
+            JOIN application_measurements_registry r
+            ON m.registry_id = r.id
+            JOIN sensors s
+            ON m.sensor_id = s.id
+            ORDER BY m.timestamp DESC
+        """
+        cur.execute(
+            query,
+        )
+
+        rows = cur.fetchall()
+
+        # Step 2: Group measurement_ids by table name
+        table_to_ids = defaultdict(list)
+        for measurement_id, table_name, sensor_name in rows:
+            print(
+                f"Sensor name in DB: {sensor_name}, measurement_id: {measurement_id}, table_name: {table_name}"
+            )
+            table_to_ids[table_name].append((measurement_id, sensor_name))
+
+        locations = []
+
+        # Step 3: Query each table for its measurement_ids
+        for table_name, info in table_to_ids.items():
+            measurement_ids, sensor_names = zip(*info)
+            ids = list(measurement_ids)
+            print(f"IDs: {ids}")
+            # Check if latitude and longitude columns exist in the table
+            cur.execute(
+                sql.SQL(
+                    "SELECT column_name FROM information_schema.columns WHERE table_name = %s AND column_name IN ('latitude', 'longitude')"
+                ),
+                (table_name,),
+            )
+            columns_present = {row[0] for row in cur.fetchall()}
+            if not {"latitude", "longitude"}.issubset(columns_present):
+                continue
+            select_fields = sql.SQL("id, timestamp, latitude, longitude")
+            query = sql.SQL(
+                "SELECT {} FROM {} WHERE id = ANY(%s) ORDER BY timestamp DESC"
+            ).format(select_fields, sql.Identifier(table_name))
+            # print(f"Executing query on table {table_name} for IDs: {ids}")
+            # print(f"Query: {query.as_string(conn)}")
+
+            try:
+                cur.execute(query, (ids,))
+                columns = [desc[0] for desc in cur.description]
+                rows = cur.fetchall()
+                for row in rows:
+                    data = {col: val for col, val in zip(columns[1:], row[1:])}
+                    filtered_data = {k: v for k, v in data.items() if v is not None}
+
+                    # sensor_name = sensor_names[ids.index(row[0])]
+                    filtered_data["sensor_name"] = sensor_name
+                    print(f"Sensor name: {sensor_name}")
+                    if filtered_data:
+                        if (
+                            filtered_data.get("latitude") is not None
+                            and filtered_data.get("longitude") is not None
+                        ):
+                            locations.append(
+                                Measurement(
+                                    table_name=table_name,
+                                    timestamp=str(row[1]),
+                                    data=filtered_data,
+                                )
+                            )
+                            break
+            except Exception as e:
+                print(f"Error executing query on table {table_name}: {e}")
+
+        return locations
+
+
 def fetch_sensors_for_application(conn, app_name: str) -> list | None:
     """Fetch all sensors for a given application name."""
     with conn.cursor() as cur: