piccolo 1.9.0__py3-none-any.whl → 1.11.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,6 +1,7 @@
 from .aggregate import Avg, Count, Max, Min, Sum
+from .datetime import Day, Extract, Hour, Month, Second, Strftime, Year
 from .math import Abs, Ceil, Floor, Round
-from .string import Length, Lower, Ltrim, Reverse, Rtrim, Upper
+from .string import Concat, Length, Lower, Ltrim, Reverse, Rtrim, Upper
 from .type_conversion import Cast
 
 __all__ = (
@@ -8,16 +9,25 @@ __all__ = (
     "Avg",
     "Cast",
     "Ceil",
+    "Concat",
     "Count",
+    "Day",
+    "Extract",
+    "Extract",
     "Floor",
+    "Hour",
     "Length",
     "Lower",
     "Ltrim",
     "Max",
     "Min",
+    "Month",
     "Reverse",
     "Round",
     "Rtrim",
+    "Second",
+    "Strftime",
     "Sum",
     "Upper",
+    "Year",
 )

piccolo/query/functions/datetime.py

@@ -0,0 +1,260 @@
+import typing as t
+
+from piccolo.columns.base import Column
+from piccolo.columns.column_types import (
+    Date,
+    Integer,
+    Time,
+    Timestamp,
+    Timestamptz,
+)
+from piccolo.querystring import QueryString
+
+from .type_conversion import Cast
+
+###############################################################################
+# Postgres / Cockroach
+
+ExtractComponent = t.Literal[
+    "century",
+    "day",
+    "decade",
+    "dow",
+    "doy",
+    "epoch",
+    "hour",
+    "isodow",
+    "isoyear",
+    "julian",
+    "microseconds",
+    "millennium",
+    "milliseconds",
+    "minute",
+    "month",
+    "quarter",
+    "second",
+    "timezone",
+    "timezone_hour",
+    "timezone_minute",
+    "week",
+    "year",
+]
+
+
+class Extract(QueryString):
+    def __init__(
+        self,
+        identifier: t.Union[Date, Time, Timestamp, Timestamptz, QueryString],
+        datetime_component: ExtractComponent,
+        alias: t.Optional[str] = None,
+    ):
+        """
+        .. note:: This is for Postgres / Cockroach only.
+
+        Extract a date or time component from a ``Date`` / ``Time`` /
+        ``Timestamp`` / ``Timestamptz`` column. For example, getting the month
+        from a timestamp:
+
+        .. code-block:: python
+
+            >>> from piccolo.query.functions import Extract
+            >>> await Concert.select(
+            ...     Extract(Concert.starts, "month", alias="start_month")
+            ... )
+            [{"start_month": 12}]
+
+        :param identifier:
+            Identifies the column.
+        :param datetime_component:
+            The date or time component to extract from the column.
+
+        """
+        if datetime_component.lower() not in t.get_args(ExtractComponent):
+            raise ValueError("The date time component isn't recognised.")
+
+        super().__init__(
+            f"EXTRACT({datetime_component} FROM {{}})",
+            identifier,
+            alias=alias,
+        )
+
+
+###############################################################################
+# SQLite
+
+
+class Strftime(QueryString):
+    def __init__(
+        self,
+        identifier: t.Union[Date, Time, Timestamp, Timestamptz, QueryString],
+        datetime_format: str,
+        alias: t.Optional[str] = None,
+    ):
+        """
+        .. note:: This is for SQLite only.
+
+        Format a datetime value. For example:
+
+        .. code-block:: python
+
+            >>> from piccolo.query.functions import Strftime
+            >>> await Concert.select(
+            ...     Strftime(Concert.starts, "%Y", alias="start_year")
+            ... )
+            [{"start_month": "2024"}]
+
+        :param identifier:
+            Identifies the column.
+        :param datetime_format:
+            A string describing the output format (see SQLite's
+            `documentation <https://www.sqlite.org/lang_datefunc.html>`_
+            for more info).
+
+        """
+        super().__init__(
+            f"strftime('{datetime_format}', {{}})",
+            identifier,
+            alias=alias,
+        )
+
+
+###############################################################################
+# Database agnostic
+
+
+def _get_engine_type(identifier: t.Union[Column, QueryString]) -> str:
+    if isinstance(identifier, Column):
+        return identifier._meta.engine_type
+    elif isinstance(identifier, QueryString) and (
+        columns := identifier.columns
+    ):
+        return columns[0]._meta.engine_type
+    else:
+        raise ValueError("Unable to determine the engine type")
+
+
+def _extract_component(
+    identifier: t.Union[Date, Time, Timestamp, Timestamptz, QueryString],
+    sqlite_format: str,
+    postgres_format: ExtractComponent,
+    alias: t.Optional[str],
+) -> QueryString:
+    engine_type = _get_engine_type(identifier=identifier)
+
+    return Cast(
+        (
+            Strftime(
+                identifier=identifier,
+                datetime_format=sqlite_format,
+            )
+            if engine_type == "sqlite"
+            else Extract(
+                identifier=identifier,
+                datetime_component=postgres_format,
+            )
+        ),
+        Integer(),
+        alias=alias,
+    )
+
+
+def Year(
+    identifier: t.Union[Date, Timestamp, Timestamptz, QueryString],
+    alias: t.Optional[str] = None,
+) -> QueryString:
+    """
+    Extract the year as an integer.
+    """
+    return _extract_component(
+        identifier=identifier,
+        sqlite_format="%Y",
+        postgres_format="year",
+        alias=alias,
+    )
+
+
+def Month(
+    identifier: t.Union[Date, Timestamp, Timestamptz, QueryString],
+    alias: t.Optional[str] = None,
+) -> QueryString:
+    """
+    Extract the month as an integer.
+    """
+    return _extract_component(
+        identifier=identifier,
+        sqlite_format="%m",
+        postgres_format="month",
+        alias=alias,
+    )
+
+
+def Day(
+    identifier: t.Union[Date, Timestamp, Timestamptz, QueryString],
+    alias: t.Optional[str] = None,
+) -> QueryString:
+    """
+    Extract the day as an integer.
+    """
+    return _extract_component(
+        identifier=identifier,
+        sqlite_format="%d",
+        postgres_format="day",
+        alias=alias,
+    )
+
+
+def Hour(
+    identifier: t.Union[Time, Timestamp, Timestamptz, QueryString],
+    alias: t.Optional[str] = None,
+) -> QueryString:
+    """
+    Extract the hour as an integer.
+    """
+    return _extract_component(
+        identifier=identifier,
+        sqlite_format="%H",
+        postgres_format="hour",
+        alias=alias,
+    )
+
+
+def Minute(
+    identifier: t.Union[Time, Timestamp, Timestamptz, QueryString],
+    alias: t.Optional[str] = None,
+) -> QueryString:
+    """
+    Extract the minute as an integer.
+    """
+    return _extract_component(
+        identifier=identifier,
+        sqlite_format="%M",
+        postgres_format="minute",
+        alias=alias,
+    )
+
+
+def Second(
+    identifier: t.Union[Time, Timestamp, Timestamptz, QueryString],
+    alias: t.Optional[str] = None,
+) -> QueryString:
+    """
+    Extract the second as an integer.
+    """
+    return _extract_component(
+        identifier=identifier,
+        sqlite_format="%S",
+        postgres_format="second",
+        alias=alias,
+    )
+
+
+__all__ = (
+    "Extract",
+    "Strftime",
+    "Year",
+    "Month",
+    "Day",
+    "Hour",
+    "Minute",
+    "Second",
+)

piccolo/query/functions/string.py

@@ -5,6 +5,12 @@ https://www.postgresql.org/docs/current/functions-string.html
 
 """
 
+import typing as t
+
+from piccolo.columns.base import Column
+from piccolo.columns.column_types import Text, Varchar
+from piccolo.querystring import QueryString
+
 from .base import Function
 
 
@@ -63,6 +69,44 @@ class Upper(Function):
     function_name = "UPPER"
 
 
+class Concat(QueryString):
+    def __init__(
+        self,
+        *args: t.Union[Column, QueryString, str],
+        alias: t.Optional[str] = None,
+    ):
+        """
+        Concatenate multiple values into a single string.
+
+        .. note::
+            Null values are ignored, so ``null + '!!!'`` returns ``!!!``,
+            not ``null``.
+
+        .. warning::
+            For SQLite, this is only available in version 3.44.0 and above.
+
+        """
+        if len(args) < 2:
+            raise ValueError("At least two values must be passed in.")
+
+        placeholders = ", ".join("{}" for _ in args)
+
+        processed_args: t.List[t.Union[QueryString, Column]] = []
+
+        for arg in args:
+            if isinstance(arg, str) or (
+                isinstance(arg, Column)
+                and not isinstance(arg, (Varchar, Text))
+            ):
+                processed_args.append(QueryString("CAST({} AS TEXT)", arg))
+            else:
+                processed_args.append(arg)
+
+        super().__init__(
+            f"CONCAT({placeholders})", *processed_args, alias=alias
+        )
+
+
 __all__ = (
     "Length",
     "Lower",
@@ -70,4 +114,5 @@ __all__ = (
     "Reverse",
     "Rtrim",
     "Upper",
+    "Concat",
 )

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: