implement SQLAlchemySubscription

Author: mobuild-io

eventsourcing_sqlalchemy/datastore.py

@@ -2,10 +2,13 @@
 from __future__ import annotations
 
 import sqlite3
+from contextlib import contextmanager
 from contextvars import ContextVar, Token
 from threading import Lock, Semaphore
-from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union, cast
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union, cast
 
+import psycopg
+import psycopg2
 import sqlalchemy.exc
 from eventsourcing.persistence import (
     DatabaseError,
@@ -23,6 +26,7 @@
 from sqlalchemy.future import create_engine
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 from sqlalchemy.pool import StaticPool
+from typing_extensions import TypeVar
 
 from eventsourcing_sqlalchemy.models import (  # type: ignore
     EventRecord,
@@ -263,3 +267,30 @@ def define_record_class(
             )
             cls.record_classes[record_classes_key] = (record_class, base_cls)
         return cast(Type[TEventRecord], record_class)
+
+    @contextmanager
+    def get_connection(self) -> Iterator[Connection]:
+        try:
+            assert self.engine
+            conn = self.engine.connect()
+            yield conn
+        except (psycopg.InterfaceError, psycopg2.InterfaceError) as e:
+            raise InterfaceError(str(e)) from e
+        except (psycopg.OperationalError, psycopg2.OperationalError) as e:
+            raise OperationalError(str(e)) from e
+        except (psycopg.DataError, psycopg2.DataError) as e:
+            raise DataError(str(e)) from e
+        except (psycopg.IntegrityError, psycopg2.IntegrityError) as e:
+            raise IntegrityError(str(e)) from e
+        except (psycopg.InternalError, psycopg2.InternalError) as e:
+            raise InternalError(str(e)) from e
+        except (psycopg.ProgrammingError, psycopg2.ProgrammingError) as e:
+            raise ProgrammingError(str(e)) from e
+        except (psycopg.NotSupportedError, psycopg2.NotSupportedError) as e:
+            raise NotSupportedError(str(e)) from e
+        except (psycopg.DatabaseError, psycopg2.DatabaseError) as e:
+            raise DatabaseError(str(e)) from e
+        except (psycopg.Error, psycopg2.Error) as e:
+            raise PersistenceError(str(e)) from e
+        except Exception:
+            raise

eventsourcing_sqlalchemy/recorders.py

@@ -1,15 +1,20 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-from typing import Any, List, Optional, Sequence, Type, cast
+import select
+import time
+from threading import Thread
+from typing import Any, Callable, List, Optional, Sequence, Type, cast
 from uuid import UUID
 
 from eventsourcing.persistence import (
     AggregateRecorder,
     ApplicationRecorder,
     IntegrityError,
+    ListenNotifySubscription,
     Notification,
     ProcessRecorder,
+    ProgrammingError,
     StoredEvent,
     Subscription,
     Tracking,
@@ -29,6 +34,8 @@
 class SQLAlchemyRecorder:
     """Base class for recorders that use SQLAlchemy."""
 
+    POSTGRES_MAX_IDENTIFIER_LEN = 63
+
     def __init__(
         self,
         datastore: SQLAlchemyDatastore,
@@ -38,6 +45,13 @@ def __init__(
         self.schema_name = schema_name
         self.tables: List[Table] = []
 
+    def check_identifier_length(self, table_name: str) -> None:
+        assert self.datastore.engine is not None
+        if self.datastore.engine.dialect.name == "postgresql":
+            if len(table_name) > SQLAlchemyRecorder.POSTGRES_MAX_IDENTIFIER_LEN:
+                msg = f"Identifier too long: {table_name}"
+                raise ProgrammingError(msg)
+
     def create_table(self) -> None:
         assert self.datastore.engine is not None
         for table in self.tables:
@@ -233,12 +247,12 @@ def _insert_stored_events(
             session.add(record)
         if self._has_autoincrementing_ids:
             session.flush()  # We want the autoincremented IDs now.
+        self._notify_channel(session)
         return [cast(StoredEventRecord, r).id for r in records]
 
     def max_notification_id(self) -> int | None:
         try:
             with self.transaction(commit=False) as session:
-                # record_class = cast(Type[StoredEventRecord], self.events_record_cls)
                 record_class = self.events_record_cls
                 q = session.query(record_class)
                 q = q.order_by(record_class.id.desc())
@@ -257,7 +271,6 @@ def select_notifications(
         inclusive_of_start: bool = True,
     ) -> list[Notification]:
         with self.transaction(commit=False) as session:
-            # record_class = cast(Type[StoredEventRecord], self.events_record_cls)
             record_class = self.events_record_cls
             q = session.query(record_class)
             if start is not None:
@@ -289,8 +302,86 @@ def select_notifications(
     def subscribe(
         self, gt: int | None = None, topics: Sequence[str] = ()
     ) -> Subscription[ApplicationRecorder]:
-        msg = "SQLAlchemyApplicationRecorder.subscribe() is not implemented"
-        raise NotImplementedError(msg)
+        assert self.datastore.engine
+        if self.datastore.engine.dialect.name == "postgresql":
+            return SQLAlchemySubscription(recorder=self, gt=gt, topics=topics)
+        else:
+            msg = "SQLAlchemyApplicationRecorder.subscribe() is not implemented for"
+            msg += f"{self.datastore.engine.dialect}"
+            raise NotImplementedError(msg)
+
+    def _notify_channel(self, session: Session) -> None:
+        """
+        Send a NOTIFY on the channel using a SQLAlchemy connection.
+        """
+        assert self.datastore.engine
+        if self.datastore.engine.dialect.name == "postgresql":
+            # Get the raw psycopg connection
+            cursor = session.connection().connection.cursor()
+            cursor.execute(f"NOTIFY {self.channel_name};")
+
+
+class SQLAlchemySubscription(ListenNotifySubscription[SQLAlchemyApplicationRecorder]):
+    def __init__(
+        self,
+        recorder: SQLAlchemyApplicationRecorder,
+        gt: int | None = None,
+        topics: Sequence[str] = (),
+    ) -> None:
+        assert isinstance(recorder, SQLAlchemyApplicationRecorder)
+        super().__init__(recorder=recorder, gt=gt, topics=topics)
+        self._listen_thread = Thread(target=self._listen)
+        self._listen_thread.start()
+
+    def __exit__(self, *args: object, **kwargs: Any) -> None:
+        super().__exit__(*args, **kwargs)
+        self._listen_thread.join()
+
+    def _listen(self) -> None:
+        assert self._recorder.datastore.engine
+        assert self._recorder.datastore.engine.dialect.name == "postgresql"
+        notification_handler = self.__get_notification_handler()
+
+        try:
+            with self._recorder.datastore.get_connection() as sa_conn:
+                sa_conn.execution_options(isolation_level="AUTOCOMMIT")
+                raw_conn = sa_conn.connection
+
+                cur = raw_conn.cursor()
+                cur.execute(f"LISTEN {self._recorder.channel_name};")
+
+                while not self._has_been_stopped and not self._thread_error:
+                    if select.select([raw_conn], [], [], 0.1)[0]:
+                        notification_handler(raw_conn)
+                    else:
+                        time.sleep(0.1)
+
+        except BaseException as e:  # noqa: B036
+            if self._thread_error is None:
+                self._thread_error = e
+            self.stop()
+
+    def __get_notification_handler(self) -> Callable[[Any], None]:
+        assert self._recorder.datastore.engine
+        driver_name = self._recorder.datastore.engine.dialect.driver
+        handlers = {
+            "psycopg": self.__handle_psycopg_notification,
+            "psycopg2": self.__handle_psycopg2_notification,
+        }
+        try:
+            return handlers[driver_name]
+        except KeyError as e:
+            raise NotImplementedError(f"Unsupported driver: {driver_name}") from e
+
+    def __handle_psycopg_notification(self, raw_conn: Any) -> None:
+        next(raw_conn.notifies())
+        self._has_been_notified.set()
+
+    def __handle_psycopg2_notification(self, raw_conn: Any) -> None:
+        raw_conn.poll()
+        if raw_conn.notifies:
+            raw_conn.notifies.pop(0)
+            self._has_been_notified.set()
 
 
 class SQLAlchemyTrackingRecorder(SQLAlchemyRecorder, TrackingRecorder):

poetry.lock

@@ -40,6 +40,7 @@ flake8-coding = "*"
 flake8-isort = "*"
 flake8-tidy-imports = "*"
 isort = "*"
+msgspec = "*"
 mypy = "*"
 pre-commit = "*"
 pre-commit-hooks = "*"