working 2.0 clone

custom_components/ltss/models.py

@@ -0,0 +1,79 @@
+"""Models for SQLAlchemy."""
+import json
+from datetime import datetime
+import logging
+
+from sqlalchemy import (
+    Column,
+    BigInteger,
+    DateTime,
+    String,
+    Text,
+)
+
+from sqlalchemy.schema import Index
+from sqlalchemy.dialects.postgresql import JSONB
+from geoalchemy2 import Geometry
+from sqlalchemy.orm import column_property, declarative_base
+
+# SQLAlchemy Schema
+# pylint: disable=invalid-name
+Base = declarative_base()
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class LTSS(Base):  # type: ignore
+    """State change history."""
+
+    __tablename__ = "ltss"
+    time = Column(DateTime(timezone=True), default=datetime.utcnow, primary_key=True)
+    entity_id = Column(String(255), primary_key=True)
+    state = Column(String(255), index=True)
+    attributes = Column(JSONB)
+    location = None  # when not activated, no location column will be added to the table/database
+
+    @classmethod
+    def activate_location_extraction(cls):
+        """
+        Method to activate proper Postgis handling of location.
+
+        After activation, this cannot be deactivated (due to how the underlying SQLAlchemy ORM works).
+        """
+        cls.location = column_property(Column(Geometry("POINT", srid=4326)))
+
+    @classmethod
+    def from_event(cls, event):
+        """Create object from a state_changed event."""
+        entity_id = event.data["entity_id"]
+        state = event.data.get("new_state")
+
+        attrs = dict(state.attributes)
+
+        location = None
+
+        if (
+            cls.location
+        ):  # if the additional column exists, use Postgis' Geometry/Point data structure
+            lat = attrs.pop("latitude", None)
+            lon = attrs.pop("longitude", None)
+
+            location = f"SRID=4326;POINT({lon} {lat})" if lon and lat else None
+
+        row = LTSS(
+            entity_id=entity_id,
+            time=event.time_fired,
+            state=state.state,
+            attributes=attrs,
+            location=location,
+        )
+
+        return row
+
+
+LTSS_attributes_index = Index(
+    "ltss_attributes_idx", LTSS.attributes, postgresql_using="gin"
+)
+LTSS_entityid_time_composite_index = Index(
+    "ltss_entityid_time_composite_idx", LTSS.entity_id, LTSS.time.desc()
+)