Merge pull request piccolo-orm#109 from piccolo-orm/frozen_queries

Frozen queries

Author: dantownsend

docs/src/piccolo/query_clauses/batch.rst

@@ -33,5 +33,5 @@ There's currently no synchronous version. However, it's easy enough to achieve:
             async for _batch in batch:
                 print(_batch)
 
-    import asyncio
-    asyncio.run(get_batch())
+    from piccolo.utils.sync import run_sync
+    run_sync(get_batch())

docs/src/piccolo/query_clauses/freeze.rst

@@ -0,0 +1,10 @@
+.. _freeze:
+
+freeze
+======
+
+You can use the ``freeze`` clause with any query type.
+
+.. currentmodule:: piccolo.query.base
+
+.. automethod:: Query.freeze

docs/src/piccolo/query_clauses/index.rst

@@ -16,3 +16,4 @@ by modifying the return values.
     ./order_by
     ./where
     ./batch
+    ./freeze

piccolo/apps/user/tables.py

@@ -41,9 +41,7 @@ class BaseUser(Table, tablename="piccolo_user"):
     )
 
     def __init__(self, **kwargs):
-        """
-        Generating passwords upfront is expensive, so might need reworking.
-        """
+        # Generating passwords upfront is expensive, so might need reworking.
         password = kwargs.get("password", None)
         if password:
             kwargs["password"] = self.__class__.hash_password(password)

piccolo/query/base.py

@@ -22,10 +22,15 @@ def __exit__(self, exception_type, exception, traceback):
 
 class Query:
 
-    __slots__ = ("table",)
+    __slots__ = ("table", "_frozen_querystrings")
 
-    def __init__(self, table: t.Type[Table]):
+    def __init__(
+        self,
+        table: t.Type[Table],
+        frozen_querystrings: t.Optional[t.Sequence[QueryString]] = None,
+    ):
         self.table = table
+        self._frozen_querystrings = frozen_querystrings
 
     @property
     def engine_type(self) -> str:
@@ -156,6 +161,9 @@ def querystrings(self) -> t.Sequence[QueryString]:
         """
         Calls the correct underlying method, depending on the current engine.
         """
+        if self._frozen_querystrings is not None:
+            return self._frozen_querystrings
+
         engine_type = self.engine_type
         if engine_type == "postgres":
             try:
@@ -174,5 +182,86 @@ def querystrings(self) -> t.Sequence[QueryString]:
 
     ###########################################################################
 
+    def freeze(self) -> FrozenQuery:
+        """
+        This is a performance optimisation when the same query is run
+        repeatedly. For example:
+
+        .. code-block:: python
+
+            TOP_BANDS = Band.select(
+                Band.name
+            ).order_by(
+                Band.popularity,
+                ascending=False
+            ).limit(
+                10
+            ).output(
+                as_json=True
+            ).freeze()
+
+            # In the corresponding view/endpoint of whichever web framework
+            # you're using:
+            async def top_bands(self, request):
+                return await TOP_BANDS.run()
+
+        It means that Piccolo doesn't have to work as hard each time the query
+        is run to generate the corresponding SQL - some of it is cached. If the
+        query is defined within the view/endpoint, it has to generate the SQL
+        from scratch each time.
+
+        Once a query is frozen, you can't apply any more clauses to it
+        (``where``, ``limit``, ``output`` etc).
+
+        Even though ``freeze`` helps with performance, there are limits to
+        how much it can help, as most of the time is still spent waiting for a
+        response from the database. However, for high throughput apps and data
+        science scripts, it's a worthwhile optimisation.
+
+        """
+        querystrings = self.querystrings
+        for querystring in querystrings:
+            querystring.freeze(engine_type=self.engine_type)
+
+        # Copy the query, so we don't store any references to the original.
+        query = self.__class__(
+            table=self.table, frozen_querystrings=self.querystrings
+        )
+
+        if hasattr(self, "limit_delegate"):
+            # Needed for `response_handler`
+            query.limit_delegate = self.limit_delegate.copy()  # type: ignore
+
+        if hasattr(self, "output_delegate"):
+            # Needed for `_process_results`
+            query.output_delegate = self.output_delegate.copy()  # type: ignore
+
+        return FrozenQuery(query=query)
+
+    ###########################################################################
+
     def __str__(self) -> str:
         return "; ".join([i.__str__() for i in self.querystrings])
