Merge pull request #47 from freol35241/sqlalchemy2-support

WIP: Updating codebase for breaking changes introduced by bumping the sqlalchemy version to

README.md

@@ -1,6 +1,8 @@
 Long time state storage (LTSS) custom component for Home Assistant
 ========================================
 
+**NOTE:** From version 2.0 LTSS requires at least Home Assistant 2023.3
+
 **NOTE:** Starting 2020-09-13 attributes are stored with type JSONB instead of as a plain string, in addition a GIN index is created on this column by default. At first startup after updating of LTSS, migration of your DB happens automatically. Note that this can take a couple of minutes and HASS will not finish starting (i.e. frontend will not be available) until migration is done.
 
 **WARNING:** I take no responsibility for any data loss that may happen as a result of this. Please make sure to backup your data before upgrading!

custom_components/ltss/__init__.py

@@ -26,7 +26,7 @@
     EVENT_HOMEASSISTANT_START,
     EVENT_HOMEASSISTANT_STOP,
     EVENT_STATE_CHANGED,
-    STATE_UNKNOWN
+    STATE_UNKNOWN,
 )
 from homeassistant.components import persistent_notification
 from homeassistant.core import CoreState, HomeAssistant, callback
@@ -56,7 +56,9 @@
         DOMAIN: INCLUDE_EXCLUDE_BASE_FILTER_SCHEMA.extend(
             {
                 vol.Required(CONF_DB_URL): cv.string,
-                vol.Optional(CONF_CHUNK_TIME_INTERVAL, default=2592000000000): cv.positive_int,  # 30 days
+                vol.Optional(
+                    CONF_CHUNK_TIME_INTERVAL, default=2592000000000
+                ): cv.positive_int,  # 30 days
             }
         )
     },
@@ -84,37 +86,15 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return await instance.async_db_ready
 
 
-@contextmanager
-def session_scope(*, session=None):
-    """Provide a transactional scope around a series of operations."""
-
-    if session is None:
-        raise RuntimeError("Session required")
-
-    need_rollback = False
-    try:
-        yield session
-        if session.transaction:
-            need_rollback = True
-            session.commit()
-    except Exception as err:  # pylint: disable=broad-except
-        _LOGGER.error("Error executing query: %s", err)
-        if need_rollback:
-            session.rollback()
-        raise
-    finally:
-        session.close()
-
-
 class LTSS_DB(threading.Thread):
     """A threaded LTSS class."""
 
     def __init__(
-            self,
-            hass: HomeAssistant,
-            uri: str,
-            chunk_time_interval: int,
-            entity_filter: Callable[[str], bool],
+        self,
+        hass: HomeAssistant,
+        uri: str,
+        chunk_time_interval: int,
+        entity_filter: Callable[[str], bool],
     ) -> None:
         """Initialize the ltss."""
         threading.Thread.__init__(self, name="LTSS")
@@ -158,6 +138,7 @@ def run(self):
                 tries += 1
 
         if not connected:
+
             @callback
             def connection_failed():
                 """Connect failed tasks."""
@@ -222,17 +203,18 @@ def notify_hass_started(event):
                 if tries != 1:
                     time.sleep(CONNECT_RETRY_WAIT)
                 try:
-                    with session_scope(session=self.get_session()) as session:
-                        try:
-                            row = LTSS.from_event(event)
-                            session.add(row)
-                        except (TypeError, ValueError):
-                            _LOGGER.warning(
-                                "State is not JSON serializable: %s",
-                                event.data.get("new_state"),
-                            )
-
-                    updated = True
+                    with self.get_session() as session:
+                        with session.begin():
+                            try:
+                                row = LTSS.from_event(event)
+                                session.add(row)
+                            except (TypeError, ValueError):
+                                _LOGGER.warning(
+                                    "State is not JSON serializable: %s",
+                                    event.data.get("new_state"),
+                                )
+
+                        updated = True
 
                 except exc.OperationalError as err:
                     _LOGGER.error(
@@ -273,36 +255,49 @@ def _setup_connection(self):
         if self.engine is not None:
             self.engine.dispose()
 
-        self.engine = create_engine(self.db_url, echo=False,
-                                    json_serializer=lambda obj: json.dumps(obj, cls=JSONEncoder))
+        self.engine = create_engine(
+            self.db_url,
+            echo=False,
+            json_serializer=lambda obj: json.dumps(obj, cls=JSONEncoder),
+        )
 
         inspector = inspect(self.engine)
 
         with self.engine.connect() as con:
-            available_extensions = {row['name']: row['installed_version'] for row in
-                                    con.execute(text("SELECT name, installed_version FROM pg_available_extensions"))}
+            con = con.execution_options(isolation_level="AUTOCOMMIT")
+            available_extensions = {
+                row.name: row.installed_version
+                for row in con.execute(
+                    text("SELECT name, installed_version FROM pg_available_extensions")
+                )
+            }
 
             # create table if necessary
             if not inspector.has_table(LTSS.__tablename__):
                 self._create_table(available_extensions)
 
-            if 'timescaledb' in available_extensions:
+            if "timescaledb" in available_extensions:
                 # chunk_time_interval can be adjusted even after first setup
                 try:
                     con.execute(
-                        text(f"SELECT set_chunk_time_interval('{LTSS.__tablename__}', {self.chunk_time_interval})")
-                            .execution_options(autocommit=True)
+                        text(
+                            f"SELECT set_chunk_time_interval('{LTSS.__tablename__}', {self.chunk_time_interval})"
+                        )
                     )
                 except exc.ProgrammingError as exception:
                     if isinstance(exception.orig, psycopg2.errors.UndefinedTable):
                         # The table does exist but is not a hypertable, not much we can do except log that fact
                         _LOGGER.exception(
-                            "TimescaleDB is available as an extension but the LTSS table is not a hypertable!")
+                            "TimescaleDB is available as an extension but the LTSS table is not a hypertable!"
+                        )
                     else:
                         raise
 
         # check if table has been set up with location extraction
