# Copyright (c) 2022 Tulir Asokan # # This Source Code Form is subject to the terms of the Mozilla Public # License, v. 2.0. If a copy of the MPL was not distributed with this # file, You can obtain one at http://mozilla.org/MPL/2.0/. from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Type, TypeVar, cast from contextlib import contextmanager from sqlalchemy import Constraint, Table from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.engine.base import Connection, Engine from sqlalchemy.ext.declarative import as_declarative, declarative_base from sqlalchemy.sql.base import ImmutableColumnCollection from sqlalchemy.sql.expression import ClauseElement, Select, and_ if TYPE_CHECKING: from sqlalchemy.engine.result import ResultProxy, RowProxy T = TypeVar("T", bound="BaseClass") class BaseClass: """ Base class for SQLAlchemy models. Provides SQLAlchemy declarative base features and some additional utilities. .. deprecated:: 0.15.0 The :mod:`mautrix.util.async_db` utility is now recommended over SQLAlchemy. """ __tablename__: str db: Engine t: Table __table__: Table c: ImmutableColumnCollection column_names: List[str] @classmethod def bind(cls, db_engine: Engine) -> None: cls.db = db_engine cls.t = cls.__table__ cls.c = cls.t.columns cls.column_names = cls.c.keys() @classmethod def copy( cls, bind: Optional[Engine] = None, rebase: Optional[declarative_base] = None ) -> Type[T]: copy = cast(Type[T], type(cls.__name__, (cls, rebase) if rebase else (cls,), {})) if bind is not None: copy.bind(db_engine=bind) return copy @classmethod def _one_or_none(cls: Type[T], rows: "ResultProxy") -> Optional[T]: """ Try scanning one row from a ResultProxy and return ``None`` if it fails. Args: rows: The SQLAlchemy result to scan. Returns: The scanned object, or ``None`` if there were no rows. """ try: return cls.scan(next(rows)) except StopIteration: return None @classmethod def _all(cls: Type[T], rows: "ResultProxy") -> Iterator[T]: """ Scan all rows from a ResultProxy. Args: rows: The SQLAlchemy result to scan. Yields: Each row scanned with :meth:`scan` """ for row in rows: yield cls.scan(row) @classmethod def scan(cls: Type[T], row: "RowProxy") -> T: """ Read the data from a row into an object. Args: row: The RowProxy object. Returns: An object containing the information in the row. """ return cls(**dict(zip(cls.column_names, row))) @classmethod def _make_simple_select(cls: Type[T], *args: ClauseElement) -> Select: """ Create a simple ``SELECT * FROM table WHERE `` statement. Args: *args: The WHERE clauses. If there are many elements, they're joined with AND. Returns: The SQLAlchemy SELECT statement object. """ if len(args) > 1: return cls.t.select().where(and_(*args)) elif len(args) == 1: return cls.t.select().where(args[0]) else: return cls.t.select() @classmethod def _select_all(cls: Type[T], *args: ClauseElement) -> Iterator[T]: """ Select all rows with given conditions. This is intended to be used by table-specific select methods. Args: *args: The WHERE clauses. If there are many elements, they're joined with AND. Yields: The objects representing the rows read with :meth:`scan` """ yield from cls._all(cls.db.execute(cls._make_simple_select(*args))) @classmethod def _select_one_or_none(cls: Type[T], *args: ClauseElement) -> T: """ Select one row with given conditions. If no row is found, return ``None``. This is intended to be used by table-specific select methods. Args: *args: The WHERE clauses. If there are many elements, they're joined with AND. Returns: The object representing the matched row read with :meth:`scan`, or ``None`` if no rows matched. """ return cls._one_or_none(cls.db.execute(cls._make_simple_select(*args))) def _constraint_to_clause(self, constraint: Constraint) -> ClauseElement: return and_( *[column == self.__dict__[name] for name, column in constraint.columns.items()] ) @property def _edit_identity(self: T) -> ClauseElement: """The SQLAlchemy WHERE clause used for editing and deleting individual rows. Usually AND of primary keys.""" return self._constraint_to_clause(self.t.primary_key) def edit(self: T, *, _update_values: bool = True, **values) -> None: """ Edit this row. Args: _update_values: Whether or not the values in memory should be updated as well as the values in the database. **values: The values to change. """ with self.db.begin() as conn: conn.execute(self.t.update().where(self._edit_identity).values(**values)) if _update_values: for key, value in values.items(): setattr(self, key, value) @contextmanager def edit_mode(self: T) -> None: """ Edit this row in a fancy context manager way. This stores the current edit identity, then yields to the context manager and finally puts the new values into the row using the old edit identity in the WHERE clause. >>> class TableClass(Base): ... ... >>> db_instance = TableClass(id="something") >>> with db_instance.edit_mode(): ... db_instance.id = "new_id" """ old_identity = self._edit_identity yield old_identity with self.db.begin() as conn: conn.execute(self.t.update().where(old_identity).values(**self._insert_values)) def delete(self: T) -> None: """Delete this row.""" with self.db.begin() as conn: conn.execute(self.t.delete().where(self._edit_identity)) @property def _insert_values(self: T) -> Dict[str, Any]: """Values for inserts. Generally you want all the values in the table.""" return { column_name: self.__dict__[column_name] for column_name in self.column_names if column_name in self.__dict__ } def insert(self) -> None: with self.db.begin() as conn: conn.execute(self.t.insert().values(**self._insert_values)) @property def _upsert_values(self: T) -> Dict[str, Any]: """The values to set when an upsert-insert conflicts and moves to the update part.""" return self._insert_values def _upsert_postgres(self: T, conn: Connection) -> None: conn.execute( pg_insert(self.t) .values(**self._insert_values) .on_conflict_do_update(constraint=self.t.primary_key, set_=self._upsert_values) ) def _upsert_sqlite(self: T, conn: Connection) -> None: conn.execute(self.t.insert().values(**self._insert_values).prefix_with("OR REPLACE")) def _upsert_generic(self: T, conn: Connection): conn.execute(self.t.delete().where(self._edit_identity)) conn.execute(self.t.insert().values(**self._insert_values)) def upsert(self: T) -> None: with self.db.begin() as conn: if self.db.dialect.name == "postgresql": self._upsert_postgres(conn) elif self.db.dialect.name == "sqlite": self._upsert_sqlite(conn) else: self._upsert_generic(conn) def __iter__(self): for key in self.column_names: yield self.__dict__[key] @as_declarative() class Base(BaseClass): """ .. deprecated:: 0.15.0 The :mod:`mautrix.util.async_db` utility is now recommended over SQLAlchemy. """ pass