ns-orm 0.0.0__py3-none-any.whl

ns_orm/query.py

@@ -0,0 +1,659 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Any, Generic, Optional, TypeVar
+
+from ns_orm.database import AsyncDatabase, Database
+from ns_orm.dialects import Dialect
+from ns_orm.exceptions import DoesNotExist, MultipleObjectsReturned, QueryError
+from ns_orm.expressions import Condition, Q, normalize_in_values
+
+TModel = TypeVar("TModel")
+
+
+def _table_alias() -> str:
+    return "t"
+
+
+def _qualified_table(model: type[Any], dialect: Dialect) -> str:
+    table = dialect.quote_ident(model.table_name())
+    schema = getattr(model._meta, "schema", None)
+    if schema:
+        return f"{dialect.quote_ident(schema)}.{table}"
+    return table
+
+
+def _where_from_q(
+    model: type[Any], dialect: Dialect, q: Q
+) -> tuple[str, dict[str, Any]]:
+    def _lookup_compiler(
+        *,
+        lookup: str,
+        value: Any,
+        table_alias: str,
+        param_prefix: str,
+        start_index: int,
+    ) -> tuple[Condition, int]:
+        parts = lookup.split("__")
+        field = parts[0]
+        op = parts[1] if len(parts) > 1 else "exact"
+        if field not in model._meta.fields:
+            raise QueryError(f"Unknown field: {field}")
+
+        col = f"{table_alias}.{dialect.quote_ident(field)}"
+
+        if op == "exact":
+            key = f"{param_prefix}{start_index}"
+            return Condition(
+                sql=f"{col} = :{key}", params={key: value}
+            ), start_index + 1
+        if op in {"gt", "gte", "lt", "lte"}:
+            key = f"{param_prefix}{start_index}"
+            op_map = {"gt": ">", "gte": ">=", "lt": "<", "lte": "<="}[op]
+            return Condition(
+                sql=f"{col} {op_map} :{key}", params={key: value}
+            ), start_index + 1
+        if op == "in":
+            values = normalize_in_values(value)
+            if not values:
+                return Condition(sql="1=0", params={}), start_index
+            keys = []
+            params: dict[str, Any] = {}
+            idx = start_index
+            for v in values:
+                key = f"{param_prefix}{idx}"
+                idx += 1
+                keys.append(f":{key}")
+                params[key] = v
+            return Condition(sql=f"{col} IN ({', '.join(keys)})", params=params), idx
+        if op in {
+            "contains",
+            "icontains",
+            "startswith",
+            "istartswith",
+            "endswith",
+            "iendswith",
+        }:
+            key = f"{param_prefix}{start_index}"
+            if op.endswith("contains"):
+                pattern = f"%{value}%"
+            elif op.endswith("startswith"):
+                pattern = f"{value}%"
+            else:
+                pattern = f"%{value}"
+            if op.startswith("i"):
+                return (
+                    Condition(
+                        sql=f"LOWER({col}) LIKE LOWER(:{key})", params={key: pattern}
+                    ),
+                    start_index + 1,
+                )
+            return Condition(
+                sql=f"{col} LIKE :{key}", params={key: pattern}
+            ), start_index + 1
+        if op == "isnull":
+            is_null = bool(value)
+            return (
+                Condition(
+                    sql=f"{col} IS {'NULL' if is_null else 'NOT NULL'}", params={}
+                ),
+                start_index,
+            )
+        raise QueryError(f"Unsupported lookup: {lookup}")
+
+    cond, _ = q.compile(
+        table_alias=_table_alias(),
+        param_prefix="p",
+        start_index=1,
+        lookup_compiler=_lookup_compiler,
+    )
+    return cond.sql, dict(cond.params)
+
+
+@dataclass(frozen=True)
+class QueryState:
+    where: Q
+    order_by: tuple[str, ...]
+    limit: Optional[int]
+    offset: Optional[int]
+    prefetch: tuple[str, ...]
+
+
+class QuerySet(Generic[TModel]):
+    def __init__(
+        self, model: type[TModel], *, db: Database, state: Optional[QueryState] = None
+    ):
+        self.model = model
+        self.db = db
+        self.state = state or QueryState(
+            where=Q(), order_by=(), limit=None, offset=None, prefetch=()
+        )
+
+    def filter(self, **lookups: Any) -> QuerySet[TModel]:
+        q = Q(**lookups)
+        combined = self.state.where & q if not self.state.where.is_empty() else q
+        return QuerySet(
+            self.model, db=self.db, state=replace(self.state, where=combined)
+        )
+
+    def exclude(self, **lookups: Any) -> QuerySet[TModel]:
+        q = ~Q(**lookups)
+        combined = self.state.where & q if not self.state.where.is_empty() else q
+        return QuerySet(
+            self.model, db=self.db, state=replace(self.state, where=combined)
+        )
+
+    def order_by(self, *fields: str) -> QuerySet[TModel]:
+        return QuerySet(
+            self.model, db=self.db, state=replace(self.state, order_by=fields)
+        )
+
+    def limit(self, n: int) -> QuerySet[TModel]:
+        return QuerySet(self.model, db=self.db, state=replace(self.state, limit=int(n)))
+
+    def offset(self, n: int) -> QuerySet[TModel]:
+        return QuerySet(
+            self.model, db=self.db, state=replace(self.state, offset=int(n))
+        )
+
+    def prefetch_related(self, *names: str) -> QuerySet[TModel]:
+        merged = tuple(dict.fromkeys(self.state.prefetch + tuple(names)))
+        return QuerySet(
+            self.model, db=self.db, state=replace(self.state, prefetch=merged)
+        )
+
+    def select_related(self, *names: str) -> QuerySet[TModel]:
+        return self.prefetch_related(*names)
+
+    def _select_columns(self) -> list[str]:
+        cols = []
+        alias = _table_alias()
+        for name in self.model._meta.fields:
+            col = self.db.dialect.quote_ident(name)
+            cols.append(f"{alias}.{col} AS {col}")
+        return cols
+
+    def _order_by_sql(self) -> str:
+        if not self.state.order_by:
+            return ""
+        alias = _table_alias()
+        parts: list[str] = []
+        for f in self.state.order_by:
+            direction = "ASC"
+            name = f
+            if f.startswith("-"):
+                direction = "DESC"
+                name = f[1:]
+            if name not in self.model._meta.fields:
+                raise QueryError(f"Unknown order_by field: {name}")
+            parts.append(f"{alias}.{self.db.dialect.quote_ident(name)} {direction}")
+        return " ORDER BY " + ", ".join(parts)
+
+    def _select_sql(self) -> tuple[str, dict[str, Any]]:
+        where_sql, params = _where_from_q(self.model, self.db.dialect, self.state.where)
+        cols = ", ".join(self._select_columns())
+        sql = (
+            f"SELECT {cols} FROM {_qualified_table(self.model, self.db.dialect)} "
+            f"{_table_alias()} WHERE {where_sql}"
+        )
+        sql += self._order_by_sql()
+        sql = self.db.dialect.apply_limit_offset(
+            sql, self.state.limit, self.state.offset
+        )
+        return sql, params
+
+    def all(self) -> list[TModel]:
+        sql, params = self._select_sql()
+        rows = self.db.fetch_all(sql, params)
+        instances = [self.model.parse_obj(r) for r in rows]
+        if self.state.prefetch and instances:
+            self._prefetch(instances)
+        return instances
+
+    def first(self) -> Optional[TModel]:
+        items = self.limit(1).all()
+        return items[0] if items else None
+
+    def get(self, **lookups: Any) -> TModel:
+        items = self.filter(**lookups).limit(2).all()
+        if not items:
+            raise DoesNotExist(f"{self.model.__name__} matching query does not exist")
+        if len(items) > 1:
+            raise MultipleObjectsReturned(f"Multiple {self.model.__name__} returned")
+        return items[0]
+
+    def create(self, **data: Any) -> TModel:
+        inst = self.model(**data)
+        self.save_instance(inst)
+        return inst
+
+    def save_instance(self, inst: Any) -> Any:
+        pk_name = inst.pk_name()
+        pk_value = getattr(inst, pk_name, None)
+        data = inst.to_db_dict()
+        if pk_value is None:
+            new_pk = self._insert(data)
+            if new_pk is not None:
+                setattr(inst, pk_name, new_pk)
+            return inst
+        self.filter(**{pk_name: pk_value}).update(**data)
+        return inst
+
+    def _insert(self, data: dict[str, Any]) -> Any:
+        cols = [
+            c
+            for c in data.keys()
+            if c in self.model._meta.fields and data[c] is not None
+        ]
+        if not cols:
+            raise QueryError("No insertable columns provided")
+        col_sql = ", ".join(self.db.dialect.quote_ident(c) for c in cols)
+        values_sql = ", ".join(f":p{i + 1}" for i in range(len(cols)))
+        params = {f"p{i + 1}": data[c] for i, c in enumerate(cols)}
+        table = _qualified_table(self.model, self.db.dialect)
+        pk = self.model.pk_name()
+        if self.db.dialect.name == "postgres":
+            sql = (
+                f"INSERT INTO {table} ({col_sql}) VALUES ({values_sql}) "
+                f"RETURNING {self.db.dialect.quote_ident(pk)}"
+            )
+            row = self.db.fetch_one(sql, params)
+            return None if row is None else row.get(pk)
+
+        sql = f"INSERT INTO {table} ({col_sql}) VALUES ({values_sql})"
+        self.db.execute(sql, params)
+        if self.db.dialect.name == "sqlite":
+            row = self.db.fetch_one("SELECT last_insert_rowid() AS id")
+            return None if row is None else row.get("id")
+        return None
+
+    def update(self, **data: Any) -> int:
+        set_cols = [k for k in data.keys() if k in self.model._meta.fields]
+        if not set_cols:
+            return 0
+        where_sql, where_params = _where_from_q(
+            self.model, self.db.dialect, self.state.where
+        )
+        sets = []
+        params: dict[str, Any] = {}
+        idx = 1
+        for c in set_cols:
+            key = f"p{idx}"
+            idx += 1
+            sets.append(f"{self.db.dialect.quote_ident(c)} = :{key}")
+            params[key] = data[c]
+        params.update(where_params)
+        sql = (
+            f"UPDATE {_qualified_table(self.model, self.db.dialect)} "
+            f"SET {', '.join(sets)} WHERE {where_sql}"
+        )
+        return self.db.execute(sql, params)
+
+    def delete(self) -> int:
+        where_sql, params = _where_from_q(self.model, self.db.dialect, self.state.where)
+        table = _qualified_table(self.model, self.db.dialect)
+        sql = f"DELETE FROM {table} WHERE {where_sql}"
+        return self.db.execute(sql, params)
+
+    def _prefetch(self, instances: list[Any]) -> None:
+        for name in self.state.prefetch:
+            if name in self.model._meta.relations:
+                self._prefetch_fk(instances, name)
+            elif name in self.model._meta.m2m:
+                self._prefetch_m2m(instances, name)
+            else:
+                raise QueryError(f"Unknown relation for prefetch: {name}")
+
+    def _prefetch_fk(self, instances: list[Any], rel_name: str) -> None:
+        fk_col, fk = self.model._meta.relations[rel_name]
+        to_model = fk.to
+        if isinstance(to_model, str):
+            to_model = self.model._resolve_model(to_model)
+        ids = [
+            getattr(i, fk_col)
+            for i in instances
+            if getattr(i, fk_col, None) is not None
+        ]
+        ids = list(dict.fromkeys(ids))
+        if not ids:
+            for i in instances:
+                setattr(i, rel_name, None)
+            return
+        related = (
+            QuerySet(to_model, db=self.db)
+            .filter(**{to_model.pk_name() + "__in": ids})
+            .all()
+        )
+        mapping = {getattr(r, to_model.pk_name()): r for r in related}
+        for i in instances:
+            setattr(i, rel_name, mapping.get(getattr(i, fk_col)))
+
+    def _prefetch_m2m(self, instances: list[Any], field_name: str) -> None:
+        rel = self.model._meta.m2m[field_name]
+        to_model = rel.to
+        if isinstance(to_model, str):
+            to_model = self.model._resolve_model(to_model)
+        from_pk = self.model.pk_name()
+        to_pk = to_model.pk_name()
+        through = rel.through or f"{self.model.table_name()}_{to_model.table_name()}"
+        from_col = rel.from_field or f"{self.model.table_name()}_{from_pk}"
+        to_col = rel.to_field or f"{to_model.table_name()}_{to_pk}"
+
+        base_ids = [
+            getattr(i, from_pk)
+            for i in instances
+            if getattr(i, from_pk, None) is not None
+        ]
+        base_ids = list(dict.fromkeys(base_ids))
+        if not base_ids:
+            for i in instances:
+                setattr(i, field_name, [])
+            return
+
+        table = self.db.dialect.quote_ident(through)
+        sql = (
+            f"SELECT {self.db.dialect.quote_ident(from_col)} AS from_id, "
+            f"{self.db.dialect.quote_ident(to_col)} AS to_id "
+            f"FROM {table} WHERE {self.db.dialect.quote_ident(from_col)} IN ("
+        )
+        params: dict[str, Any] = {}
+        keys = []
+        for idx, v in enumerate(base_ids, start=1):
+            key = f"p{idx}"
+            keys.append(f":{key}")
+            params[key] = v
+        sql += ", ".join(keys) + ")"
+        pairs = self.db.fetch_all(sql, params)
+        to_ids = list(dict.fromkeys([p["to_id"] for p in pairs]))
+        if not to_ids:
+            for i in instances:
+                setattr(i, field_name, [])
+            return
+
+        related = (
+            QuerySet(to_model, db=self.db).filter(**{to_pk + "__in": to_ids}).all()
+        )
+        to_map = {getattr(r, to_pk): r for r in related}
+        by_from: dict[Any, list[Any]] = {}
+        for p in pairs:
+            by_from.setdefault(p["from_id"], []).append(to_map.get(p["to_id"]))
+        for i in instances:
+            setattr(
+                i,
+                field_name,
+                [x for x in by_from.get(getattr(i, from_pk), []) if x is not None],
+            )
+
+
+class AsyncQuerySet(Generic[TModel]):
+    def __init__(
+        self,
+        model: type[TModel],
+        *,
+        db: AsyncDatabase,
+        state: Optional[QueryState] = None,
+    ) -> None:
+        self.model = model
+        self.db = db
+        self.state = state or QueryState(
+            where=Q(), order_by=(), limit=None, offset=None, prefetch=()
+        )
+
+    def filter(self, **lookups: Any) -> AsyncQuerySet[TModel]:
+        q = Q(**lookups)
+        combined = self.state.where & q if not self.state.where.is_empty() else q
+        return AsyncQuerySet(
+            self.model, db=self.db, state=replace(self.state, where=combined)
+        )
+
+    def exclude(self, **lookups: Any) -> AsyncQuerySet[TModel]:
+        q = ~Q(**lookups)
+        combined = self.state.where & q if not self.state.where.is_empty() else q
+        return AsyncQuerySet(
+            self.model, db=self.db, state=replace(self.state, where=combined)
+        )
+
+    def order_by(self, *fields: str) -> AsyncQuerySet[TModel]:
+        return AsyncQuerySet(
+            self.model, db=self.db, state=replace(self.state, order_by=fields)
+        )
+
+    def limit(self, n: int) -> AsyncQuerySet[TModel]:
+        return AsyncQuerySet(
+            self.model, db=self.db, state=replace(self.state, limit=int(n))
+        )
+
+    def offset(self, n: int) -> AsyncQuerySet[TModel]:
+        return AsyncQuerySet(
+            self.model, db=self.db, state=replace(self.state, offset=int(n))
+        )
+
+    def prefetch_related(self, *names: str) -> AsyncQuerySet[TModel]:
+        merged = tuple(dict.fromkeys(self.state.prefetch + tuple(names)))
+        return AsyncQuerySet(
+            self.model, db=self.db, state=replace(self.state, prefetch=merged)
+        )
+
+    def select_related(self, *names: str) -> AsyncQuerySet[TModel]:
+        return self.prefetch_related(*names)
+
+    def _select_columns(self) -> list[str]:
+        cols = []
+        alias = _table_alias()
+        for name in self.model._meta.fields:
+            col = self.db.dialect.quote_ident(name)
+            cols.append(f"{alias}.{col} AS {col}")
+        return cols
+
+    def _order_by_sql(self) -> str:
+        if not self.state.order_by:
+            return ""
+        alias = _table_alias()
+        parts: list[str] = []
+        for f in self.state.order_by:
+            direction = "ASC"
+            name = f
+            if f.startswith("-"):
+                direction = "DESC"
+                name = f[1:]
+            if name not in self.model._meta.fields:
+                raise QueryError(f"Unknown order_by field: {name}")
+            parts.append(f"{alias}.{self.db.dialect.quote_ident(name)} {direction}")
+        return " ORDER BY " + ", ".join(parts)
+
+    def _select_sql(self) -> tuple[str, dict[str, Any]]:
+        where_sql, params = _where_from_q(self.model, self.db.dialect, self.state.where)
+        cols = ", ".join(self._select_columns())
+        sql = (
+            f"SELECT {cols} FROM {_qualified_table(self.model, self.db.dialect)} "
+            f"{_table_alias()} WHERE {where_sql}"
+        )
+        sql += self._order_by_sql()
+        sql = self.db.dialect.apply_limit_offset(
+            sql, self.state.limit, self.state.offset
+        )
+        return sql, params
+
+    async def all(self) -> list[TModel]:
+        sql, params = self._select_sql()
+        rows = await self.db.fetch_all(sql, params)
+        instances = [self.model.parse_obj(r) for r in rows]
+        if self.state.prefetch and instances:
+            await self._prefetch(instances)
+        return instances
+
+    async def first(self) -> Optional[TModel]:
+        items = await self.limit(1).all()
+        return items[0] if items else None
+
+    async def get(self, **lookups: Any) -> TModel:
+        items = await self.filter(**lookups).limit(2).all()
+        if not items:
+            raise DoesNotExist(f"{self.model.__name__} matching query does not exist")
+        if len(items) > 1:
+            raise MultipleObjectsReturned(f"Multiple {self.model.__name__} returned")
+        return items[0]
+
+    async def create(self, **data: Any) -> TModel:
+        inst = self.model(**data)
+        await self.save_instance(inst)
+        return inst
+
+    async def save_instance(self, inst: Any) -> Any:
+        pk_name = inst.pk_name()
+        pk_value = getattr(inst, pk_name, None)
+        data = inst.to_db_dict()
+        if pk_value is None:
+            new_pk = await self._insert(data)
+            if new_pk is not None:
+                setattr(inst, pk_name, new_pk)
+            return inst
+        await self.filter(**{pk_name: pk_value}).update(**data)
+        return inst
+
+    async def _insert(self, data: dict[str, Any]) -> Any:
+        cols = [
+            c
+            for c in data.keys()
+            if c in self.model._meta.fields and data[c] is not None
+        ]
+        if not cols:
+            raise QueryError("No insertable columns provided")
+        col_sql = ", ".join(self.db.dialect.quote_ident(c) for c in cols)
+        values_sql = ", ".join(f":p{i + 1}" for i in range(len(cols)))
+        params = {f"p{i + 1}": data[c] for i, c in enumerate(cols)}
+        table = _qualified_table(self.model, self.db.dialect)
+        pk = self.model.pk_name()
+        if self.db.dialect.name == "postgres":
+            sql = (
+                f"INSERT INTO {table} ({col_sql}) VALUES ({values_sql}) "
+                f"RETURNING {self.db.dialect.quote_ident(pk)}"
+            )
+            row = await self.db.fetch_one(sql, params)
+            return None if row is None else row.get(pk)
+
+        sql = f"INSERT INTO {table} ({col_sql}) VALUES ({values_sql})"
+        await self.db.execute(sql, params)
+        if self.db.dialect.name == "sqlite":
+            row = await self.db.fetch_one("SELECT last_insert_rowid() AS id")
+            return None if row is None else row.get("id")
+        return None
+
+    async def update(self, **data: Any) -> int:
+        set_cols = [k for k in data.keys() if k in self.model._meta.fields]
+        if not set_cols:
+            return 0
+        where_sql, where_params = _where_from_q(
+            self.model, self.db.dialect, self.state.where
+        )
+        sets = []
+        params: dict[str, Any] = {}
+        idx = 1
+        for c in set_cols:
+            key = f"p{idx}"
+            idx += 1
+            sets.append(f"{self.db.dialect.quote_ident(c)} = :{key}")
+            params[key] = data[c]
+        params.update(where_params)
+        sql = (
+            f"UPDATE {_qualified_table(self.model, self.db.dialect)} "
+            f"SET {', '.join(sets)} WHERE {where_sql}"
+        )
+        return await self.db.execute(sql, params)
+
+    async def delete(self) -> int:
+        where_sql, params = _where_from_q(self.model, self.db.dialect, self.state.where)
+        table = _qualified_table(self.model, self.db.dialect)
+        sql = f"DELETE FROM {table} WHERE {where_sql}"
+        return await self.db.execute(sql, params)
+
+    async def _prefetch(self, instances: list[Any]) -> None:
+        for name in self.state.prefetch:
+            if name in self.model._meta.relations:
+                await self._prefetch_fk(instances, name)
+            elif name in self.model._meta.m2m:
+                await self._prefetch_m2m(instances, name)
+            else:
+                raise QueryError(f"Unknown relation for prefetch: {name}")
+
+    async def _prefetch_fk(self, instances: list[Any], rel_name: str) -> None:
+        fk_col, fk = self.model._meta.relations[rel_name]
+        to_model = fk.to
+        if isinstance(to_model, str):
+            to_model = self.model._resolve_model(to_model)
+        ids = [
+            getattr(i, fk_col)
+            for i in instances
+            if getattr(i, fk_col, None) is not None
+        ]
+        ids = list(dict.fromkeys(ids))
+        if not ids:
+            for i in instances:
+                setattr(i, rel_name, None)
+            return
+        related = (
+            await AsyncQuerySet(to_model, db=self.db)
+            .filter(**{to_model.pk_name() + "__in": ids})
+            .all()
+        )
+        mapping = {getattr(r, to_model.pk_name()): r for r in related}
+        for i in instances:
+            setattr(i, rel_name, mapping.get(getattr(i, fk_col)))
+
+    async def _prefetch_m2m(self, instances: list[Any], field_name: str) -> None:
+        rel = self.model._meta.m2m[field_name]
+        to_model = rel.to
+        if isinstance(to_model, str):
+            to_model = self.model._resolve_model(to_model)
+        from_pk = self.model.pk_name()
+        to_pk = to_model.pk_name()
+        through = rel.through or f"{self.model.table_name()}_{to_model.table_name()}"
+        from_col = rel.from_field or f"{self.model.table_name()}_{from_pk}"
+        to_col = rel.to_field or f"{to_model.table_name()}_{to_pk}"
+
+        base_ids = [
+            getattr(i, from_pk)
+            for i in instances
+            if getattr(i, from_pk, None) is not None
+        ]
+        base_ids = list(dict.fromkeys(base_ids))
+        if not base_ids:
+            for i in instances:
+                setattr(i, field_name, [])
+            return
+
+        table = self.db.dialect.quote_ident(through)
+        sql = (
+            f"SELECT {self.db.dialect.quote_ident(from_col)} AS from_id, "
+            f"{self.db.dialect.quote_ident(to_col)} AS to_id "
+            f"FROM {table} WHERE {self.db.dialect.quote_ident(from_col)} IN ("
+        )
+        params: dict[str, Any] = {}
+        keys = []
+        for idx, v in enumerate(base_ids, start=1):
+            key = f"p{idx}"
+            keys.append(f":{key}")
+            params[key] = v
+        sql += ", ".join(keys) + ")"
+        pairs = await self.db.fetch_all(sql, params)
+        to_ids = list(dict.fromkeys([p["to_id"] for p in pairs]))
+        if not to_ids:
+            for i in instances:
+                setattr(i, field_name, [])
+            return
+
+        related = (
+            await AsyncQuerySet(to_model, db=self.db)
+            .filter(**{to_pk + "__in": to_ids})
+            .all()
+        )
+        to_map = {getattr(r, to_pk): r for r in related}
+        by_from: dict[Any, list[Any]] = {}
+        for p in pairs:
+            by_from.setdefault(p["from_id"], []).append(to_map.get(p["to_id"]))
+        for i in instances:
+            setattr(
+                i,
+                field_name,
+                [x for x in by_from.get(getattr(i, from_pk), []) if x is not None],
+            )