109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
|
import logging
|
||
|
|
import os
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
try:
|
||
|
|
from cassandra.cluster import Cluster
|
||
|
|
from cassandra.query import dict_factory
|
||
|
|
except Exception:
|
||
|
|
Cluster = None
|
||
|
|
dict_factory = None
|
||
|
|
|
||
|
|
|
||
|
|
class CassandraService:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.enabled = os.environ.get("CASSANDRA_ENABLED", "0") == "1"
|
||
|
|
self.contact_points = [
|
||
|
|
item.strip()
|
||
|
|
for item in os.environ.get("CASSANDRA_HOSTS", "cassandra").split(",")
|
||
|
|
if item.strip()
|
||
|
|
]
|
||
|
|
self.port = int(os.environ.get("CASSANDRA_PORT", "9042"))
|
||
|
|
self.keyspace = os.environ.get("CASSANDRA_KEYSPACE", "sensor_hub")
|
||
|
|
self.replication = os.environ.get(
|
||
|
|
"CASSANDRA_REPLICATION",
|
||
|
|
"{'class': 'SimpleStrategy', 'replication_factor': 1}",
|
||
|
|
)
|
||
|
|
self._session = None
|
||
|
|
|
||
|
|
def is_available(self) -> bool:
|
||
|
|
return self.enabled and Cluster is not None
|
||
|
|
|
||
|
|
def connect(self):
|
||
|
|
if self._session is not None:
|
||
|
|
return self._session
|
||
|
|
if not self.is_available():
|
||
|
|
return None
|
||
|
|
|
||
|
|
cluster = Cluster(contact_points=self.contact_points, port=self.port)
|
||
|
|
session = cluster.connect()
|
||
|
|
if dict_factory is not None:
|
||
|
|
session.row_factory = dict_factory
|
||
|
|
session.execute(
|
||
|
|
f"CREATE KEYSPACE IF NOT EXISTS {self.keyspace} "
|
||
|
|
f"WITH replication = {self.replication}"
|
||
|
|
)
|
||
|
|
session.set_keyspace(self.keyspace)
|
||
|
|
session.execute(
|
||
|
|
"""
|
||
|
|
CREATE TABLE IF NOT EXISTS sensor_payloads (
|
||
|
|
sensor_uuid text,
|
||
|
|
reading_at timestamp,
|
||
|
|
payload_id uuid,
|
||
|
|
original_payload text,
|
||
|
|
translated_payload text,
|
||
|
|
translation_status text,
|
||
|
|
created_at timestamp,
|
||
|
|
PRIMARY KEY ((sensor_uuid), reading_at, payload_id)
|
||
|
|
) WITH CLUSTERING ORDER BY (reading_at DESC, payload_id DESC)
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
self._session = session
|
||
|
|
return self._session
|
||
|
|
|
||
|
|
def save_payload(self, *, sensor_uuid: str, payload_id: Any, reading_at, original_payload: str, translated_payload: str, translation_status: str, created_at):
|
||
|
|
session = self.connect()
|
||
|
|
if session is None:
|
||
|
|
logger.warning("Cassandra is disabled or unavailable; skipping payload persistence.")
|
||
|
|
return False
|
||
|
|
|
||
|
|
session.execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO sensor_payloads (
|
||
|
|
sensor_uuid,
|
||
|
|
reading_at,
|
||
|
|
payload_id,
|
||
|
|
original_payload,
|
||
|
|
translated_payload,
|
||
|
|
translation_status,
|
||
|
|
created_at
|
||
|
|
) VALUES (%s, %s, %s, %s, %s, %s, %s)
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
sensor_uuid,
|
||
|
|
reading_at,
|
||
|
|
payload_id,
|
||
|
|
original_payload,
|
||
|
|
translated_payload,
|
||
|
|
translation_status,
|
||
|
|
created_at,
|
||
|
|
),
|
||
|
|
)
|
||
|
|
return True
|
||
|
|
|
||
|
|
def get_payloads(self, sensor_uuid: str, limit: int = 50):
|
||
|
|
session = self.connect()
|
||
|
|
if session is None:
|
||
|
|
return []
|
||
|
|
rows = session.execute(
|
||
|
|
"SELECT sensor_uuid, reading_at, payload_id, original_payload, translated_payload, translation_status, created_at "
|
||
|
|
"FROM sensor_payloads WHERE sensor_uuid = %s LIMIT %s",
|
||
|
|
(sensor_uuid, limit),
|
||
|
|
)
|
||
|
|
return list(rows)
|
||
|
|
|
||
|
|
|
||
|
|
cassandra_service = CassandraService()
|