diff --git a/docs/src/piccolo/functions/array.rst b/docs/src/piccolo/functions/array.rst new file mode 100644 index 000000000..521778dc4 --- /dev/null +++ b/docs/src/piccolo/functions/array.rst @@ -0,0 +1,34 @@ +Array functions +=============== + +.. currentmodule:: piccolo.query.functions.array + +ArrayCat +-------- + +.. autoclass:: ArrayCat + :class-doc-from: class + +ArrayAppend +----------- + +.. autoclass:: ArrayAppend + :class-doc-from: class + +ArrayPrepend +------------ + +.. autoclass:: ArrayPrepend + :class-doc-from: class + +ArrayRemove +----------- + +.. autoclass:: ArrayRemove + :class-doc-from: class + +ArrayReplace +------------ + +.. autoclass:: ArrayReplace + :class-doc-from: class diff --git a/docs/src/piccolo/functions/index.rst b/docs/src/piccolo/functions/index.rst index da3dfe43d..d4063dd8f 100644 --- a/docs/src/piccolo/functions/index.rst +++ b/docs/src/piccolo/functions/index.rst @@ -1,14 +1,19 @@ Functions ========= +.. hint:: + This is an advanced topic - if you're new to Piccolo you can skip this for + now. + Functions can be used to modify how queries are run, and what is returned. .. toctree:: :maxdepth: 1 ./basic_usage - ./string - ./math + ./aggregate + ./array ./datetime + ./math + ./string ./type_conversion - ./aggregate diff --git a/piccolo/columns/column_types.py b/piccolo/columns/column_types.py index 79f00460b..ff10ba8cf 100644 --- a/piccolo/columns/column_types.py +++ b/piccolo/columns/column_types.py @@ -77,6 +77,7 @@ class Band(Table): ) from piccolo.columns.operators.string import Concat from piccolo.columns.reference import LazyTableReference +from piccolo.query.functions.array import ArrayMethodsMixin from piccolo.querystring import QueryString from piccolo.utils.encoding import dump_json from piccolo.utils.warnings import colored_warning @@ -2577,7 +2578,7 @@ def __repr__(self): return "list" -class Array(Column): +class Array(Column, ArrayMethodsMixin): """ Used for storing lists of data. @@ -2811,103 +2812,6 @@ def all(self, value: Any) -> Where: else: raise ValueError("Unrecognised engine type") - def cat(self, value: list[Any]) -> QueryString: - """ - Used in an ``update`` query to concatenate two arrays. - - .. code-block:: python - - >>> await Ticket.update({ - ... Ticket.seat_numbers: Ticket.seat_numbers.cat([1000]) - ... }).where(Ticket.id == 1) - - You can also use the ``+`` symbol if you prefer: - - .. code-block:: python - - >>> await Ticket.update({ - ... Ticket.seat_numbers: Ticket.seat_numbers + [1000] - ... }).where(Ticket.id == 1) - - .. note:: Postgres / CockroachDB only - - """ - from piccolo.query.functions.array import ArrayCat - - return ArrayCat(self, value=value) - - def remove(self, value: Any) -> QueryString: - """ - Used in an ``update`` query to remove an item from an array. - - .. code-block:: python - - >>> await Ticket.update({ - ... Ticket.seat_numbers: Ticket.seat_numbers.remove(1000) - ... }).where(Ticket.id == 1) - - .. note:: Postgres / CockroachDB only - - """ - from piccolo.query.functions.array import ArrayRemove - - return ArrayRemove(self, value=value) - - def prepend(self, value: Any) -> QueryString: - """ - Used in an ``update`` query to prepend an item to an array. - - .. code-block:: python - - >>> await Ticket.update({ - ... Ticket.seat_numbers: Ticket.seat_numbers.prepend(1000) - ... }).where(Ticket.id == 1) - - .. note:: Postgres / CockroachDB only - - """ - from piccolo.query.functions.array import ArrayPrepend - - return ArrayPrepend(self, value=value) - - def append(self, value: Any) -> QueryString: - """ - Used in an ``update`` query to append an item to an array. - - .. code-block:: python - - >>> await Ticket.update({ - ... Ticket.seat_numbers: Ticket.seat_numbers.append(1000) - ... }).where(Ticket.id == 1) - - .. note:: Postgres / CockroachDB only - - """ - from piccolo.query.functions.array import ArrayAppend - - return ArrayAppend(self, value=value) - - def replace(self, old_value: Any, new_value: Any) -> QueryString: - """ - Used in an ``update`` query to replace each array item - equal to the given value with a new value. - - .. code-block:: python - - >>> await Ticket.update({ - ... Ticket.seat_numbers: Ticket.seat_numbers.replace(1000, 500) - ... }).where(Ticket.id == 1) - - .. note:: Postgres / CockroachDB only - - """ - from piccolo.query.functions.array import ArrayReplace - - return ArrayReplace(self, old_value=old_value, new_value=new_value) - - def __add__(self, value: list[Any]) -> QueryString: - return self.cat(value) - ########################################################################### # Descriptors @@ -2922,3 +2826,5 @@ def __get__(self, obj, objtype=None): def __set__(self, obj, value: list[Any]): obj.__dict__[self._meta.name] = value + obj.__dict__[self._meta.name] = value + obj.__dict__[self._meta.name] = value diff --git a/piccolo/query/base.py b/piccolo/query/base.py index d45d885dc..ed1878f4d 100644 --- a/piccolo/query/base.py +++ b/piccolo/query/base.py @@ -4,7 +4,6 @@ from time import time from typing import TYPE_CHECKING, Any, Generic, Optional, Union, cast -from piccolo.columns.column_types import JSON, JSONB from piccolo.custom_types import QueryResponseType, TableInstance from piccolo.query.mixins import ColumnsDelegate from piccolo.query.operators.json import JSONQueryString @@ -62,6 +61,8 @@ async def _process_results(self, results) -> QueryResponseType: ####################################################################### + from piccolo.columns.column_types import JSON, JSONB + if output and output._output.load_json: columns_delegate: Optional[ColumnsDelegate] = getattr( self, "columns_delegate", None diff --git a/piccolo/query/functions/array.py b/piccolo/query/functions/array.py index 02b0fda0f..b610796e3 100644 --- a/piccolo/query/functions/array.py +++ b/piccolo/query/functions/array.py @@ -1,125 +1,247 @@ -from typing import Any, Union +from typing import TYPE_CHECKING, Any, Union -from piccolo.columns.base import Column -from piccolo.querystring import QueryString +from typing_extensions import TypeAlias, TypeGuard +if TYPE_CHECKING: + from piccolo.columns.base import Column + from piccolo.querystring import QueryString -class ArrayCat(QueryString): - def __init__(self, column: Union[Column, QueryString], value: list[Any]): + ArrayType: TypeAlias = Union[Column, QueryString, list[Any]] + + +def is_array_type(value) -> TypeGuard[ArrayType]: + from piccolo.columns.base import Column + from piccolo.querystring import QueryString + + return isinstance(value, QueryString) or isinstance(value, Column) + + +class ArrayMethodsMixin: + + def cat(self, array: ArrayType): """ - Concatenate two arrays. + Used in an ``update`` query to concatenate two arrays. - :param column: - Identifies the column. - :param value: - The value to concatenate. + .. code-block:: python + + >>> await Ticket.update({ + ... Ticket.seat_numbers: Ticket.seat_numbers.cat([1000]) + ... }).where(Ticket.id == 1) + + You can also use the ``+`` symbol if you prefer. To concatenate to + the end: + + .. code-block:: python + + >>> await Ticket.update({ + ... Ticket.seat_numbers: Ticket.seat_numbers + [1000] + ... }).where(Ticket.id == 1) + + To concatenate to the start: + + .. code-block:: python + + >>> await Ticket.update({ + ... Ticket.seat_numbers: [1000] + Ticket.seat_numbers + ... }).where(Ticket.id == 1) + + You can concatenate multiple arrays in one go: + + .. code-block:: python + + >>> await Ticket.update({ + ... Ticket.seat_numbers: [1000] + Ticket.seat_numbers + [2000] + ... }).where(Ticket.id == 1) + + .. note:: Postgres / CockroachDB only + + """ + assert is_array_type(self) + return ArrayCat(array_1=self, array_2=array) + + def __add__(self, array: ArrayType): + """ + QueryString will use the ``+`` operator by default for addition, but + for arrays we want to concatenate them instead. + """ + assert is_array_type(self) + return ArrayCat(array_1=self, array_2=array) + + def __radd__(self, array: ArrayType): + assert is_array_type(self) + return ArrayCat(array_1=array, array_2=self) + + def remove(self, value: Any): + """ + Used in an ``update`` query to remove an item from an array. + + .. code-block:: python + + >>> await Ticket.update({ + ... Ticket.seat_numbers: Ticket.seat_numbers.remove(1000) + ... }).where(Ticket.id == 1) + + .. note:: Postgres / CockroachDB only + + """ + assert is_array_type(self) + return ArrayRemove(array=self, value=value) + + def replace(self, old_value: Any, new_value: Any): + """ + Used in an ``update`` query to replace each array item + equal to the given value with a new value. + + .. code-block:: python + + >>> await Ticket.update({ + ... Ticket.seat_numbers: Ticket.seat_numbers.replace(1000, 500) + ... }).where(Ticket.id == 1) + + .. note:: Postgres / CockroachDB only """ - if isinstance(column, Column): - engine_type = column._meta.engine_type - if engine_type not in ("postgres", "cockroach"): - raise ValueError( - "Only Postgres and Cockroach support array concatenating." - ) + assert is_array_type(self) + return ArrayReplace( + array=self, old_value=old_value, new_value=new_value + ) + + def append(self, value: Any): + """ + Used in an ``update`` query to append an item to an array. + + .. code-block:: python + + >>> await Ticket.update({ + ... Ticket.seat_numbers: Ticket.seat_numbers.append(1000) + ... }).where(Ticket.id == 1) + + .. note:: Postgres / CockroachDB only + + """ + assert is_array_type(self) + return ArrayAppend(array=self, value=value) + + def prepend(self, value: Any): + """ + Used in an ``update`` query to prepend an item to an array. + + .. code-block:: python + + >>> await Ticket.update({ + ... Ticket.seat_numbers: Ticket.seat_numbers.prepend(1000) + ... }).where(Ticket.id == 1) + + .. note:: Postgres / CockroachDB only + + """ + assert is_array_type(self) + return ArrayPrepend(array=self, value=value) + + +class ArrayQueryString(QueryString, ArrayMethodsMixin): + pass + + +def validate_engine(value): + from piccolo.columns.base import Column - if not isinstance(value, list): - value = [value] + if isinstance(value, Column): + engine_type = value._meta.engine_type + if engine_type not in ("postgres", "cockroach"): + raise ValueError( + "Only Postgres and CockroachDB support this feature." + ) - super().__init__("array_cat({}, {})", column, value) + +class ArrayCat(ArrayQueryString): + def __init__( + self, + array_1: ArrayType, + array_2: ArrayType, + ): + """ + Concatenate two arrays. + + :param array_1: + These values will be at the start of the new array. + :param array_2: + These values will be at the end of the new array. + + """ + for value in (array_1, array_2): + validate_engine(value) + + super().__init__("array_cat({}, {})", array_1, array_2) -class ArrayAppend(QueryString): - def __init__(self, column: Union[Column, QueryString], value: Any): +class ArrayAppend(ArrayQueryString): + def __init__(self, array: ArrayType, value: Any): """ Append an element to the end of an array. - :param column: - Identifies the column. + :param array: + Identifies the array. :param value: The value to append. """ - if isinstance(column, Column): - engine_type = column._meta.engine_type - if engine_type not in ("postgres", "cockroach"): - raise ValueError( - "Only Postgres and Cockroach support array appending." - ) + validate_engine(array) - super().__init__("array_append({}, {})", column, value) + super().__init__("array_append({}, {})", array, value) -class ArrayPrepend(QueryString): - def __init__(self, column: Union[Column, QueryString], value: Any): +class ArrayPrepend(ArrayQueryString): + def __init__(self, array: ArrayType, value: Any): """ Append an element to the beginning of an array. - :param value: - The value to prepend. + :param array: + Identifies the array. :param column: - Identifies the column. + The value to prepend. """ - if isinstance(column, Column): - engine_type = column._meta.engine_type - if engine_type not in ("postgres", "cockroach"): - raise ValueError( - "Only Postgres and Cockroach support array prepending." - ) + validate_engine(value) - super().__init__("array_prepend({}, {})", value, column) + super().__init__("array_prepend({}, {})", value, array) -class ArrayReplace(QueryString): - def __init__( - self, - column: Union[Column, QueryString], - old_value: Any, - new_value: Any, - ): +class ArrayReplace(ArrayQueryString): + def __init__(self, array: ArrayType, old_value: Any, new_value: Any): """ Replace each array element equal to the given value with a new value. - :param column: - Identifies the column. + :param array: + Identifies the array. :param old_value: The old value to be replaced. :param new_value: The new value we are replacing with. """ - if isinstance(column, Column): - engine_type = column._meta.engine_type - if engine_type not in ("postgres", "cockroach"): - raise ValueError( - "Only Postgres and Cockroach support array substitution." - ) + validate_engine(array) super().__init__( - "array_replace({}, {}, {})", column, old_value, new_value + "array_replace({}, {}, {})", array, old_value, new_value ) -class ArrayRemove(QueryString): - def __init__(self, column: Union[Column, QueryString], value: Any): +class ArrayRemove(ArrayQueryString): + def __init__(self, array: ArrayType, value: Any): """ Remove all elements equal to the given value from the array (array must be one-dimensional). - :param column: - Identifies the column. + :param array: + Identifies the array. :param value: The value to remove. """ - if isinstance(column, Column): - engine_type = column._meta.engine_type - if engine_type not in ("postgres", "cockroach"): - raise ValueError( - "Only Postgres and Cockroach support array removing." - ) + validate_engine(array) - super().__init__("array_remove({}, {})", column, value) + super().__init__("array_remove({}, {})", array, value) __all__ = ( diff --git a/piccolo/query/functions/type_conversion.py b/piccolo/query/functions/type_conversion.py index cc4bd8b39..b9213f8cb 100644 --- a/piccolo/query/functions/type_conversion.py +++ b/piccolo/query/functions/type_conversion.py @@ -29,11 +29,11 @@ def __init__( """ # Make sure the identifier is a supported type. - if not isinstance(identifier, (Column, QueryString)): - raise ValueError( - "The identifier is an unsupported type - only Column and " - "QueryString instances are allowed." - ) + # if not isinstance(identifier, (Column, QueryString)): + # raise ValueError( + # "The identifier is an unsupported type - only Column and " + # "QueryString instances are allowed." + # ) ####################################################################### # Convert `as_type` to a string which can be used in the query.