piccolo 1.8.0__py3-none-any.whl → 1.10.0__py3-none-any.whl

piccolo/engine/postgres.py

@@ -4,7 +4,15 @@ import contextvars
 import typing as t
 from dataclasses import dataclass
 
-from piccolo.engine.base import Batch, Engine, validate_savepoint_name
+from typing_extensions import Self
+
+from piccolo.engine.base import (
+    BaseAtomic,
+    BaseBatch,
+    BaseTransaction,
+    Engine,
+    validate_savepoint_name,
+)
 from piccolo.engine.exceptions import TransactionError
 from piccolo.query.base import DDL, Query
 from piccolo.querystring import QueryString
@@ -18,16 +26,17 @@ if t.TYPE_CHECKING:  # pragma: no cover
     from asyncpg.connection import Connection
     from asyncpg.cursor import Cursor
     from asyncpg.pool import Pool
+    from asyncpg.transaction import Transaction
 
 
 @dataclass
-class AsyncBatch(Batch):
+class AsyncBatch(BaseBatch):
     connection: Connection
     query: Query
     batch_size: int
 
     # Set internally
-    _transaction = None
+    _transaction: t.Optional[Transaction] = None
     _cursor: t.Optional[Cursor] = None
 
     @property
@@ -36,20 +45,26 @@ class AsyncBatch(Batch):
             raise ValueError("_cursor not set")
         return self._cursor
 
+    @property
+    def transaction(self) -> Transaction:
+        if not self._transaction:
+            raise ValueError("The transaction can't be found.")
+        return self._transaction
+
     async def next(self) -> t.List[t.Dict]:
         data = await self.cursor.fetch(self.batch_size)
         return await self.query._process_results(data)
 
-    def __aiter__(self):
+    def __aiter__(self: Self) -> Self:
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> t.List[t.Dict]:
         response = await self.next()
         if response == []:
             raise StopAsyncIteration()
         return response
 
-    async def __aenter__(self):
+    async def __aenter__(self: Self) -> Self:
         self._transaction = self.connection.transaction()
         await self._transaction.start()
         querystring = self.query.querystrings[0]
@@ -60,9 +75,9 @@ class AsyncBatch(Batch):
 
     async def __aexit__(self, exception_type, exception, traceback):
         if exception:
-            await self._transaction.rollback()
+            await self.transaction.rollback()
         else:
-            await self._transaction.commit()
+            await self.transaction.commit()
 
         await self.connection.close()
 
@@ -72,7 +87,7 @@ class AsyncBatch(Batch):
 ###############################################################################
 
 
-class Atomic:
+class Atomic(BaseAtomic):
     """
     This is useful if you want to build up a transaction programatically, by
     adding queries to it.
@@ -140,7 +155,7 @@ class Savepoint:
         )
 
 
-class PostgresTransaction:
+class PostgresTransaction(BaseTransaction):
     """
     Used for wrapping queries in a transaction, using a context manager.
     Currently it's async only.
@@ -243,7 +258,7 @@ class PostgresTransaction:
 
     ###########################################################################
 
-    async def __aexit__(self, exception_type, exception, traceback):
+    async def __aexit__(self, exception_type, exception, traceback) -> bool:
         if self._parent:
             return exception is None
 
@@ -269,7 +284,7 @@ class PostgresTransaction:
 ###############################################################################
 
 
-class PostgresEngine(Engine[t.Optional[PostgresTransaction]]):
+class PostgresEngine(Engine[PostgresTransaction]):
     """
     Used to connect to PostgreSQL.
 
@@ -331,16 +346,10 @@ class PostgresEngine(Engine[t.Optional[PostgresTransaction]]):
     __slots__ = (
         "config",
         "extensions",
-        "log_queries",
-        "log_responses",
         "extra_nodes",
         "pool",
-        "current_transaction",
     )
 
-    engine_type = "postgres"
-    min_version_number = 10
-
     def __init__(
         self,
         config: t.Dict[str, t.Any],
@@ -362,7 +371,12 @@ class PostgresEngine(Engine[t.Optional[PostgresTransaction]]):
         self.current_transaction = contextvars.ContextVar(
             f"pg_current_transaction_{database_name}", default=None
         )
-        super().__init__()
+        super().__init__(
+            engine_type="postgres",
+            log_queries=log_queries,
+            log_responses=log_responses,
+            min_version_number=10,
+        )
 
     @staticmethod
     def _parse_raw_version_string(version_string: str) -> float:

piccolo/engine/sqlite.py

@@ -11,7 +11,15 @@ from dataclasses import dataclass
 from decimal import Decimal
 from functools import partial, wraps
 
-from piccolo.engine.base import Batch, Engine, validate_savepoint_name
+from typing_extensions import Self
+
+from piccolo.engine.base import (
+    BaseAtomic,
+    BaseBatch,
+    BaseTransaction,
+    Engine,
+    validate_savepoint_name,
+)
 from piccolo.engine.exceptions import TransactionError
 from piccolo.query.base import DDL, Query
 from piccolo.querystring import QueryString
@@ -309,7 +317,7 @@ for column_name in ("TIMESTAMP", "TIMESTAMPTZ", "DATE", "TIME"):
 
 
 @dataclass
-class AsyncBatch(Batch):
+class AsyncBatch(BaseBatch):
     connection: Connection
     query: Query
     batch_size: int
@@ -327,16 +335,16 @@ class AsyncBatch(Batch):
         data = await self.cursor.fetchmany(self.batch_size)
         return await self.query._process_results(data)
 
-    def __aiter__(self):
+    def __aiter__(self: Self) -> Self:
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> t.List[t.Dict]:
         response = await self.next()
         if response == []:
             raise StopAsyncIteration()
         return response
 
-    async def __aenter__(self):
+    async def __aenter__(self: Self) -> Self:
         querystring = self.query.querystrings[0]
         template, template_args = querystring.compile_string()
 
@@ -344,7 +352,7 @@ class AsyncBatch(Batch):
         return self
 
     async def __aexit__(self, exception_type, exception, traceback):
-        await self._cursor.close()
+        await self.cursor.close()
         await self.connection.close()
         return exception is not None
 
@@ -363,7 +371,7 @@ class TransactionType(enum.Enum):
     exclusive = "EXCLUSIVE"
 
 
-class Atomic:
+class Atomic(BaseAtomic):
     """
     Usage:
 
@@ -384,9 +392,9 @@ class Atomic:
     ):
         self.engine = engine
         self.transaction_type = transaction_type
-        self.queries: t.List[Query] = []
+        self.queries: t.List[t.Union[Query, DDL]] = []
 
-    def add(self, *query: Query):
+    def add(self, *query: t.Union[Query, DDL]):
         self.queries += list(query)
 
     async def run(self):
@@ -434,7 +442,7 @@ class Savepoint:
         )
 
 
