Add enum for db schemes

Author: tulir

mautrix/client/state_store/asyncpg/store.py

@@ -16,7 +16,7 @@
     RoomID,
     UserID,
 )
-from mautrix.util.async_db import Database
+from mautrix.util.async_db import Database, Scheme
 
 from ..abstract import StateStore
 from .upgrade import upgrade_table
@@ -80,7 +80,7 @@ async def get_members(
         memberships: tuple[Membership, ...] = (Membership.JOIN, Membership.INVITE),
     ) -> list[UserID]:
         membership_values = [membership.value for membership in memberships]
-        if self.db.scheme == "postgres":
+        if self.db.scheme in (Scheme.POSTGRES, Scheme.COCKROACH):
             q = "SELECT user_id FROM mx_user_profile WHERE room_id=$1 AND membership=ANY($2)"
             res = await self.db.fetch(q, room_id, membership_values)
         else:
@@ -98,7 +98,7 @@ async def get_member_profiles(
         memberships: tuple[Membership, ...] = (Membership.JOIN, Membership.INVITE),
     ) -> dict[UserID, Member]:
         membership_values = [membership.value for membership in memberships]
-        if self.db.scheme == "postgres":
+        if self.db.scheme in (Scheme.POSTGRES, Scheme.COCKROACH):
             q = (
                 "SELECT user_id, membership, displayname, avatar_url FROM mx_user_profile "
                 "WHERE room_id=$1 AND membership=ANY($2)"
@@ -123,7 +123,7 @@ async def get_members_filtered(
     ) -> list[UserID]:
         not_like = f"{not_prefix}%{not_suffix}"
         membership_values = [membership.value for membership in memberships]
-        if self.db.scheme == "postgres":
+        if self.db.scheme in (Scheme.POSTGRES, Scheme.COCKROACH):
             q = (
                 "SELECT user_id FROM mx_user_profile "
                 "WHERE room_id=$1 AND membership=ANY($2)"
@@ -155,15 +155,15 @@ async def set_members(
             del_q = "DELETE FROM mx_user_profile WHERE room_id=$1"
             if only_membership is None:
                 await conn.execute(del_q, room_id)
-            elif self.db.scheme == "postgres":
+            elif self.db.scheme in (Scheme.POSTGRES, Scheme.COCKROACH):
                 del_q = f"{del_q} AND (membership=$2 OR user_id = ANY($3))"
                 await conn.execute(del_q, room_id, only_membership.value, list(members.keys()))
             else:
                 member_placeholders = ("?," * len(members)).rstrip(",")
                 del_q = f"{del_q} AND (membership=? OR user_id IN ({member_placeholders}))"
                 await conn.execute(del_q, room_id, only_membership.value, *members.keys())
 
-            if self.db.scheme == "postgres":
+            if self.db.scheme == Scheme.POSTGRES:
                 await conn.copy_records_to_table(
                     "mx_user_profile", records=records, columns=columns
                 )

mautrix/client/state_store/asyncpg/upgrade.py

@@ -5,9 +5,7 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 import logging
 
-from asyncpg import Connection
-
-from mautrix.util.async_db.upgrade import UpgradeTable
+from mautrix.util.async_db import Connection, Scheme, UpgradeTable
 
 upgrade_table = UpgradeTable(
     version_table_name="mx_version",
@@ -17,7 +15,7 @@
 
 
 @upgrade_table.register(description="Latest revision", upgrades_to=2)
-async def upgrade_blank_to_v2(conn: Connection, scheme: str) -> None:
+async def upgrade_blank_to_v2(conn: Connection, scheme: Scheme) -> None:
     await conn.execute(
         """CREATE TABLE mx_room_state (
             room_id              TEXT PRIMARY KEY,
@@ -27,7 +25,7 @@ async def upgrade_blank_to_v2(conn: Connection, scheme: str) -> None:
             power_levels         TEXT
         )"""
     )
-    if scheme != "sqlite":
+    if scheme != Scheme.SQLITE:
         await conn.execute(
             "CREATE TYPE membership AS ENUM ('join', 'leave', 'invite', 'ban', 'knock')"
         )
@@ -44,8 +42,8 @@ async def upgrade_blank_to_v2(conn: Connection, scheme: str) -> None:
 
 
 @upgrade_table.register(description="Stop using size-limited string fields")
-async def upgrade_v2(conn: Connection, scheme: str) -> None:
-    if scheme == "sqlite":
+async def upgrade_v2(conn: Connection, scheme: Scheme) -> None:
+    if scheme == Scheme.SQLITE:
         # SQLite doesn't care about types
         return
     await conn.execute("ALTER TABLE mx_room_state ALTER COLUMN room_id TYPE TEXT")

mautrix/crypto/store/asyncpg/store.py

@@ -5,14 +5,13 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from __future__ import annotations
 
-from typing import cast
 from collections import defaultdict
 from datetime import timedelta
 
 from mautrix.client.state_store import SyncStore
 from mautrix.client.state_store.asyncpg import PgStateStore
 from mautrix.types import DeviceID, EventID, IdentityKey, RoomID, SessionID, SyncToken, UserID
-from mautrix.util.async_db import Database
+from mautrix.util.async_db import Database, Scheme
 from mautrix.util.logging import TraceLogger
 
 from ... import (
@@ -339,7 +338,7 @@ async def remove_outbound_group_session(self, room_id: RoomID) -> None:
         )
 
     async def remove_outbound_group_sessions(self, rooms: list[RoomID]) -> None:
-        if self.db.scheme == "postgres":
+        if self.db.scheme in (Scheme.POSTGRES, Scheme.COCKROACH):
             await self.db.execute(
                 "DELETE FROM crypto_megolm_outbound_session "
                 "WHERE account_id=$1 AND room_id=ANY($2)",
@@ -374,7 +373,7 @@ async def validate_message_index(
         index: int,
         timestamp: int,
     ) -> bool:
-        if self.db.scheme == "postgres":
+        if self.db.scheme in (Scheme.POSTGRES, Scheme.COCKROACH):
             row = await self.db.fetchrow(
                 self._validate_message_index_query,
                 sender_key,
@@ -499,7 +498,7 @@ async def put_devices(self, user_id: UserID, devices: dict[DeviceID, DeviceIdent
                 user_id,
             )
             await conn.execute("DELETE FROM crypto_device WHERE user_id=$1", user_id)
-            if self.db.scheme == "postgres":
+            if self.db.scheme == Scheme.POSTGRES:
                 await conn.copy_records_to_table("crypto_device", records=data, columns=columns)
             else:
                 await conn.executemany(
@@ -510,7 +509,7 @@ async def put_devices(self, user_id: UserID, devices: dict[DeviceID, DeviceIdent
                 )
 
     async def filter_tracked_users(self, users: list[UserID]) -> list[UserID]:
-        if self.db.scheme == "postgres":
+        if self.db.scheme in (Scheme.POSTGRES, Scheme.COCKROACH):
             rows = await self.db.fetch(
                 "SELECT user_id FROM crypto_tracked_user WHERE user_id = ANY($1)", users
             )

mautrix/crypto/store/asyncpg/upgrade.py

@@ -7,7 +7,7 @@
 
 import logging
 
-from mautrix.util.async_db import Connection, UpgradeTable
+from mautrix.util.async_db import Connection, Scheme, UpgradeTable
 
 upgrade_table = UpgradeTable(
     version_table_name="crypto_version",
@@ -96,8 +96,8 @@ async def upgrade_blank_to_v4(conn: Connection) -> None:
 
 
 @upgrade_table.register(description="Add account_id primary key column")
-async def upgrade_v2(conn: Connection, scheme: str) -> None:
-    if scheme == "sqlite":
+async def upgrade_v2(conn: Connection, scheme: Scheme) -> None:
+    if scheme == Scheme.SQLITE:
         await conn.execute("DROP TABLE crypto_account")
         await conn.execute("DROP TABLE crypto_olm_session")
         await conn.execute("DROP TABLE crypto_megolm_inbound_session")
@@ -170,8 +170,8 @@ async def add_account_id_column(table: str, pkey_columns: list[str]) -> None:
 
 
 @upgrade_table.register(description="Stop using size-limited string fields")
-async def upgrade_v3(conn: Connection, scheme: str) -> None:
-    if scheme == "sqlite":
+async def upgrade_v3(conn: Connection, scheme: Scheme) -> None:
+    if scheme == Scheme.SQLITE:
         return
     await conn.execute("ALTER TABLE crypto_account ALTER COLUMN account_id TYPE TEXT")
     await conn.execute("ALTER TABLE crypto_account ALTER COLUMN device_id TYPE TEXT")
@@ -192,11 +192,11 @@ async def upgrade_v3(conn: Connection, scheme: str) -> None:
 
 
 @upgrade_table.register(description="Split last_used into last_encrypted and last_decrypted")
-async def upgrade_v4(conn: Connection, scheme: str) -> None:
+async def upgrade_v4(conn: Connection, scheme: Scheme) -> None:
     await conn.execute("ALTER TABLE crypto_olm_session RENAME COLUMN last_used TO last_decrypted")
     await conn.execute("ALTER TABLE crypto_olm_session ADD COLUMN last_encrypted timestamp")
     await conn.execute("UPDATE crypto_olm_session SET last_encrypted=last_decrypted")
-    if scheme == "postgres":
+    if scheme == Scheme.POSTGRES:
         # This is too hard to do on sqlite, so let's just do it on postgres
         await conn.execute(
             "ALTER TABLE crypto_olm_session ALTER COLUMN last_encrypted SET NOT NULL"

mautrix/util/async_db/__init__.py

@@ -2,6 +2,7 @@
 
 from .connection import LoggingConnection as Connection
 from .database import Database
+from .scheme import Scheme
 from .upgrade import UpgradeTable, register_upgrade
 
 try:

mautrix/util/async_db/aiosqlite.py

@@ -7,16 +7,17 @@
 
 from typing import Any
 from contextlib import asynccontextmanager
-from urllib.parse import urlparse
 import asyncio
 import logging
 import re
 import sqlite3
 
+from yarl import URL
 import aiosqlite
 
 from .connection import LoggingConnection
 from .database import Database
+from .scheme import Scheme
 from .upgrade import UpgradeTable
 
 POSITIONAL_PARAM_PATTERN = re.compile(r"\$(\d+)")
@@ -73,20 +74,20 @@ async def fetchval(
 
 
 class SQLiteDatabase(Database):
-    scheme = "sqlite"
+    scheme = Scheme.SQLITE
     _pool: asyncio.Queue[TxnConnection]
     _stopped: bool
     _conns: int
 
     def __init__(
         self,
-        url: str,
+        url: URL,
         upgrade_table: UpgradeTable,
         db_args: dict[str, Any] | None = None,
         log: logging.Logger | None = None,
     ) -> None:
         super().__init__(url, db_args=db_args, upgrade_table=upgrade_table, log=log)
-        self._path = urlparse(url).path
+        self._path = url.path
         if self._path.startswith("/"):
             self._path = self._path[1:]
         self._pool = asyncio.Queue(self._db_args.pop("min_size", 1))

mautrix/util/async_db/asyncpg.py

@@ -10,25 +10,31 @@
 import asyncio
 import logging
 
+from yarl import URL
 import asyncpg
 
 from .connection import LoggingConnection
 from .database import Database
+from .scheme import Scheme
 from .upgrade import UpgradeTable
 
 
 class PostgresDatabase(Database):
-    scheme = "postgres"
+    scheme = Scheme.POSTGRES
     _pool: asyncpg.pool.Pool | None
     _pool_override: bool
 
     def __init__(
         self,
-        url: str,
+        url: URL,
         upgrade_table: UpgradeTable,
         db_args: dict[str, Any] = None,
         log: logging.Logger | None = None,
     ) -> None:
+        if url.scheme in ("cockroach", "cockroachdb"):
+            self.scheme = Scheme.COCKROACH
+            # Send postgres scheme to asyncpg
+            url = url.with_scheme("postgres")
         super().__init__(url, db_args=db_args, upgrade_table=upgrade_table, log=log)
         self._pool = None
         self._pool_override = False
@@ -41,7 +47,7 @@ async def start(self) -> None:
         if not self._pool_override:
             self._db_args["loop"] = asyncio.get_running_loop()
             self.log.debug(f"Connecting to {self.url}")
-            self._pool = await asyncpg.create_pool(self.url, **self._db_args)
+            self._pool = await asyncpg.create_pool(str(self.url), **self._db_args)
         await super().start()
 
     @property
@@ -62,3 +68,5 @@ async def acquire(self) -> LoggingConnection:
 
 Database.schemes["postgres"] = PostgresDatabase
 Database.schemes["postgresql"] = PostgresDatabase
+Database.schemes["cockroach"] = PostgresDatabase
+Database.schemes["cockroachdb"] = PostgresDatabase

mautrix/util/async_db/connection.py

@@ -14,6 +14,8 @@
 from mautrix import __optional_imports__
 from mautrix.util.logging import SILLY, TraceLogger
 
+from .scheme import Scheme
+
 if __optional_imports__:
     from sqlite3 import Row
 
@@ -41,12 +43,15 @@ async def wrapper(self: LoggingConnection, arg: str, *args: Any, **kwargs: str)
 
 
 class LoggingConnection:
-    scheme: str
+    scheme: Scheme
     wrapped: aiosqlite.TxnConnection | asyncpg.Connection
     log: TraceLogger
 
     def __init__(
-        self, scheme: str, wrapped: aiosqlite.TxnConnection | asyncpg.Connection, log: TraceLogger
+        self,
+        scheme: Scheme,
+        wrapped: aiosqlite.TxnConnection | asyncpg.Connection,
+        log: TraceLogger,
     ) -> None:
         self.scheme = scheme
         self.wrapped = wrapped
@@ -97,8 +102,8 @@ async def copy_records_to_table(
         schema_name: str | None = None,
         timeout: float | None = None,
     ) -> None:
-        if self.scheme != "postgres":
-            raise RuntimeError("copy_records_to_table is not supported on SQLite")
+        if self.scheme != Scheme.POSTGRES:
+            raise RuntimeError("copy_records_to_table is only supported on Postgres")
         return await self.wrapped.copy_records_to_table(
             table_name, records=records, columns=columns, schema_name=schema_name, timeout=timeout
         )