+
+
+class FrozenQuery:
+    def __init__(self, query: Query):
+        self.query = query
+
+    async def run(self, *args, **kwargs):
+        return await self.query.run(*args, **kwargs)
+
+    def run_sync(self, *args, **kwargs):
+        return self.query.run_sync(*args, **kwargs)
+
+    def __getattr__(self, name: str):
+        if hasattr(self.query, name):
+            raise AttributeError(
+                f"This query is frozen - {name} is only available on "
+                "unfrozen queries."
+            )
+        else:
+            raise AttributeError("Unrecognised attribute name.")
+
+    def __str__(self) -> str:
+        return self.query.__str__()

piccolo/query/methods/alter.py

@@ -293,8 +293,8 @@ class Alter(Query):
         "_set_unique",
     )
 
-    def __init__(self, table: t.Type[Table]):
-        super().__init__(table)
+    def __init__(self, table: t.Type[Table], **kwargs):
+        super().__init__(table, **kwargs)
         self._add_foreign_key_constraint: t.List[AddForeignKeyConstraint] = []
         self._add: t.List[AddColumn] = []
         self._drop_contraint: t.List[DropConstraint] = []
@@ -506,7 +506,7 @@ def set_digits(
         return self
 
     @property
-    def querystrings(self) -> t.Sequence[QueryString]:
+    def default_querystrings(self) -> t.Sequence[QueryString]:
         if self._drop_table is not None:
             return [self._drop_table.querystring]
 

piccolo/query/methods/count.py

@@ -15,8 +15,8 @@
 class Count(Query):
     __slots__ = ("where_delegate",)
 
-    def __init__(self, table: t.Type[Table]):
-        super().__init__(table)
+    def __init__(self, table: t.Type[Table], **kwargs):
+        super().__init__(table, **kwargs)
         self.where_delegate = WhereDelegate()
 
     def where(self, where: Combinable) -> Count:
@@ -27,7 +27,7 @@ async def response_handler(self, response) -> bool:
         return response[0]["count"]
 
     @property
-    def querystrings(self) -> t.Sequence[QueryString]:
+    def default_querystrings(self) -> t.Sequence[QueryString]:
         select = Select(self.table)
         select.where_delegate._where = self.where_delegate._where
         return [

piccolo/query/methods/create.py

@@ -21,13 +21,14 @@ def __init__(
         table: t.Type[Table],
         if_not_exists: bool = False,
         only_default_columns: bool = False,
+        **kwargs,
     ):
-        super().__init__(table)
+        super().__init__(table, **kwargs)
         self.if_not_exists = if_not_exists
         self.only_default_columns = only_default_columns
 
     @property
-    def querystrings(self) -> t.Sequence[QueryString]:
+    def default_querystrings(self) -> t.Sequence[QueryString]:
         prefix = "CREATE TABLE"
         if self.if_not_exists:
             prefix += " IF NOT EXISTS"

piccolo/query/methods/create_index.py

@@ -17,11 +17,12 @@ def __init__(
         columns: t.List[t.Union[Column, str]],
         method: IndexMethod = IndexMethod.btree,
         if_not_exists: bool = False,
+        **kwargs,
     ):
         self.columns = columns
         self.method = method
         self.if_not_exists = if_not_exists
-        super().__init__(table)
+        super().__init__(table, **kwargs)
 
     @property
     def column_names(self) -> t.List[str]:

piccolo/query/methods/delete.py

@@ -18,8 +18,8 @@ class Delete(Query):
 
     __slots__ = ("force", "where_delegate")
 
-    def __init__(self, table: t.Type[Table], force: bool = False):
-        super().__init__(table)
+    def __init__(self, table: t.Type[Table], force: bool = False, **kwargs):
+        super().__init__(table, **kwargs)
         self.force = force
         self.where_delegate = WhereDelegate()
 
@@ -40,7 +40,7 @@ def _validate(self):
             )
 
     @property
-    def querystrings(self) -> t.Sequence[QueryString]:
+    def default_querystrings(self) -> t.Sequence[QueryString]:
         query = f"DELETE FROM {self.table._meta.tablename}"
         if self.where_delegate._where:
             query += " WHERE {}"