From f5ef7e13d7e60bd9b901f90da250c7c6d1ba1f53 Mon Sep 17 00:00:00 2001 From: Patrick Cloke Date: Fri, 29 Jul 2022 13:57:43 -0400 Subject: [PATCH] Initial abstraction. --- synapse/config/database.py | 2 +- synapse/storage/engines/__init__.py | 33 +++++++++++-------- synapse/storage/engines/_base.py | 12 +++++++ synapse/storage/engines/postgres.py | 10 ++++++ synapse/storage/engines/psycopg.py | 12 ++++++- synapse/storage/engines/sqlite.py | 10 ++++++ synapse/storage/prepare_database.py | 6 ++-- .../main/delta/56/unique_user_filter_index.py | 3 +- .../68/05partial_state_rooms_triggers.py | 9 +++-- .../schema/main/delta/69/01as_txn_seq.py | 2 +- tests/utils.py | 5 ++- 11 files changed, 80 insertions(+), 24 deletions(-) diff --git a/synapse/config/database.py b/synapse/config/database.py index 928fec8dfe..5092bd1fb2 100644 --- a/synapse/config/database.py +++ b/synapse/config/database.py @@ -50,7 +50,7 @@ class DatabaseConnectionConfig: def __init__(self, name: str, db_config: dict): db_engine = db_config.get("name", "sqlite3") - if db_engine not in ("sqlite3", "psycopg2"): + if db_engine not in ("sqlite3", "psycopg2", "psycopg"): raise ConfigError("Unsupported database type %r" % (db_engine,)) if db_engine == "sqlite3": diff --git a/synapse/storage/engines/__init__.py b/synapse/storage/engines/__init__.py index a182e8a098..a9fdcdcef7 100644 --- a/synapse/storage/engines/__init__.py +++ b/synapse/storage/engines/__init__.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Mapping, NoReturn +from typing import Any, Mapping, NoReturn, cast from ._base import BaseDatabaseEngine, IncorrectDatabaseSetup @@ -20,26 +20,30 @@ from ._base import BaseDatabaseEngine, IncorrectDatabaseSetup # and sqlite. But the database driver modules are both optional: they may not be # installed. To account for this, create dummy classes on import failure so we can # still run `isinstance()` checks. +def dummy_engine(name: str, module: str) -> BaseDatabaseEngine: + class Engine(BaseDatabaseEngine): # type: ignore[no-redef] + def __new__(cls, *args: object, **kwargs: object) -> NoReturn: # type: ignore[misc] + raise RuntimeError( + f"Cannot create {name}Engine -- {module} module is not installed" + ) + + return cast(BaseDatabaseEngine, Engine) + + try: from .postgres import PostgresEngine except ImportError: + PostgresEngine = dummy_engine("PostgresEngine", "psycopg2") - class PostgresEngine(BaseDatabaseEngine): # type: ignore[no-redef] - def __new__(cls, *args: object, **kwargs: object) -> NoReturn: # type: ignore[misc] - raise RuntimeError( - f"Cannot create {cls.__name__} -- psycopg2 module is not installed" - ) - +try: + from .psycopg import PsycopgEngine +except ImportError: + PsycopgEngine = dummy_engine("PsycopgEngine", "psycopg") try: from .sqlite import Sqlite3Engine except ImportError: - - class Sqlite3Engine(BaseDatabaseEngine): # type: ignore[no-redef] - def __new__(cls, *args: object, **kwargs: object) -> NoReturn: # type: ignore[misc] - raise RuntimeError( - f"Cannot create {cls.__name__} -- sqlite3 module is not installed" - ) + Sqlite3Engine = dummy_engine("Sqlite3Engine", "sqlite3") def create_engine(database_config: Mapping[str, Any]) -> BaseDatabaseEngine: @@ -51,6 +55,9 @@ def create_engine(database_config: Mapping[str, Any]) -> BaseDatabaseEngine: if name == "psycopg2": return PostgresEngine(database_config) + if name == "psycopg": + return PsycopgEngine(database_config) + raise RuntimeError("Unsupported database engine '%s'" % (name,)) diff --git a/synapse/storage/engines/_base.py b/synapse/storage/engines/_base.py index 70e594a68f..cfc1f1c66b 100644 --- a/synapse/storage/engines/_base.py +++ b/synapse/storage/engines/_base.py @@ -58,6 +58,18 @@ class BaseDatabaseEngine(Generic[ConnectionType, CursorType], metaclass=abc.ABCM """Do we support the `RETURNING` clause in insert/update/delete?""" ... + @property + @abc.abstractmethod + def supports_select_distinct_on(self) -> bool: + """Do we support the `DISTINCT ON` clause in SELECT?""" + ... + + @property + @abc.abstractmethod + def supports_sequences(self) -> bool: + """Do we support the `CREATE SEQUENCE` clause?""" + ... + @abc.abstractmethod def check_database( self, db_conn: ConnectionType, allow_outdated_version: bool = False diff --git a/synapse/storage/engines/postgres.py b/synapse/storage/engines/postgres.py index 719a517336..177d3e796c 100644 --- a/synapse/storage/engines/postgres.py +++ b/synapse/storage/engines/postgres.py @@ -170,6 +170,16 @@ class PostgresEngine( """Do we support the `RETURNING` clause in insert/update/delete?""" return True + @property + def supports_select_distinct_on(self) -> bool: + """Do we support the `DISTINCT ON` clause in SELECT?""" + return True + + @property + def supports_sequences(self) -> bool: + """Do we support the `CREATE SEQUENCE` clause?""" + return True + def is_deadlock(self, error: Exception) -> bool: if isinstance(error, psycopg2.DatabaseError): # https://www.postgresql.org/docs/current/static/errcodes-appendix.html diff --git a/synapse/storage/engines/psycopg.py b/synapse/storage/engines/psycopg.py index 517f9d5f98..bf3cf94777 100644 --- a/synapse/storage/engines/psycopg.py +++ b/synapse/storage/engines/psycopg.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -class PostgresEngine(BaseDatabaseEngine[psycopg2.extensions.connection]): +class PsycopgEngine(BaseDatabaseEngine[psycopg2.extensions.connection]): def __init__(self, database_config: Mapping[str, Any]): super().__init__(psycopg2, database_config) psycopg2.extensions.register_type(psycopg2.extensions.UNICODE) @@ -175,6 +175,16 @@ class PostgresEngine(BaseDatabaseEngine[psycopg2.extensions.connection]): """Do we support the `RETURNING` clause in insert/update/delete?""" return True + @property + def supports_select_distinct_on(self) -> bool: + """Do we support the `DISTINCT ON` clause in SELECT?""" + return True + + @property + def supports_sequences(self) -> bool: + """Do we support the `CREATE SEQUENCE` clause?""" + return True + def is_deadlock(self, error: Exception) -> bool: if isinstance(error, psycopg2.DatabaseError): # https://www.postgresql.org/docs/current/static/errcodes-appendix.html diff --git a/synapse/storage/engines/sqlite.py b/synapse/storage/engines/sqlite.py index 14260442b6..ea8b0b4d14 100644 --- a/synapse/storage/engines/sqlite.py +++ b/synapse/storage/engines/sqlite.py @@ -58,6 +58,16 @@ class Sqlite3Engine(BaseDatabaseEngine[sqlite3.Connection, sqlite3.Cursor]): """Do we support the `RETURNING` clause in insert/update/delete?""" return sqlite3.sqlite_version_info >= (3, 35, 0) + @property + def supports_select_distinct_on(self) -> bool: + """Do we support the `DISTINCT ON` clause in SELECT?""" + return False + + @property + def supports_sequences(self) -> bool: + """Do we support the `CREATE SEQUENCE` clause?""" + return False + def check_database( self, db_conn: sqlite3.Connection, allow_outdated_version: bool = False ) -> None: diff --git a/synapse/storage/prepare_database.py b/synapse/storage/prepare_database.py index 3acdb39da7..bc5f8e4159 100644 --- a/synapse/storage/prepare_database.py +++ b/synapse/storage/prepare_database.py @@ -23,7 +23,7 @@ from typing_extensions import Counter as CounterType from synapse.config.homeserver import HomeServerConfig from synapse.storage.database import LoggingDatabaseConnection -from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine +from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine, PsycopgEngine from synapse.storage.schema import SCHEMA_COMPAT_VERSION, SCHEMA_VERSION from synapse.storage.types import Cursor @@ -255,7 +255,7 @@ def _setup_new_database( for file_name in os.listdir(directory) ) - if isinstance(database_engine, PostgresEngine): + if isinstance(database_engine, (PostgresEngine, PsycopgEngine)): specific = "postgres" else: specific = "sqlite" @@ -399,7 +399,7 @@ def _upgrade_existing_database( logger.debug("applied_delta_files: %s", current_schema_state.applied_deltas) - if isinstance(database_engine, PostgresEngine): + if isinstance(database_engine, (PostgresEngine, PsycopgEngine)): specific_engine_extension = ".postgres" else: specific_engine_extension = ".sqlite" diff --git a/synapse/storage/schema/main/delta/56/unique_user_filter_index.py b/synapse/storage/schema/main/delta/56/unique_user_filter_index.py index bb7296852a..9983c77548 100644 --- a/synapse/storage/schema/main/delta/56/unique_user_filter_index.py +++ b/synapse/storage/schema/main/delta/56/unique_user_filter_index.py @@ -1,7 +1,6 @@ import logging from io import StringIO -from synapse.storage.engines import PostgresEngine from synapse.storage.prepare_database import execute_statements_from_stream logger = logging.getLogger(__name__) @@ -21,7 +20,7 @@ def run_upgrade(cur, database_engine, *args, **kwargs): def run_create(cur, database_engine, *args, **kwargs): - if isinstance(database_engine, PostgresEngine): + if database_engine.supports_select_distinct_on: select_clause = """ SELECT DISTINCT ON (user_id, filter_id) user_id, filter_id, filter_json FROM user_filters diff --git a/synapse/storage/schema/main/delta/68/05partial_state_rooms_triggers.py b/synapse/storage/schema/main/delta/68/05partial_state_rooms_triggers.py index a2ec4fc26e..ae6e7137b8 100644 --- a/synapse/storage/schema/main/delta/68/05partial_state_rooms_triggers.py +++ b/synapse/storage/schema/main/delta/68/05partial_state_rooms_triggers.py @@ -18,7 +18,12 @@ This migration adds triggers to the partial_state_events tables to enforce uniqu Triggers cannot be expressed in .sql files, so we have to use a separate file. """ -from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine, Sqlite3Engine +from synapse.storage.engines import ( + BaseDatabaseEngine, + PostgresEngine, + PsycopgEngine, + Sqlite3Engine, +) from synapse.storage.types import Cursor @@ -43,7 +48,7 @@ def run_create(cur: Cursor, database_engine: BaseDatabaseEngine, *args, **kwargs END; """ ) - elif isinstance(database_engine, PostgresEngine): + elif isinstance(database_engine, (PostgresEngine, PsycopgEngine)): cur.execute( """ CREATE OR REPLACE FUNCTION check_partial_state_events() RETURNS trigger AS $BODY$ diff --git a/synapse/storage/schema/main/delta/69/01as_txn_seq.py b/synapse/storage/schema/main/delta/69/01as_txn_seq.py index 24bd4b391e..4856569ceb 100644 --- a/synapse/storage/schema/main/delta/69/01as_txn_seq.py +++ b/synapse/storage/schema/main/delta/69/01as_txn_seq.py @@ -21,7 +21,7 @@ from synapse.storage.engines import PostgresEngine def run_create(cur, database_engine, *args, **kwargs): - if isinstance(database_engine, PostgresEngine): + if database_engine.supports_sequences: # If we already have some AS TXNs we want to start from the current # maximum value. There are two potential places this is stored - the # actual TXNs themselves *and* the AS state table. At time of migration diff --git a/tests/utils.py b/tests/utils.py index 045a8b5fa7..77e4b18011 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -59,7 +59,10 @@ def setupdb() -> None: # If we're using PostgreSQL, set up the db once if USE_POSTGRES_FOR_TESTS: # create a PostgresEngine - db_engine = create_engine({"name": "psycopg2", "args": {}}) + if USE_POSTGRES_FOR_TESTS == "psycopg": + db_engine = create_engine({"name": "psycopg", "args": {}}) + else: + db_engine = create_engine({"name": "psycopg2", "args": {}}) # connect to postgres to create the base database. db_conn = db_engine.module.connect( user=POSTGRES_USER,