-class SQLiteTransaction:
+class SQLiteTransaction(BaseTransaction):
     """
     Used for wrapping queries in a transaction, using a context manager.
     Currently it's async only.
@@ -534,7 +542,7 @@ class SQLiteTransaction:
 
     ###########################################################################
 
-    async def __aexit__(self, exception_type, exception, traceback):
+    async def __aexit__(self, exception_type, exception, traceback) -> bool:
         if self._parent:
             return exception is None
 
@@ -560,16 +568,8 @@ def dict_factory(cursor, row) -> t.Dict:
     return {col[0]: row[idx] for idx, col in enumerate(cursor.description)}
 
 
-class SQLiteEngine(Engine[t.Optional[SQLiteTransaction]]):
-    __slots__ = (
-        "connection_kwargs",
-        "current_transaction",
-        "log_queries",
-        "log_responses",
-    )
-
-    engine_type = "sqlite"
-    min_version_number = 3.25
+class SQLiteEngine(Engine[SQLiteTransaction]):
+    __slots__ = ("connection_kwargs",)
 
     def __init__(
         self,
@@ -613,7 +613,12 @@ class SQLiteEngine(Engine[t.Optional[SQLiteTransaction]]):
             f"sqlite_current_transaction_{path}", default=None
         )
 
-        super().__init__()
+        super().__init__(
+            engine_type="sqlite",
+            min_version_number=3.25,
+            log_queries=log_queries,
+            log_responses=log_responses,
+        )
 
     @property
     def path(self):

piccolo/query/functions/__init__.py

@@ -1,17 +1,22 @@
 from .aggregate import Avg, Count, Max, Min, Sum
+from .math import Abs, Ceil, Floor, Round
 from .string import Length, Lower, Ltrim, Reverse, Rtrim, Upper
 from .type_conversion import Cast
 
 __all__ = (
+    "Abs",
     "Avg",
     "Cast",
+    "Ceil",
     "Count",
+    "Floor",
     "Length",
     "Lower",
     "Ltrim",
     "Max",
     "Min",
     "Reverse",
+    "Round",
     "Rtrim",
     "Sum",
     "Upper",

piccolo/query/functions/math.py

@@ -0,0 +1,48 @@
+"""
+These functions mirror their counterparts in the Postgresql docs:
+
+https://www.postgresql.org/docs/current/functions-math.html
+
+"""
+
+from .base import Function
+
+
+class Abs(Function):
+    """
+    Absolute value.
+    """
+
+    function_name = "ABS"
+
+
+class Ceil(Function):
+    """
+    Nearest integer greater than or equal to argument.
+    """
+
+    function_name = "CEIL"
+
+
+class Floor(Function):
+    """
+    Nearest integer less than or equal to argument.
+    """
+
+    function_name = "FLOOR"
+
+
+class Round(Function):
+    """
+    Rounds to nearest integer.
+    """
+
+    function_name = "ROUND"
+
+
+__all__ = (
+    "Abs",
+    "Ceil",
+    "Floor",
+    "Round",
+)

piccolo/query/methods/create_index.py

@@ -14,7 +14,7 @@ class CreateIndex(DDL):
     def __init__(
         self,
         table: t.Type[Table],
-        columns: t.List[t.Union[Column, str]],
+        columns: t.Union[t.List[Column], t.List[str]],
         method: IndexMethod = IndexMethod.btree,
         if_not_exists: bool = False,
         **kwargs,

piccolo/query/methods/drop_index.py

@@ -14,7 +14,7 @@ class DropIndex(Query):
     def __init__(
         self,
         table: t.Type[Table],
-        columns: t.List[t.Union[Column, str]],
+        columns: t.Union[t.List[Column], t.List[str]],
         if_exists: bool = True,
         **kwargs,
     ):

piccolo/query/methods/objects.py

@@ -5,7 +5,7 @@ import typing as t
 from piccolo.columns.column_types import ForeignKey
 from piccolo.columns.combination import And, Where
 from piccolo.custom_types import Combinable, TableInstance
-from piccolo.engine.base import Batch
+from piccolo.engine.base import BaseBatch
 from piccolo.query.base import Query
 from piccolo.query.methods.select import Select
 from piccolo.query.mixins import (
@@ -268,17 +268,17 @@ class Objects(
 
     ###########################################################################
 
-    def first(self: Self) -> First[TableInstance]:
+    def first(self) -> First[TableInstance]:
         self.limit_delegate.limit(1)
         return First[TableInstance](query=self)
 
-    def get(self: Self, where: Combinable) -> Get[TableInstance]:
+    def get(self, where: Combinable) -> Get[TableInstance]:
         self.where_delegate.where(where)
         self.limit_delegate.limit(1)
         return Get[TableInstance](query=First[TableInstance](query=self))
 
     def get_or_create(
-        self: Self,
+        self,
         where: Combinable,
         defaults: t.Optional[t.Dict[Column, t.Any]] = None,
     ) -> GetOrCreate[TableInstance]:
@@ -288,17 +288,17 @@ class Objects(
             query=self, table_class=self.table, where=where, defaults=defaults
         )
 
-    def create(self: Self, **columns: t.Any) -> Create[TableInstance]:
+    def create(self, **columns: t.Any) -> Create[TableInstance]:
         return Create[TableInstance](table_class=self.table, columns=columns)
 
     ###########################################################################
 
     async def batch(
-        self: Self,
+        self,
         batch_size: t.Optional[int] = None,
         node: t.Optional[str] = None,
         **kwargs,
-    ) -> Batch:
+    ) -> BaseBatch:
         if batch_size:
             kwargs.update(batch_size=batch_size)
         if node:

piccolo/query/methods/select.py

@@ -5,11 +5,11 @@ import typing as t
 from collections import OrderedDict
 
 from piccolo.columns import Column, Selectable
-from piccolo.columns.column_types import JSON, JSONB, PrimaryKey
+from piccolo.columns.column_types import JSON, JSONB
 from piccolo.columns.m2m import M2MSelect
 from piccolo.columns.readable import Readable
 from piccolo.custom_types import TableInstance
-from piccolo.engine.base import Batch
+from piccolo.engine.base import BaseBatch
 from piccolo.query.base import Query
 from piccolo.query.mixins import (
     AsOfDelegate,
@@ -217,7 +217,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         self,
         response: t.List[t.Dict[str, t.Any]],
         secondary_table: t.Type[Table],
-        secondary_table_pk: PrimaryKey,
+        secondary_table_pk: Column,
         m2m_name: str,
         m2m_select: M2MSelect,
         as_list: bool = False,
@@ -386,14 +386,20 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         return self
 
     @t.overload
-    def output(self: Self, *, as_list: bool) -> SelectList: ...
+    def output(self: Self, *, as_list: bool) -> SelectList:  # type: ignore
+        ...
 
     @t.overload
-    def output(self: Self, *, as_json: bool) -> SelectJSON: ...
+    def output(self: Self, *, as_json: bool) -> SelectJSON:  # type: ignore
+        ...
 
     @t.overload
     def output(self: Self, *, load_json: bool) -> Self: ...
 
+    @t.overload
+    def output(self: Self, *, load_json: bool, as_list: bool) -> SelectJSON:  # type: ignore  # noqa: E501
+        ...
+
     @t.overload
     def output(self: Self, *, nested: bool) -> Self: ...
 
@@ -404,7 +410,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         as_json: bool = False,
         load_json: bool = False,
         nested: bool = False,
-    ):
+    ) -> t.Union[Self, SelectJSON, SelectList]:
         self.output_delegate.output(
             as_list=as_list,
             as_json=as_json,
@@ -436,7 +442,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         batch_size: t.Optional[int] = None,
         node: t.Optional[str] = None,
         **kwargs,
-    ) -> Batch:
+    ) -> BaseBatch:
         if batch_size:
             kwargs.update(batch_size=batch_size)
         if node:

piccolo/query/mixins.py

@@ -207,7 +207,6 @@ class Returning:
 
 @dataclass
 class Output:
-
     as_json: bool = False
     as_list: bool = False
     as_objects: bool = False
@@ -236,7 +235,6 @@ class Callback:
 
 @dataclass
 class WhereDelegate:
-
     _where: t.Optional[Combinable] = None
     _where_columns: t.List[Column] = field(default_factory=list)
 
@@ -246,7 +244,8 @@ class WhereDelegate:
         needed.
         """
         self._where_columns = []
