Reverse lookup

docs/src/piccolo/schema/index.rst

@@ -9,5 +9,6 @@ The schema is how you define your database tables, columns and relationships.
     ./defining
     ./column_types
     ./m2m
+    ./reverse_lookup
     ./one_to_one
     ./advanced

docs/src/piccolo/schema/reverse_lookup.rst

@@ -0,0 +1,105 @@
+.. currentmodule:: piccolo.columns.reverse_lookup
+
+##############
+Reverse Lookup
+##############
+
+For example, we might have our ``Manager`` table, and we want to 
+get all the bands associated with the same manager. 
+For this we can use reverse foreign key lookup.
+
+We create it in Piccolo like this:
+
+.. code-block:: python
+
+    from piccolo.columns.column_types import (
+        ForeignKey,
+        LazyTableReference,
+        Varchar
+    )
+    from piccolo.columns.reverse_lookup import ReverseLookup
+    from piccolo.table import Table
+
+
+    class Manager(Table):
+        name = Varchar()
+        bands = ReverseLookup(
+            LazyTableReference(
+                "Band",
+                module_path=__name__,
+            ),
+            reverse_fk="manager",
+        )
+
+
+    class Band(Table):
+        name = Varchar()
+        manager = ForeignKey(Manager)
+
+-------------------------------------------------------------------------------
+
+Select queries
+==============
+
+If we want to select each manager, along with a list of associated band names,
+we can do this:
+
+.. code-block:: python
+
+    >>> await Manager.select(Manager.name, Manager.bands(Band.name, as_list=True))
+    [
+        {'name': 'John', 'bands': ['C-Sharps']},
+        {'name': 'Guido', 'bands': ['Pythonistas', 'Rustaceans']},
+    ]
+
+You can request whichever column you like from the reverse lookup:
+
+.. code-block:: python
+
+    >>> await Manager.select(Manager.name, Manager.bands(Band.id, as_list=True))
+    [
+        {'name': 'John', 'bands': [3]}, 
+        {'name': 'Guido', 'bands': [1, 2]},
+    ]
+
+You can also request multiple columns from the reverse lookup:
+
+.. code-block:: python
+
+    >>> await Manager.select(Manager.name, Manager.bands(Band.id, Band.name))
+    [
+        {
+            'name': 'John',
+            'bands': [
+                {'id': 3, 'name': 'C-Sharps'},
+            ]
+        },
+        {   
+            'name': 'Guido', 
+            'bands': [
+                {'id': 1, 'name': 'Pythonistas'}, 
+                {'id': 2, 'name': 'Rustaceans'},
+            ]
+        }
+    ]
+
+If you omit the columns argument, then all of the columns are returned.
+
+.. code-block:: python
+
+    >>> await Manager.select(Manager.name, Manager.bands())
+    [
+        {
+            'name': 'John',
+            'bands': [
+                {'id': 3, 'name': 'C-Sharps'},
+            ]
+        },
+        {   
+            'name': 'Guido', 
+            'bands': [
+                {'id': 1, 'name': 'Pythonistas'}, 
+                {'id': 2, 'name': 'Rustaceans'},
+            ]
+        }
+    ]

piccolo/columns/m2m.py

@@ -15,7 +15,7 @@
 from piccolo.utils.list import flatten
 from piccolo.utils.sync import run_sync
 
-if t.TYPE_CHECKING:
+if t.TYPE_CHECKING:  # pragma: no cover
     from piccolo.table import Table
 
 

