From 1b491ad9f4057bcd4d920e7f3c78077b16329b09 Mon Sep 17 00:00:00 2001
From: Daniel Townsend <dan@dantownsend.co.uk>
Date: Sun, 21 Jul 2024 08:24:52 +0100
Subject: [PATCH] Refactor `_process_results` to support new DB engines

---
 piccolo/columns/base.py         |  6 +++---
 piccolo/engine/base.py          |  7 +++++++
 piccolo/engine/postgres.py      |  7 +++++++
 piccolo/query/base.py           | 19 +++++--------------
 piccolo/query/methods/select.py |  2 +-
 piccolo/querystring.py          |  2 +-
 6 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/piccolo/columns/base.py b/piccolo/columns/base.py
index 9d3e2b1cc..4725b78ad 100644
--- a/piccolo/columns/base.py
+++ b/piccolo/columns/base.py
@@ -250,11 +250,11 @@ def get_default_alias(self):
 
         if self.call_chain:
             column_name = (
-                "$".join(
+                ".".join(
                     t.cast(str, i._meta.db_column_name)
                     for i in self.call_chain
                 )
-                + f"${column_name}"
+                + f".{column_name}"
             )
 
         return column_name
@@ -291,7 +291,7 @@ def get_full_name(
                 'band$manager.name'
 
                 >>> Band.manager.name._meta.get_full_name(with_alias=True)
-                'band$manager.name AS "manager$name"'
+                'band$manager.name AS "manager.name"'
 
         :param include_quotes:
             If you're using the name in a SQL query, each component needs to be
diff --git a/piccolo/engine/base.py b/piccolo/engine/base.py
index bf59426ad..0181a29b5 100644
--- a/piccolo/engine/base.py
+++ b/piccolo/engine/base.py
@@ -132,6 +132,13 @@ async def run_querystring(
     ):
         pass
 
+    def transform_response_to_dicts(self, results) -> t.List[t.Dict]:
+        """
+        If the database adapter returns something other than a list of
+        dictionaries, it should perform the transformation here.
+        """
+        return results
+
     @abstractmethod
     async def run_ddl(self, ddl: str, in_pool: bool = True):
         pass
diff --git a/piccolo/engine/postgres.py b/piccolo/engine/postgres.py
index 970623535..117bf9a53 100644
--- a/piccolo/engine/postgres.py
+++ b/piccolo/engine/postgres.py
@@ -579,6 +579,13 @@ async def run_ddl(self, ddl: str, in_pool: bool = True):
 
         return response
 
+    def transform_response_to_dicts(self, results) -> t.List[t.Dict]:
+        """
+        asyncpg returns a special Record object, so we need to convert it to
+        a dict.
+        """
+        return [dict(i) for i in results]
+
     def atomic(self) -> Atomic:
         return Atomic(engine=self)
 
diff --git a/piccolo/query/base.py b/piccolo/query/base.py
index c169fde0e..444736f03 100644
--- a/piccolo/query/base.py
+++ b/piccolo/query/base.py
@@ -45,20 +45,11 @@ def engine_type(self) -> str:
             raise ValueError("Engine isn't defined.")
 
     async def _process_results(self, results) -> QueryResponseType:
-        if results:
-            keys = results[0].keys()
-            keys = [i.replace("$", ".") for i in keys]
-            if self.engine_type in ("postgres", "cockroach"):
-                # asyncpg returns a special Record object. We can pass it
-                # directly into zip without calling `values` on it. This can
-                # save us hundreds of microseconds, depending on the number of
-                # results.
-                raw = [dict(zip(keys, i)) for i in results]
-            else:
-                # SQLite returns a list of dictionaries.
-                raw = [dict(zip(keys, i.values())) for i in results]
-        else:
-            raw = []
+        raw = (
+            self.table._meta.db.transform_response_to_dicts(results)
+            if results
+            else []
+        )
 
         if hasattr(self, "_raw_response_callback"):
             self._raw_response_callback(raw)
diff --git a/piccolo/query/methods/select.py b/piccolo/query/methods/select.py
index 4c6390fae..55267b703 100644
--- a/piccolo/query/methods/select.py
+++ b/piccolo/query/methods/select.py
@@ -574,7 +574,7 @@ def default_querystrings(self) -> t.Sequence[QueryString]:
         query += "{}"
         args.append(distinct.querystring)
 
-        columns_str = ", ".join("{}" for i in select_strings)
+        columns_str = ", ".join("{}" for _ in select_strings)
         query += f" {columns_str} FROM {self.table._meta.get_formatted_tablename()}"  # noqa: E501
         args.extend(select_strings)
 
diff --git a/piccolo/querystring.py b/piccolo/querystring.py
index 1f7282a7d..22f5f215a 100644
--- a/piccolo/querystring.py
+++ b/piccolo/querystring.py
@@ -246,7 +246,7 @@ def get_select_string(
         self, engine_type: str, with_alias: bool = True
     ) -> QueryString:
         if with_alias and self._alias:
-            return QueryString("{} AS " + self._alias, self)
+            return QueryString("{} AS " + f'"{self._alias}"', self)
         else:
             return self