-        self._extract_columns(self._where)
+        if self._where is not None:
+            self._extract_columns(self._where)
         return self._where_columns
 
     def _extract_columns(self, combinable: Combinable):
@@ -277,7 +276,6 @@ class WhereDelegate:
 
 @dataclass
 class OrderByDelegate:
-
     _order_by: OrderBy = field(default_factory=OrderBy)
 
     def get_order_by_columns(self) -> t.List[Column]:
@@ -303,7 +301,6 @@ class OrderByDelegate:
 
 @dataclass
 class LimitDelegate:
-
     _limit: t.Optional[Limit] = None
     _first: bool = False
 
@@ -330,7 +327,6 @@ class AsOfDelegate:
 
 @dataclass
 class DistinctDelegate:
-
     _distinct: Distinct = field(
         default_factory=lambda: Distinct(enabled=False, on=None)
     )
@@ -356,7 +352,6 @@ class ReturningDelegate:
 
 @dataclass
 class CountDelegate:
-
     _count: bool = False
 
     def count(self):
@@ -365,7 +360,6 @@ class CountDelegate:
 
 @dataclass
 class AddDelegate:
-
     _add: t.List[Table] = field(default_factory=list)
 
     def add(self, *instances: Table, table_class: t.Type[Table]):
@@ -421,8 +415,7 @@ class OutputDelegate:
             self._output.nested = bool(nested)
 
     def copy(self) -> OutputDelegate:
