eventsourcing 9.3.5__py3-none-any.whl → 9.4.0a2__py3-none-any.whl

eventsourcing/postgres.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
 import logging
+from asyncio import CancelledError
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Callable, Iterator, List, Sequence
+from threading import Thread
+from typing import TYPE_CHECKING, Any, Callable, Iterator, List, Sequence, Type
 
 import psycopg
 import psycopg.errors
 import psycopg_pool
-from psycopg import Connection, Cursor
+from psycopg import Connection, Cursor, Error
+from psycopg.generators import notifies
 from psycopg.rows import DictRow, dict_row
 
 from eventsourcing.persistence import (
@@ -19,6 +22,7 @@ from eventsourcing.persistence import (
     IntegrityError,
     InterfaceError,
     InternalError,
+    ListenNotifySubscription,
     Notification,
     NotSupportedError,
     OperationalError,
@@ -27,10 +31,11 @@ from eventsourcing.persistence import (
     ProgrammingError,
     StoredEvent,
     Tracking,
+    TrackingRecorder,
 )
 from eventsourcing.utils import Environment, resolve_topic, retry, strtobool
 
-if TYPE_CHECKING:  # pragma: nocover
+if TYPE_CHECKING:
     from uuid import UUID
 
     from typing_extensions import Self
@@ -38,6 +43,14 @@ if TYPE_CHECKING:  # pragma: nocover
 logging.getLogger("psycopg.pool").setLevel(logging.CRITICAL)
 logging.getLogger("psycopg").setLevel(logging.CRITICAL)
 
+# Copy of "private" psycopg.errors._NO_TRACEBACK (in case it changes)
+# From psycopg: "Don't show a complete traceback upon raising these exception.
+# Usually the traceback starts from internal functions (for instance in the
+# server communication callbacks) but, for the end user, it's more important
+# to get the high level information about where the exception was raised, for
+# instance in a certain `Cursor.execute()`."
+NO_TRACEBACK = (Error, KeyboardInterrupt, CancelledError)
+
 
 class ConnectionPool(psycopg_pool.ConnectionPool[Any]):
     def __init__(
@@ -60,7 +73,7 @@ class PostgresDatastore:
         self,
         dbname: str,
         host: str,
-        port: str,
+        port: str | int,
         user: str,
         password: str,
         *,
@@ -95,7 +108,7 @@ class PostgresDatastore:
             min_size=pool_size,
             max_size=pool_size + max_overflow,
             open=False,
-            configure=self.after_connect,
+            configure=self.after_connect_func(),
             timeout=connect_timeout,
             max_waiting=max_waiting,
             max_lifetime=conn_max_age,
@@ -104,13 +117,18 @@ class PostgresDatastore:
         self.lock_timeout = lock_timeout
         self.schema = schema.strip()
 
-    def after_connect(self, conn: Connection[DictRow]) -> None:
-        conn.autocommit = True
-        conn.cursor().execute(
+    def after_connect_func(self) -> Callable[[Connection[Any]], None]:
+        statement = (
             "SET idle_in_transaction_session_timeout = "
             f"'{self.idle_in_transaction_session_timeout}s'"
         )
 
+        def after_connect(conn: Connection[DictRow]) -> None:
+            conn.autocommit = True
+            conn.cursor().execute(statement)
+
+        return after_connect
+
     @contextmanager
     def get_connection(self) -> Iterator[Connection[DictRow]]:
         try:
@@ -154,9 +172,6 @@ class PostgresDatastore:
     def close(self) -> None:
         self.pool.close()
 
-    def __del__(self) -> None:
-        self.close()
-
     def __enter__(self) -> Self:
         return self
 
@@ -164,33 +179,18 @@ class PostgresDatastore:
         self.close()
 
 
-class PostgresAggregateRecorder(AggregateRecorder):
+class PostgresRecorder:
+    """Base class for recorders that use PostgreSQL."""
+
     def __init__(
         self,
         datastore: PostgresDatastore,
-        events_table_name: str,
     ):
-        self.check_table_name_length(events_table_name, datastore.schema)
         self.datastore = datastore
-        self.events_table_name = events_table_name
-        # Index names can't be qualified names, but
-        # are created in the same schema as the table.
-        if "." in self.events_table_name:
-            unqualified_table_name = self.events_table_name.split(".")[-1]
-        else:
-            unqualified_table_name = self.events_table_name
-        self.notification_id_index_name = (
-            f"{unqualified_table_name}_notification_id_idx "
-        )
-
         self.create_table_statements = self.construct_create_table_statements()
-        self.insert_events_statement = (
-            f"INSERT INTO {self.events_table_name} VALUES (%s, %s, %s, %s)"
-        )
-        self.select_events_statement = (
-            f"SELECT * FROM {self.events_table_name} WHERE originator_id = %s"
-        )
-        self.lock_table_statements: List[str] = []
+
+    def construct_create_table_statements(self) -> List[str]:
+        return []
 
     @staticmethod
     def check_table_name_length(table_name: str, schema_name: str) -> None:
@@ -203,8 +203,32 @@ class PostgresAggregateRecorder(AggregateRecorder):
             msg = f"Table name too long: {unqualified_table_name}"
             raise ProgrammingError(msg)
 
-    def construct_create_table_statements(self) -> List[str]:
-        statement = (
+    def create_table(self) -> None:
+        with self.datastore.transaction(commit=True) as curs:
+            for statement in self.create_table_statements:
+                curs.execute(statement, prepare=False)
+
+
+class PostgresAggregateRecorder(PostgresRecorder, AggregateRecorder):
+    def __init__(
+        self,
+        datastore: PostgresDatastore,
+        *,
+        events_table_name: str = "stored_events",
+    ):
+        super().__init__(datastore)
+        self.check_table_name_length(events_table_name, datastore.schema)
+        self.events_table_name = events_table_name
+        # Index names can't be qualified names, but
+        # are created in the same schema as the table.
+        if "." in self.events_table_name:
+            unqualified_table_name = self.events_table_name.split(".")[-1]
+        else:
+            unqualified_table_name = self.events_table_name
+        self.notification_id_index_name = (
+            f"{unqualified_table_name}_notification_id_idx "
+        )
+        self.create_table_statements.append(
             "CREATE TABLE IF NOT EXISTS "
             f"{self.events_table_name} ("
             "originator_id uuid NOT NULL, "
@@ -215,12 +239,14 @@ class PostgresAggregateRecorder(AggregateRecorder):
             "(originator_id, originator_version)) "
             "WITH (autovacuum_enabled=false)"
         )
-        return [statement]
 
-    def create_table(self) -> None:
-        with self.datastore.transaction(commit=True) as curs:
-            for statement in self.create_table_statements:
-                curs.execute(statement, prepare=False)
+        self.insert_events_statement = (
+            f"INSERT INTO {self.events_table_name} VALUES (%s, %s, %s, %s)"
+        )
+        self.select_events_statement = (
+            f"SELECT * FROM {self.events_table_name} WHERE originator_id = %s"
+        )
+        self.lock_table_statements: List[str] = []
 
     @retry((InterfaceError, OperationalError), max_attempts=10, wait=0.2)
     def insert_events(
@@ -257,7 +283,7 @@ class PostgresAggregateRecorder(AggregateRecorder):
         self,
         c: Cursor[DictRow],
         stored_events: List[StoredEvent],
-        **kwargs: Any,
+        **_: Any,
     ) -> None:
         pass
 
@@ -271,6 +297,8 @@ class PostgresAggregateRecorder(AggregateRecorder):
         if len(stored_events) > 0:
             self._lock_table(c)
 
+            self._notify_channel(c)
+
             # Insert events.
             c.executemany(
                 query=self.insert_events_statement,
@@ -289,6 +317,9 @@ class PostgresAggregateRecorder(AggregateRecorder):
     def _lock_table(self, c: Cursor[DictRow]) -> None:
         pass
 
+    def _notify_channel(self, c: Cursor[DictRow]) -> None:
+        pass
+
     def _fetch_ids_after_insert_events(
         self,
         c: Cursor[DictRow],
@@ -341,9 +372,28 @@ class PostgresApplicationRecorder(PostgresAggregateRecorder, ApplicationRecorder
     def __init__(
         self,
         datastore: PostgresDatastore,
+        *,
         events_table_name: str = "stored_events",
     ):
-        super().__init__(datastore, events_table_name)
+        super().__init__(datastore, events_table_name=events_table_name)
+        self.create_table_statements[-1] = (
+            "CREATE TABLE IF NOT EXISTS "
+            f"{self.events_table_name} ("
+            "originator_id uuid NOT NULL, "
+            "originator_version bigint NOT NULL, "
+            "topic text, "
+            "state bytea, "
+            "notification_id bigserial, "
+            "PRIMARY KEY "
+            "(originator_id, originator_version)) "
+            "WITH (autovacuum_enabled=false)"
+        )
+        self.create_table_statements.append(
+            "CREATE UNIQUE INDEX IF NOT EXISTS "
+            f"{self.notification_id_index_name}"
+            f"ON {self.events_table_name} (notification_id ASC);"
+        )
+        self.channel_name = self.events_table_name.replace(".", "_")
         self.insert_events_statement += " RETURNING notification_id"
         self.max_notification_id_statement = (
             f"SELECT MAX(notification_id) FROM {self.events_table_name}"
@@ -353,50 +403,50 @@ class PostgresApplicationRecorder(PostgresAggregateRecorder, ApplicationRecorder
             f"LOCK TABLE {self.events_table_name} IN EXCLUSIVE MODE",
         ]
 
-    def construct_create_table_statements(self) -> List[str]:
-        return [
-            (
-                "CREATE TABLE IF NOT EXISTS "
-                f"{self.events_table_name} ("
-                "originator_id uuid NOT NULL, "
-                "originator_version bigint NOT NULL, "
-                "topic text, "
-                "state bytea, "
-                "notification_id bigserial, "
-                "PRIMARY KEY "
-                "(originator_id, originator_version)) "
-                "WITH (autovacuum_enabled=false)"
-            ),
-            (
-                "CREATE UNIQUE INDEX IF NOT EXISTS "
-                f"{self.notification_id_index_name}"
-                f"ON {self.events_table_name} (notification_id ASC);"
-            ),
-        ]
-
     @retry((InterfaceError, OperationalError), max_attempts=10, wait=0.2)
     def select_notifications(
         self,
-        start: int,
+        start: int | None,
         limit: int,
         stop: int | None = None,
         topics: Sequence[str] = (),
+        *,
+        inclusive_of_start: bool = True,
     ) -> List[Notification]:
         """
         Returns a list of event notifications
         from 'start', limited by 'limit'.
         """
 
-        params: List[int | str | Sequence[str]] = [start]
-        statement = f"SELECT * FROM {self.events_table_name} WHERE notification_id>=%s"
+        params: List[int | str | Sequence[str]] = []
+        statement = f"SELECT * FROM {self.events_table_name}"
+        has_where = False
+        if start is not None:
+            statement += " WHERE"
+            has_where = True
+            params.append(start)
+            if inclusive_of_start:
+                statement += " notification_id>=%s"
+            else:
+                statement += " notification_id>%s"
 
         if stop is not None:
+            if not has_where:
+                has_where = True
+                statement += " WHERE"
+            else:
+                statement += " AND"
+
             params.append(stop)
-            statement += " AND notification_id <= %s"
+            statement += " notification_id <= %s"
 
         if topics:
+            if not has_where:
+                statement += " WHERE"
+            else:
+                statement += " AND"
             params.append(topics)
-            statement += " AND topic = ANY(%s)"
+            statement += " topic = ANY(%s)"
 
         params.append(limit)
         statement += " ORDER BY notification_id LIMIT %s"
@@ -416,7 +466,7 @@ class PostgresApplicationRecorder(PostgresAggregateRecorder, ApplicationRecorder
             ]
 
     @retry((InterfaceError, OperationalError), max_attempts=10, wait=0.2)
-    def max_notification_id(self) -> int:
+    def max_notification_id(self) -> int | None:
         """
         Returns the maximum notification ID.
         """
@@ -425,7 +475,7 @@ class PostgresApplicationRecorder(PostgresAggregateRecorder, ApplicationRecorder
             curs.execute(self.max_notification_id_statement)
             fetchone = curs.fetchone()
             assert fetchone is not None
-            return fetchone["max"] or 0
+            return fetchone["max"]
 
     def _lock_table(self, c: Cursor[DictRow]) -> None:
         # Acquire "EXCLUSIVE" table lock, to serialize transactions that insert
@@ -451,6 +501,9 @@ class PostgresApplicationRecorder(PostgresAggregateRecorder, ApplicationRecorder
         for lock_statement in self.lock_table_statements:
             c.execute(lock_statement, prepare=True)
 
+    def _notify_channel(self, c: Cursor[DictRow]) -> None:
+        c.execute("NOTIFY " + self.channel_name)
+
     def _fetch_ids_after_insert_events(
         self,
         c: Cursor[DictRow],
@@ -460,31 +513,75 @@ class PostgresApplicationRecorder(PostgresAggregateRecorder, ApplicationRecorder
         notification_ids: List[int] = []
         len_events = len(stored_events)
         if len_events:
-            if (
-                (c.statusmessage == "SET")
-                and c.nextset()
-                and (c.statusmessage == "LOCK TABLE")
-            ):
-                while c.nextset() and len(notification_ids) != len_events:
+            while c.nextset() and len(notification_ids) != len_events:
+                if c.statusmessage and c.statusmessage.startswith("INSERT"):
                     row = c.fetchone()
                     assert row is not None
                     notification_ids.append(row["notification_id"])
             if len(notification_ids) != len(stored_events):
-                msg = "Couldn't get all notification IDs"
+                msg = "Couldn't get all notification IDs "
+                msg += f"(got {len(notification_ids)}, expected {len(stored_events)}"
                 raise ProgrammingError(msg)
         return notification_ids
 
+    def subscribe(self, gt: int | None = None) -> PostgresSubscription:
+        return PostgresSubscription(self, gt)
+
+
+class PostgresSubscription(ListenNotifySubscription[PostgresApplicationRecorder]):
+    def __init__(
+        self, recorder: PostgresApplicationRecorder, gt: int | None = None
+    ) -> None:
+        assert isinstance(recorder, PostgresApplicationRecorder)
+        super().__init__(recorder=recorder, gt=gt)
+        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()
 
-class PostgresProcessRecorder(PostgresApplicationRecorder, ProcessRecorder):
+    def _listen(self) -> None:
+        try:
+            with self._recorder.datastore.get_connection() as conn:
+                conn.execute("LISTEN " + self._recorder.channel_name)
+                while not self._has_been_stopped and not self._thread_error:
+                    # This block simplifies psycopg's conn.notifies(), because
+                    # we aren't interested in the actual notify messages, and
+                    # also we want to stop consuming notify messages when the
+                    # subscription has an error or is otherwise stopped.
+                    with conn.lock:
+                        try:
+                            if conn.wait(notifies(conn.pgconn), interval=0.1):
+                                self._has_been_notified.set()
+                        except NO_TRACEBACK as ex:  # pragma: no cover
+                            raise ex.with_traceback(None) from None
+
+        except BaseException as e:
+            if self._thread_error is None:
+                self._thread_error = e
+            self.stop()
+
+
+class PostgresTrackingRecorder(PostgresRecorder, TrackingRecorder):
     def __init__(
         self,
         datastore: PostgresDatastore,
-        events_table_name: str,
-        tracking_table_name: str,
+        *,
+        tracking_table_name: str = "notification_tracking",
+        **kwargs: Any,
     ):
+        super().__init__(datastore, **kwargs)
         self.check_table_name_length(tracking_table_name, datastore.schema)
         self.tracking_table_name = tracking_table_name
-        super().__init__(datastore, events_table_name)
+        self.create_table_statements.append(
+            "CREATE TABLE IF NOT EXISTS "
+            f"{self.tracking_table_name} ("
+            "application_name text, "
+            "notification_id bigint, "
+            "PRIMARY KEY "
+            "(application_name, notification_id))"
+        )
         self.insert_tracking_statement = (
             f"INSERT INTO {self.tracking_table_name} VALUES (%s, %s)"
         )
@@ -499,20 +596,28 @@ class PostgresProcessRecorder(PostgresApplicationRecorder, ProcessRecorder):
             "WHERE application_name=%s AND notification_id=%s"
         )
 
-    def construct_create_table_statements(self) -> List[str]:
-        statements = super().construct_create_table_statements()
-        statements.append(
-            "CREATE TABLE IF NOT EXISTS "
-            f"{self.tracking_table_name} ("
-            "application_name text, "
-            "notification_id bigint, "
-            "PRIMARY KEY "
-            "(application_name, notification_id))"
+    @retry((InterfaceError, OperationalError), max_attempts=10, wait=0.2)
+    def insert_tracking(self, tracking: Tracking) -> None:
+        c: Connection[DictRow]
+        with self.datastore.get_connection() as c, c.transaction(), c.cursor() as curs:
+            self._insert_tracking(curs, tracking)
+
+    def _insert_tracking(
+        self,
+        c: Cursor[DictRow],
+        tracking: Tracking,
+    ) -> None:
+        c.execute(
+            query=self.insert_tracking_statement,
+            params=(
+                tracking.application_name,
+                tracking.notification_id,
+            ),
+            prepare=True,
         )
-        return statements
 
     @retry((InterfaceError, OperationalError), max_attempts=10, wait=0.2)
-    def max_tracking_id(self, application_name: str) -> int:
+    def max_tracking_id(self, application_name: str) -> int | None:
         with self.datastore.get_connection() as conn, conn.cursor() as curs:
             curs.execute(
                 query=self.max_tracking_id_statement,
@@ -521,7 +626,7 @@ class PostgresProcessRecorder(PostgresApplicationRecorder, ProcessRecorder):
             )
             fetchone = curs.fetchone()
             assert fetchone is not None
-            return fetchone["max"] or 0
+            return fetchone["max"]
 
     @retry((InterfaceError, OperationalError), max_attempts=10, wait=0.2)
     def has_tracking_id(self, application_name: str, notification_id: int) -> bool:
@@ -536,6 +641,23 @@ class PostgresProcessRecorder(PostgresApplicationRecorder, ProcessRecorder):
             assert fetchone is not None
             return bool(fetchone["count"])
 
+
+class PostgresProcessRecorder(
+    PostgresTrackingRecorder, PostgresApplicationRecorder, ProcessRecorder
+):
+    def __init__(
+        self,
+        datastore: PostgresDatastore,
+        *,
+        events_table_name: str = "stored_events",
+        tracking_table_name: str = "notification_tracking",
+    ):
+        super().__init__(
+            datastore,
+            tracking_table_name=tracking_table_name,
+            events_table_name=events_table_name,
+        )
+
     def _insert_events(
         self,
         c: Cursor[DictRow],
@@ -544,18 +666,11 @@ class PostgresProcessRecorder(PostgresApplicationRecorder, ProcessRecorder):
     ) -> None:
         tracking: Tracking | None = kwargs.get("tracking", None)
         if tracking is not None:
-            c.execute(
-                query=self.insert_tracking_statement,
-                params=(
-                    tracking.application_name,
-                    tracking.notification_id,
-                ),
-                prepare=True,
-            )
+            self._insert_tracking(c, tracking=tracking)
         super()._insert_events(c, stored_events, **kwargs)
 
 
-class Factory(InfrastructureFactory):
+class PostgresFactory(InfrastructureFactory[PostgresTrackingRecorder]):
     POSTGRES_DBNAME = "POSTGRES_DBNAME"
     POSTGRES_HOST = "POSTGRES_HOST"
     POSTGRES_PORT = "POSTGRES_PORT"
@@ -577,6 +692,7 @@ class Factory(InfrastructureFactory):
 
     aggregate_recorder_class = PostgresAggregateRecorder
     application_recorder_class = PostgresApplicationRecorder
+    tracking_recorder_class = PostgresTrackingRecorder
     process_recorder_class = PostgresProcessRecorder
 
     def __init__(self, env: Environment):
@@ -768,7 +884,17 @@ class Factory(InfrastructureFactory):
         events_table_name = prefix + "_events"
         if self.datastore.schema:
             events_table_name = f"{self.datastore.schema}.{events_table_name}"
-        recorder = type(self).application_recorder_class(
+
+        application_recorder_topic = self.env.get(self.APPLICATION_RECORDER_TOPIC)
+        if application_recorder_topic:
+            application_recorder_class: Type[PostgresApplicationRecorder] = (
+                resolve_topic(application_recorder_topic)
+            )
+            assert issubclass(application_recorder_class, PostgresApplicationRecorder)
+        else:
+            application_recorder_class = type(self).application_recorder_class
+
+        recorder = application_recorder_class(
             datastore=self.datastore,
             events_table_name=events_table_name,
         )
@@ -776,6 +902,29 @@ class Factory(InfrastructureFactory):
             recorder.create_table()
         return recorder
 
+    def tracking_recorder(
+        self, tracking_recorder_class: Type[PostgresTrackingRecorder] | None = None
+    ) -> PostgresTrackingRecorder:
+        prefix = self.env.name.lower() or "notification"
+        tracking_table_name = prefix + "_tracking"
+        if self.datastore.schema:
+            tracking_table_name = f"{self.datastore.schema}.{tracking_table_name}"
+        if tracking_recorder_class is None:
+            tracking_recorder_topic = self.env.get(self.TRACKING_RECORDER_TOPIC)
+            if tracking_recorder_topic:
+                tracking_recorder_class = resolve_topic(tracking_recorder_topic)
+            else:
+                tracking_recorder_class = type(self).tracking_recorder_class
+        assert tracking_recorder_class is not None
+        assert issubclass(tracking_recorder_class, PostgresTrackingRecorder)
+        recorder = tracking_recorder_class(
+            datastore=self.datastore,
+            tracking_table_name=tracking_table_name,
+        )
+        if self.env_create_table():
+            recorder.create_table()
+        return recorder
+
     def process_recorder(self) -> ProcessRecorder:
         prefix = self.env.name.lower() or "stored"
         events_table_name = prefix + "_events"
@@ -784,7 +933,17 @@ class Factory(InfrastructureFactory):
         if self.datastore.schema:
             events_table_name = f"{self.datastore.schema}.{events_table_name}"
             tracking_table_name = f"{self.datastore.schema}.{tracking_table_name}"
-        recorder = type(self).process_recorder_class(
+
+        process_recorder_topic = self.env.get(self.PROCESS_RECORDER_TOPIC)
+        if process_recorder_topic:
+            process_recorder_class: Type[PostgresTrackingRecorder] = resolve_topic(
+                process_recorder_topic
+            )
+            assert issubclass(process_recorder_class, PostgresProcessRecorder)
+        else:
+            process_recorder_class = type(self).process_recorder_class
+
+        recorder = process_recorder_class(
             datastore=self.datastore,
             events_table_name=events_table_name,
             tracking_table_name=tracking_table_name,
@@ -799,3 +958,6 @@ class Factory(InfrastructureFactory):
 
     def __del__(self) -> None:
         self.close()
+
+
+Factory = PostgresFactory