mega chore: update sqlalchemy (#2611)

* update sqlalchemy to 1.4 * async postgres --------- Co-authored-by: Pavol Rusnak <pavol@rusnak.io>
2024-09-24 10:56:03 +02:00
parent c637e8d31e
commit 21d87adc52
17 changed files with 1020 additions and 951 deletions
@@ -7,14 +7,13 @@ import re
 import time
 from contextlib import asynccontextmanager
 from enum import Enum
-from sqlite3 import Row
 from typing import Any, Generic, Literal, Optional, TypeVar

 from loguru import logger
 from pydantic import BaseModel, ValidationError, root_validator
-from sqlalchemy import create_engine
-from sqlalchemy_aio.base import AsyncConnection
-from sqlalchemy_aio.strategy import ASYNCIO_STRATEGY
+from sqlalchemy import event
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
+from sqlalchemy.sql import text

 from lnbits.settings import settings

@@ -24,31 +23,15 @@ SQLITE = "SQLITE"

 if settings.lnbits_database_url:
    database_uri = settings.lnbits_database_url
-
    if database_uri.startswith("cockroachdb://"):
        DB_TYPE = COCKROACH
    else:
+        if not database_uri.startswith("postgres://"):
+            raise ValueError(
+                "Please use the 'postgres://...' " "format for the database URL."
+            )
        DB_TYPE = POSTGRES

-    from psycopg2.extensions import DECIMAL, new_type, register_type
-
-    def _parse_timestamp(value, _):
-        if value is None:
-            return None
-        f = "%Y-%m-%d %H:%M:%S.%f"
-        if "." not in value:
-            f = "%Y-%m-%d %H:%M:%S"
-        return time.mktime(datetime.datetime.strptime(value, f).timetuple())
-
-    register_type(
-        new_type(
-            DECIMAL.values,
-            "DEC2FLOAT",
-            lambda value, curs: float(value) if value is not None else None,
-        )
-    )
-
-    register_type(new_type((1184, 1114), "TIMESTAMP2INT", _parse_timestamp))
 else:
    if not os.path.isdir(settings.lnbits_data_folder):
        os.mkdir(settings.lnbits_data_folder)
@@ -56,21 +39,21 @@ else:
    DB_TYPE = SQLITE


-def compat_timestamp_placeholder():
+def compat_timestamp_placeholder(key: str):
    if DB_TYPE == POSTGRES:
-        return "to_timestamp(?)"
+        return f"to_timestamp(:{key})"
    elif DB_TYPE == COCKROACH:
-        return "cast(? AS timestamp)"
+        return f"cast(:{key} AS timestamp)"
    else:
-        return "?"
+        return f":{key}"


 def get_placeholder(model: Any, field: str) -> str:
    type_ = model.__fields__[field].type_
    if type_ == datetime.datetime:
-        return compat_timestamp_placeholder()
+        return compat_timestamp_placeholder(field)
    else:
-        return "?"
+        return f":{field}"


 class Compat:
@@ -127,15 +110,13 @@ class Compat:
            return "BIGINT"
        return "INT"

-    @property
-    def timestamp_placeholder(self) -> str:
-        return compat_timestamp_placeholder()
+    def timestamp_placeholder(self, key: str) -> str:
+        return compat_timestamp_placeholder(key)


 class Connection(Compat):
-    def __init__(self, conn: AsyncConnection, txn, typ, name, schema):
+    def __init__(self, conn: AsyncConnection, typ, name, schema):
        self.conn = conn
-        self.txn = txn
        self.type = typ
        self.name = name
        self.schema = schema
@@ -146,45 +127,42 @@ class Connection(Compat):
            query = query.replace("?", "%s")
        return query

-    def rewrite_values(self, values):
+    def rewrite_values(self, values: dict) -> dict:
        # strip html
        clean_regex = re.compile("<.*?>|&([a-z0-9]+|#[0-9]{1,6}|#x[0-9a-f]{1,6});")
-
-        # tuple to list and back to tuple
-        raw_values = [values] if isinstance(values, str) else list(values)
-        values = []
-        for raw_value in raw_values:
+        clean_values: dict = {}
+        for key, raw_value in values.items():
            if isinstance(raw_value, str):
-                values.append(re.sub(clean_regex, "", raw_value))
+                clean_values[key] = re.sub(clean_regex, "", raw_value)
            elif isinstance(raw_value, datetime.datetime):
                ts = raw_value.timestamp()
                if self.type == SQLITE:
-                    values.append(int(ts))
+                    clean_values[key] = int(ts)
                else:
-                    values.append(ts)
+                    clean_values[key] = ts
            else:
-                values.append(raw_value)
-        return tuple(values)
+                clean_values[key] = raw_value
+        return clean_values

-    async def fetchall(self, query: str, values: tuple = ()) -> list:
-        result = await self.conn.execute(
-            self.rewrite_query(query), self.rewrite_values(values)
-        )
-        return await result.fetchall()
+    async def fetchall(self, query: str, values: Optional[dict] = None) -> list[dict]:
+        params = self.rewrite_values(values) if values else {}
+        result = await self.conn.execute(text(self.rewrite_query(query)), params)
+        row = result.mappings().all()
+        result.close()
+        return row

-    async def fetchone(self, query: str, values: tuple = ()):
-        result = await self.conn.execute(
-            self.rewrite_query(query), self.rewrite_values(values)
-        )
-        row = await result.fetchone()
-        await result.close()
+    async def fetchone(self, query: str, values: Optional[dict] = None) -> dict:
+        params = self.rewrite_values(values) if values else {}
+        result = await self.conn.execute(text(self.rewrite_query(query)), params)
+        row = result.mappings().first()
+        result.close()
        return row

    async def fetch_page(
        self,
        query: str,
        where: Optional[list[str]] = None,
-        values: Optional[list[str]] = None,
+        values: Optional[dict] = None,
        filters: Optional[Filters] = None,
        model: Optional[type[TRowModel]] = None,
        group_by: Optional[list[str]] = None,
@@ -211,14 +189,14 @@ class Connection(Compat):
            {filters.order_by()}
            {filters.pagination()}
            """,
-            parsed_values,
+            self.rewrite_values(parsed_values),
        )
        if rows:
            # no need for extra query if no pagination is specified
            if filters.offset or filters.limit:
-                count = await self.fetchone(
+                result = await self.fetchone(
                    f"""
-                    SELECT COUNT(*) FROM (
+                    SELECT COUNT(*) as count FROM (
                        {query}
                        {clause}
                        {group_by_string}
@@ -226,21 +204,22 @@ class Connection(Compat):
                    """,
                    parsed_values,
                )
-                count = int(count[0])
+                count = int(result.get("count", 0))
            else:
                count = len(rows)
        else:
            count = 0

        return Page(
-            data=[model.from_row(row) for row in rows] if model else rows,
+            data=[model.from_row(row) for row in rows] if model else [],
            total=count,
        )

-    async def execute(self, query: str, values: tuple = ()):
-        return await self.conn.execute(
-            self.rewrite_query(query), self.rewrite_values(values)
-        )
+    async def execute(self, query: str, values: Optional[dict] = None):
+        params = self.rewrite_values(values) if values else {}
+        result = await self.conn.execute(text(self.rewrite_query(query)), params)
+        await self.conn.commit()
+        return result


 class Database(Compat):
@@ -253,18 +232,44 @@ class Database(Compat):
            self.path = os.path.join(
                settings.lnbits_data_folder, f"{self.name}.sqlite3"
            )
-            database_uri = f"sqlite:///{self.path}"
+            database_uri = f"sqlite+aiosqlite:///{self.path}"
        else:
-            database_uri = settings.lnbits_database_url
+            database_uri = settings.lnbits_database_url.replace(
+                "postgres://", "postgresql+asyncpg://"
+            )

        if self.name.startswith("ext_"):
            self.schema = self.name[4:]
        else:
            self.schema = None

-        self.engine = create_engine(
-            database_uri, strategy=ASYNCIO_STRATEGY, echo=settings.debug_database
+        self.engine: AsyncEngine = create_async_engine(
+            database_uri, echo=settings.debug_database
        )
+
+        if self.type in {POSTGRES, COCKROACH}:
+
+            @event.listens_for(self.engine.sync_engine, "connect")
+            def register_custom_types(dbapi_connection, *_):
+                def _parse_timestamp(value):
+                    if value is None:
+                        return None
+                    f = "%Y-%m-%d %H:%M:%S.%f"
+                    if "." not in value:
+                        f = "%Y-%m-%d %H:%M:%S"
+                    return int(
+                        time.mktime(datetime.datetime.strptime(value, f).timetuple())
+                    )
+
+                dbapi_connection.run_async(
+                    lambda connection: connection.set_type_codec(
+                        "TIMESTAMP",
+                        encoder=datetime.datetime,
+                        decoder=_parse_timestamp,
+                        schema="pg_catalog",
+                    )
+                )
+
        self.lock = asyncio.Lock()

        logger.trace(f"database {self.type} added for {self.name}")
@@ -273,41 +278,37 @@ class Database(Compat):
    async def connect(self):
        await self.lock.acquire()
        try:
-            async with self.engine.connect() as conn:  # type: ignore
-                async with conn.begin() as txn:
-                    wconn = Connection(conn, txn, self.type, self.name, self.schema)
+            async with self.engine.connect() as conn:
+                if not conn:
+                    raise Exception("Could not connect to the database")

-                    if self.schema:
-                        if self.type in {POSTGRES, COCKROACH}:
-                            await wconn.execute(
-                                f"CREATE SCHEMA IF NOT EXISTS {self.schema}"
-                            )
-                        elif self.type == SQLITE:
-                            await wconn.execute(
-                                f"ATTACH '{self.path}' AS {self.schema}"
-                            )
+                wconn = Connection(conn, self.type, self.name, self.schema)

-                    yield wconn
+                if self.schema:
+                    if self.type in {POSTGRES, COCKROACH}:
+                        await wconn.execute(
+                            f"CREATE SCHEMA IF NOT EXISTS {self.schema}"
+                        )
+                    elif self.type == SQLITE:
+                        await wconn.execute(f"ATTACH '{self.path}' AS {self.schema}")
+
+                yield wconn
        finally:
            self.lock.release()

-    async def fetchall(self, query: str, values: tuple = ()) -> list:
+    async def fetchall(self, query: str, values: Optional[dict] = None) -> list[dict]:
        async with self.connect() as conn:
-            result = await conn.execute(query, values)
-            return await result.fetchall()
+            return await conn.fetchall(query, values)

-    async def fetchone(self, query: str, values: tuple = ()):
+    async def fetchone(self, query: str, values: Optional[dict] = None) -> dict:
        async with self.connect() as conn:
-            result = await conn.execute(query, values)
-            row = await result.fetchone()
-            await result.close()
-            return row
+            return await conn.fetchone(query, values)

    async def fetch_page(
        self,
        query: str,
        where: Optional[list[str]] = None,
-        values: Optional[list[str]] = None,
+        values: Optional[dict] = None,
        filters: Optional[Filters] = None,
        model: Optional[type[TRowModel]] = None,
        group_by: Optional[list[str]] = None,
@@ -315,7 +316,7 @@ class Database(Compat):
        async with self.connect() as conn:
            return await conn.fetch_page(query, where, values, filters, model, group_by)

-    async def execute(self, query: str, values: tuple = ()):
+    async def execute(self, query: str, values: Optional[dict] = None):
        async with self.connect() as conn:
            return await conn.execute(query, values)

@@ -373,8 +374,8 @@ class Operator(Enum):

 class FromRowModel(BaseModel):
    @classmethod
-    def from_row(cls, row: Row):
-        return cls(**dict(row))
+    def from_row(cls, row: dict):
+        return cls(**row)


 class FilterModel(BaseModel):
@@ -396,12 +397,13 @@ class Page(BaseModel, Generic[T]):
 class Filter(BaseModel, Generic[TFilterModel]):
    field: str
    op: Operator = Operator.EQ
-    values: list[Any]
-
    model: Optional[type[TFilterModel]]
+    values: Optional[dict] = None

    @classmethod
-    def parse_query(cls, key: str, raw_values: list[Any], model: type[TFilterModel]):
+    def parse_query(
+        cls, key: str, raw_values: list[Any], model: type[TFilterModel], i: int = 0
+    ):
        # Key format:
        # key[operator]
        # e.g. name[eq]
@@ -417,12 +419,12 @@ class Filter(BaseModel, Generic[TFilterModel]):

        if field in model.__fields__:
            compare_field = model.__fields__[field]
-            values = []
+            values: dict = {}
            for raw_value in raw_values:
                validated, errors = compare_field.validate(raw_value, {}, loc="none")
                if errors:
                    raise ValidationError(errors=[errors], model=model)
-                values.append(validated)
+                values[f"{field}__{i}"] = validated
        else:
            raise ValueError("Unknown filter field")

@@ -430,13 +432,17 @@ class Filter(BaseModel, Generic[TFilterModel]):

    @property
    def statement(self):
-        assert self.model, "Model is required for statement generation"
-        placeholder = get_placeholder(self.model, self.field)
-        if self.op in (Operator.INCLUDE, Operator.EXCLUDE):
-            placeholders = ", ".join([placeholder] * len(self.values))
-            stmt = [f"{self.field} {self.op.as_sql} ({placeholders})"]
-        else:
-            stmt = [f"{self.field} {self.op.as_sql} {placeholder}"] * len(self.values)
+        stmt = []
+        for key in self.values.keys() if self.values else []:
+            clean_key = key.split("__")[0]
+            if (
+                self.model
+                and self.model.__fields__[clean_key].type_ == datetime.datetime
+            ):
+                placeholder = compat_timestamp_placeholder(key)
+            else:
+                placeholder = f":{key}"
+            stmt.append(f"{clean_key} {self.op.as_sql} {placeholder}")
        return " OR ".join(stmt)


@@ -487,14 +493,11 @@ class Filters(BaseModel, Generic[TFilterModel]):
            for page_filter in self.filters:
                where_stmts.append(page_filter.statement)
        if self.search and self.model:
+            fields = self.model.__search_fields__
            if DB_TYPE == POSTGRES:
-                where_stmts.append(
-                    f"lower(concat({', '.join(self.model.__search_fields__)})) LIKE ?"
-                )
+                where_stmts.append(f"lower(concat({', '.join(fields)})) LIKE :search")
            elif DB_TYPE == SQLITE:
-                where_stmts.append(
-                    f"lower({'||'.join(self.model.__search_fields__)}) LIKE ?"
-                )
+                where_stmts.append(f"lower({'||'.join(fields)}) LIKE :search")
        if where_stmts:
            return "WHERE " + " AND ".join(where_stmts)
        return ""
@@ -504,12 +507,14 @@ class Filters(BaseModel, Generic[TFilterModel]):
            return f"ORDER BY {self.sortby} {self.direction or 'asc'}"
        return ""

-    def values(self, values: Optional[list[str]] = None) -> tuple:
+    def values(self, values: Optional[dict] = None) -> dict:
        if not values:
-            values = []
+            values = {}
        if self.filters:
            for page_filter in self.filters:
-                values.extend(page_filter.values)
+                if page_filter.values:
+                    for key, value in page_filter.values.items():
+                        values[key] = value
        if self.search and self.model:
-            values.append(f"%{self.search}%")
-        return tuple(values)
+            values["search"] = f"%{self.search}%"
+        return values