-        _output = self._output.copy() if self._output is not None else None
-        return self.__class__(_output=_output)
+        return self.__class__(_output=self._output.copy())
 
 
 @dataclass

piccolo/querystring.py

@@ -282,12 +282,30 @@ class QueryString(Selectable):
     def __le__(self, value) -> QueryString:
         return QueryString("{} <= {}", self, value)
 
+    def __truediv__(self, value) -> QueryString:
+        return QueryString("{} / {}", self, value)
+
+    def __mul__(self, value) -> QueryString:
+        return QueryString("{} * {}", self, value)
+
+    def __pow__(self, value) -> QueryString:
+        return QueryString("{} ^ {}", self, value)
+
+    def __mod__(self, value) -> QueryString:
+        return QueryString("{} % {}", self, value)
+
     def is_in(self, value) -> QueryString:
         return QueryString("{} IN {}", self, value)
 
     def not_in(self, value) -> QueryString:
         return QueryString("{} NOT IN {}", self, value)
 
+    def like(self, value: str) -> QueryString:
+        return QueryString("{} LIKE {}", self, value)
+
+    def ilike(self, value: str) -> QueryString:
+        return QueryString("{} ILIKE {}", self, value)
+
 
 class Unquoted(QueryString):
     """

piccolo/schema.py

@@ -10,10 +10,10 @@ from piccolo.utils.sync import run_sync
 
 
 class SchemaDDLBase(abc.ABC):
-
     db: Engine
 
-    @abc.abstractproperty
+    @property
+    @abc.abstractmethod
     def ddl(self) -> str:
         pass
 
@@ -131,16 +131,19 @@ class ListTables:
         self.db = db
         self.schema_name = schema_name
 
-    async def run(self):
-        response = await self.db.run_querystring(
-            QueryString(
-                """
+    async def run(self) -> t.List[str]:
+        response = t.cast(
+            t.List[t.Dict],
+            await self.db.run_querystring(
+                QueryString(
+                    """
                 SELECT table_name
                 FROM information_schema.tables
                 WHERE table_schema = {}
                 """,
-                self.schema_name,
-            )
+                    self.schema_name,
+                )
+            ),
         )
         return [i["table_name"] for i in response]
 
@@ -155,9 +158,14 @@ class ListSchemas:
     def __init__(self, db: Engine):
         self.db = db
 
-    async def run(self):
-        response = await self.db.run_querystring(
-            QueryString("SELECT schema_name FROM information_schema.schemata")
+    async def run(self) -> t.List[str]:
+        response = t.cast(
+            t.List[t.Dict],
+            await self.db.run_querystring(
+                QueryString(
+                    "SELECT schema_name FROM information_schema.schemata"
+                )
+            ),
         )
         return [i["schema_name"] for i in response]
 
@@ -179,7 +187,7 @@ class SchemaManager:
         """
         db = db or engine_finder()
 
-        if not db:
+        if db is None:
             raise ValueError("The DB can't be found.")
 
         self.db = db