-        if "location" in [column_conf["name"] for column_conf in inspector.get_columns(LTSS.__tablename__)]:
+        if "location" in [
+            column_conf["name"]
+            for column_conf in inspector.get_columns(LTSS.__tablename__)
+        ]:
             # activate location extraction in model/ORM
             LTSS.activate_location_extraction()
 
@@ -314,28 +309,33 @@ def _setup_connection(self):
     def _create_table(self, available_extensions):
         _LOGGER.info("Creating LTSS table")
         with self.engine.connect() as con:
-            if 'postgis' in available_extensions:
-                _LOGGER.info("PostGIS extension is available, activating location extraction...")
-                con.execute(
-                    text("CREATE EXTENSION IF NOT EXISTS postgis CASCADE"
-                         ).execution_options(autocommit=True))
+            con = con.execution_options(isolation_level="AUTOCOMMIT")
+            if "postgis" in available_extensions:
+                _LOGGER.info(
+                    "PostGIS extension is available, activating location extraction..."
+                )
+                con.execute(text("CREATE EXTENSION IF NOT EXISTS postgis CASCADE"))
 
                 # activate location extraction in model/ORM to add necessary column when calling create_all()
                 LTSS.activate_location_extraction()
 
             Base.metadata.create_all(self.engine)
 
-            if 'timescaledb' in available_extensions:
-                _LOGGER.info("TimescaleDB extension is available, creating hypertable...")
-                con.execute(
-                    text("CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE"
-                         ).execution_options(autocommit=True))
+            if "timescaledb" in available_extensions:
+                _LOGGER.info(
+                    "TimescaleDB extension is available, creating hypertable..."
+                )
+                con.execute(text("CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE"))
 
                 # Create hypertable
-                con.execute(text(f"""SELECT create_hypertable(
-                                '{LTSS.__tablename__}', 
-                                'time', 
-                                if_not_exists => TRUE);""").execution_options(autocommit=True))
+                con.execute(
+                    text(
+                        f"""SELECT create_hypertable(
+                                '{LTSS.__tablename__}',
+                                'time',
+                                if_not_exists => TRUE);"""
+                    )
+                )
 
     def _close_connection(self):
         """Close the connection."""

custom_components/ltss/manifest.json

@@ -1,15 +1,15 @@
 {
-    "domain": "ltss",
-    "version": "1.1.0",
-    "name": "Long Time State Storage (LTSS)",
-    "documentation": "https://github.com/freol35241/ltss",
-    "requirements": [
-      "sqlalchemy>=1.0,<2.0",
-      "psycopg2>=2.8,<3.0",
-      "geoalchemy2>=0.8,<0.9"
-    ],
-    "dependencies": [],
-    "codeowners": [
-      "@freol35241"
-    ]
-  }
+  "domain": "ltss",
+  "version": "2.0.0",
+  "name": "Long Time State Storage (LTSS)",
+  "documentation": "https://github.com/freol35241/ltss",
+  "requirements": [
+    "sqlalchemy>=2.0,<3.0",
+    "psycopg2>=2.8,<3.0",
+    "geoalchemy2>=0.8,<1.0"
+  ],
+  "dependencies": [],
+  "codeowners": [
+    "@freol35241"
+  ]
+}

custom_components/ltss/models.py

@@ -14,8 +14,7 @@
 from sqlalchemy.schema import Index
 from sqlalchemy.dialects.postgresql import JSONB
 from geoalchemy2 import Geometry
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import column_property
+from sqlalchemy.orm import column_property, declarative_base
 
 # SQLAlchemy Schema
 # pylint: disable=invalid-name

info.md

@@ -1,3 +1,5 @@
+**NOTE:** From version 2.0 LTSS requires at least Home Assistant 2023.3
+
 **NOTE:** Starting 2020-09-13 attributes are stored with type JSONB instead of as a plain string, in addition a GIN index is created on this column by default. At first startup after updating of LTSS, migration of your DB happens automatically. Note that this can take a couple of minutes and HASS will not finish starting (i.e. frontend will not be available) until migration is done.
 
 **WARNING:** I take no responsibility for any data loss that may happen as a result of this. Please make sure to backup your data before upgrading!

tests/requirements.txt

@@ -1,7 +1,8 @@
 docker==5.0.3
 GeoAlchemy2==0.11.1
-homeassistant==2021.9.7
 psycopg2==2.9.3
+SQLAlchemy==2.0.4
+homeassistant==2021.9.7
+
 pytest-docker==0.11.0
 pytest==7.1.1
-SQLAlchemy==1.4.32

tests/test_databases.py

@@ -88,9 +88,9 @@ def test_default_db(self):
 
     @staticmethod
     def _is_hypertable(con):
-        timescaledb_version = con.execute("SELECT installed_version "
+        timescaledb_version = con.execute(text("SELECT installed_version "
                                           "FROM pg_available_extensions "
-                                          "WHERE name = 'timescaledb'").scalar()
+                                          "WHERE name = 'timescaledb'")).scalar()
 
         # timescaledb's table/column name changed with v2
         if int(timescaledb_version.split('.')[0]) >= 2:
@@ -100,7 +100,7 @@ def _is_hypertable(con):
             query = f"SELECT 1 FROM timescaledb_information.hypertable "
             f"WHERE table_name = '{LTSS.__tablename__}'"
 
-        return 1 == con.execute(query).rowcount
+        return 1 == con.execute(text(query)).rowcount
 
     @staticmethod
     def _has_columns(con):