piccolo/columns/reverse_lookup.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+import inspect
+import typing as t
+from dataclasses import dataclass
+
+from piccolo.columns.base import QueryString, Selectable
+from piccolo.columns.column_types import (
+    JSON,
+    JSONB,
+    Column,
+    LazyTableReference,
+)
+
+if t.TYPE_CHECKING:  # pragma: no cover
+    from piccolo.table import Table
+
+
+class ReverseLookupSelect(Selectable):
+    """
+    This is a subquery used within a select to fetch reverse lookup data.
+    """
+
+    def __init__(
+        self,
+        *columns: Column,
+        reverse_lookup: ReverseLookup,
+        as_list: bool = False,
+        load_json: bool = False,
+    ):
+        """
+        :param columns:
+            Which columns to include from the related table.
+        :param as_list:
+            If a single column is provided, and ``as_list`` is ``True`` a
+            flattened list will be returned, rather than a list of objects.
+        :param load_json:
+            If ``True``, any JSON strings are loaded as Python objects.
+
+        """
+        self.as_list = as_list
+        self.columns = columns
+        self.reverse_lookup = reverse_lookup
+        self.load_json = load_json
+
+        safe_types = [int, str]
+
+        # If the columns can be serialised / deserialise as JSON, then we
+        # can fetch the data all in one go.
+        self.serialisation_safe = all(
+            (column.__class__.value_type in safe_types)
+            and (type(column) not in (JSON, JSONB))
+            for column in columns
+        )
+
+    def get_select_string(
+        self, engine_type: str, with_alias=True
+    ) -> QueryString:
+        reverse_lookup_name = self.reverse_lookup._meta.name
+
+        table1 = self.reverse_lookup._meta.table
+        table1_pk = table1._meta.primary_key._meta.name
+        table1_name = table1._meta.tablename
+
+        table2 = self.reverse_lookup._meta.resolved_reverse_joining_table
+        table2_name = table2._meta.tablename
+        table2_pk = table2._meta.primary_key._meta.name
+        table2_fk = self.reverse_lookup._meta.reverse_fk
+
+        reverse_select = f"""
+            "{table2_name}"
+            WHERE "{table2_name}"."{table2_fk}"
+                = "{table1_name}"."{table1_pk}"
+        """
+
+        if engine_type in ("postgres", "cockroach"):
+            if self.as_list:
+                column_name = self.columns[0]._meta.db_column_name
+                return QueryString(
+                    f"""
+                    ARRAY(
+                        SELECT
+                            "{table2_name}"."{column_name}"
+                        FROM {reverse_select}
+                    ) AS "{reverse_lookup_name}"
+                """
+                )
+            elif not self.serialisation_safe:
+                column_name = table2_pk
+                return QueryString(
+                    f"""
+                    ARRAY(
+                        SELECT
+                            "{table2_name}"."{column_name}"
+                        FROM {reverse_select}
+                    ) AS "{reverse_lookup_name}"
+                """
+                )
+            else:
+                if len(self.columns) > 0:
+                    column_names = ", ".join(
+                        f'"{table2_name}"."{column._meta.db_column_name}"'  # noqa: E501
+                        for column in self.columns
+                    )
+                else:
+                    column_names = ", ".join(
+                        f'"{table2_name}"."{column._meta.db_column_name}"'  # noqa: E501
+                        for column in table2._meta.columns
+                    )
+                return QueryString(
+                    f"""
+                    (
+                        SELECT JSON_AGG("{table2_name}s")
+                        FROM (
+                            SELECT {column_names} FROM {reverse_select}
+                        ) AS "{table2_name}s"
+                    ) AS "{reverse_lookup_name}"
+                """
+                )
+        elif engine_type == "sqlite":
+            if len(self.columns) > 1 or not self.serialisation_safe:
+                column_name = table2_pk
+            else:
+                try:
+                    column_name = self.columns[0]._meta.db_column_name
+                except IndexError:
+                    column_name = table2_pk
+
+            return QueryString(
+                f"""
+                (
+                    SELECT group_concat(
+                        "{table2_name}"."{column_name}"
+                    )
+                    FROM {reverse_select}
+                )
+                AS "{reverse_lookup_name} [M2M]"
+            """
+            )
+        else:
+            raise ValueError(f"{engine_type} is an unrecognised engine type")
+
+
+@dataclass
+class ReverseLookupMeta:
+    reverse_joining_table: t.Union[t.Type[Table], LazyTableReference]
+    reverse_fk: str
+
+    # Set by the Table Metaclass:
+    _name: t.Optional[str] = None
+    _table: t.Optional[t.Type[Table]] = None
+
+    @property
+    def name(self) -> str:
+        if not self._name:
+            raise ValueError(
+                "`_name` isn't defined - the Table Metaclass should set it."
+            )
+        return self._name
+
+    @property
+    def table(self) -> t.Type[Table]:
+        if not self._table:
+            raise ValueError(
+                "`_table` isn't defined - the Table Metaclass should set it."
+            )
+        return self._table
+
+    @property
+    def resolved_reverse_joining_table(self) -> t.Type[Table]:
+        """
+        Evaluates the ``reverse_joining_table`` attribute if it's a
+        ``LazyTableReference``, raising a ``ValueError`` if it fails,
+        otherwise returns a ``Table`` subclass.
+        """
+        from piccolo.table import Table
+
+        if isinstance(self.reverse_joining_table, LazyTableReference):
+            return self.reverse_joining_table.resolve()
+        elif inspect.isclass(self.reverse_joining_table) and issubclass(
+            self.reverse_joining_table, Table
+        ):
+            return self.reverse_joining_table
+        else:
+            raise ValueError(
+                "The reverse_joining_table attribute is neither a Table"
+                " subclass or a LazyTableReference instance."
+            )
+
+
+class ReverseLookup:
+    def __init__(
+        self,
+        reverse_joining_table: t.Union[t.Type[Table], LazyTableReference],
+        reverse_fk: str,
+    ):
+        """
+        :param reverse_joining_table:
+            A ``Table`` for reverse lookup.
+        :param reverse_fk:
+            The ForeignKey to be used for the reverse lookup.
+        """
+        self._meta = ReverseLookupMeta(
+            reverse_joining_table=reverse_joining_table,
+            reverse_fk=reverse_fk,
+        )
+
+    def __call__(
+        self,
+        *columns: Column,
+        as_list: bool = False,
+        load_json: bool = False,
+    ) -> ReverseLookupSelect:
+        """
+        :param columns:
+            Which columns to include from the related table. If none are
+            specified, then all of the columns are returned.
+        :param as_list:
+            If a single column is provided, and ``as_list`` is ``True`` a
+            flattened list will be returned, rather than a list of objects.
+        :param load_json:
+            If ``True``, any JSON strings are loaded as Python objects.
+
+        """
+
+        if as_list and len(columns) != 1:
+            raise ValueError(
+                "`as_list` is only valid with a single column argument"
+            )
+
+        return ReverseLookupSelect(
+            *columns,
+            reverse_lookup=self,
+            as_list=as_list,
+            load_json=load_json,
+        )