Array and JSON column migration fixes (piccolo-orm#213)

* create DDL class

* some test fixes

* fix mypy warnings

* add tests for JSON and JSONB

* add tests for arrays of varchars

* reduce code repetition

Author: dantownsend

piccolo/apps/migrations/auto/serialisation.py

@@ -21,6 +21,13 @@
 ###############################################################################
 
 
+def check_equality(self, other):
+    if getattr(other, "__hash__", None) is not None:
+        return self.__hash__() == other.__hash__()
+    else:
+        return False
+
+
 @dataclass
 class SerialisedBuiltin:
     builtin: t.Any
@@ -29,7 +36,7 @@ def __hash__(self):
         return hash(self.builtin.__name__)
 
     def __eq__(self, other):
-        return self.__hash__() == other.__hash__()
+        return check_equality(self, other)
 
     def __repr__(self):
         return self.builtin.__name__
@@ -43,7 +50,7 @@ def __hash__(self):
         return self.instance.__hash__()
 
     def __eq__(self, other):
-        return self.__hash__() == other.__hash__()
+        return check_equality(self, other)
 
     def __repr__(self):
         return repr_class_instance(self.instance)
@@ -58,7 +65,7 @@ def __hash__(self):
         return self.instance.__hash__()
 
     def __eq__(self, other):
-        return self.__hash__() == other.__hash__()
+        return check_equality(self, other)
 
     def __repr__(self):
         args = ", ".join(
@@ -78,7 +85,7 @@ def __hash__(self):
         return hash(self.__repr__())
 
     def __eq__(self, other):
-        return self.__hash__() == other.__hash__()
+        return check_equality(self, other)
 
     def __repr__(self):
         return f"{self.instance.__class__.__name__}.{self.instance.name}"
@@ -94,7 +101,7 @@ def __hash__(self):
         )
 
     def __eq__(self, other):
-        return self.__hash__() == other.__hash__()
+        return check_equality(self, other)
 
     def __repr__(self):
         tablename = self.table_type._meta.tablename
@@ -126,7 +133,7 @@ def __hash__(self):
         return hash(self.__repr__())
 
     def __eq__(self, other):
-        return self.__hash__() == other.__hash__()
+        return check_equality(self, other)
 
     def __repr__(self):
         class_name = self.enum_type.__name__
@@ -142,7 +149,7 @@ def __hash__(self):
         return hash(self.callable_.__name__)
 
     def __eq__(self, other):
-        return self.__hash__() == other.__hash__()
+        return check_equality(self, other)
 
     def __repr__(self):
         return self.callable_.__name__
@@ -156,7 +163,7 @@ def __hash__(self):
         return self.instance.int
 
     def __eq__(self, other):
-        return self.__hash__() == other.__hash__()
+        return check_equality(self, other)
 
     def __repr__(self):
         return f"UUID('{str(self.instance)}')"

piccolo/columns/base.py

@@ -32,7 +32,6 @@
     NotLike,
 )
 from piccolo.columns.reference import LazyTableReference
-from piccolo.querystring import QueryString
 from piccolo.utils.warnings import colored_warning
 
 if t.TYPE_CHECKING:  # pragma: no cover
@@ -579,7 +578,16 @@ def get_sql_value(self, value: t.Any) -> t.Any:
         elif isinstance(value, list):
             # Convert to the array syntax.
             output = (
-                "'{" + ", ".join([self.get_sql_value(i) for i in value]) + "}'"
+                "'{"
+                + ", ".join(
+                    [
+                        f'"{i}"'
+                        if isinstance(i, str)
+                        else str(self.get_sql_value(i))
+                        for i in value
+                    ]
+                )
+                + "}'"
             )
         else:
             output = value
@@ -591,7 +599,7 @@ def column_type(self):
         return self.__class__.__name__.upper()
 
     @property
-    def querystring(self) -> QueryString:
+    def ddl(self) -> str:
         """
         Used when creating tables.
         """
@@ -621,16 +629,9 @@ def querystring(self) -> QueryString:
         if not self._meta.primary_key:
             default = self.get_default_value()
             sql_value = self.get_sql_value(value=default)
-            # Escape the value if it contains a pair of curly braces, otherwise
-            # an empty value will appear in the compiled querystring.
-            sql_value = (
-                sql_value.replace("{}", "{{}}")
-                if isinstance(sql_value, str)
-                else sql_value
-            )
             query += f" DEFAULT {sql_value}"
 
-        return QueryString(query)
+        return query
 
     def copy(self) -> Column:
         column: Column = copy.copy(self)
@@ -645,7 +646,7 @@ def __deepcopy__(self, memo) -> Column:
         return self.copy()
 
     def __str__(self):
-        return self.querystring.__str__()
+        return self.ddl.__str__()
 
     def __repr__(self):
         try:

piccolo/engine/base.py

@@ -50,6 +50,10 @@ async def batch(self, query: Query, batch_size: int = 100) -> Batch:
     async def run_querystring(self, querystring: QueryString, in_pool: bool):
         pass
 
+    @abstractmethod
+    async def run_ddl(self, ddl: str, in_pool: bool = True):
+        pass
+
     async def check_version(self):
         """
         Warn if the database version isn't supported.

piccolo/engine/postgres.py

@@ -6,7 +6,7 @@
 
 from piccolo.engine.base import Batch, Engine
 from piccolo.engine.exceptions import TransactionError
-from piccolo.query.base import Query
+from piccolo.query.base import DDL, Query
 from piccolo.querystring import QueryString
 from piccolo.utils.lazy_loader import LazyLoader
 from piccolo.utils.sync import run_sync
@@ -100,11 +100,15 @@ def add(self, *query: Query):
     async def _run_queries(self, connection):
         async with connection.transaction():
             for query in self.queries:
-                for querystring in query.querystrings:
-                    _query, args = querystring.compile_string(
-                        engine_type=self.engine.engine_type
-                    )
-                    await connection.execute(_query, *args)
+                if isinstance(query, Query):
+                    for querystring in query.querystrings:
+                        _query, args = querystring.compile_string(
+                            engine_type=self.engine.engine_type
+                        )
+                        await connection.execute(_query, *args)
+                elif isinstance(query, DDL):
+                    for ddl in query.ddl:
+                        await connection.execute(ddl)
 
         self.queries = []
 
@@ -391,6 +395,19 @@ async def run_querystring(
         else:
             return await self._run_in_new_connection(query, query_args)
 
+    async def run_ddl(self, ddl: str, in_pool: bool = True):
+        if self.log_queries:
+            print(ddl)
+
+        # If running inside a transaction:
+        connection = self.transaction_connection.get()
+        if connection:
+            return await connection.fetch(ddl)
+        elif in_pool and self.pool:
+            return await self._run_in_pool(ddl)
+        else:
+            return await self._run_in_new_connection(ddl)
+
     def atomic(self) -> Atomic:
         return Atomic(engine=self)
 

piccolo/engine/sqlite.py

@@ -11,7 +11,7 @@
 
 from piccolo.engine.base import Batch, Engine
 from piccolo.engine.exceptions import TransactionError
-from piccolo.query.base import Query
+from piccolo.query.base import DDL, Query
 from piccolo.querystring import QueryString
 from piccolo.utils.encoding import dump_json, load_json
 from piccolo.utils.lazy_loader import LazyLoader
@@ -250,12 +250,17 @@ async def run(self):
 
         try:
             for query in self.queries:
-                for querystring in query.querystrings:
-                    await connection.execute(
-                        *querystring.compile_string(
-                            engine_type=self.engine.engine_type
+                if isinstance(query, Query):
+                    for querystring in query.querystrings:
+                        await connection.execute(
+                            *querystring.compile_string(
+                                engine_type=self.engine.engine_type
+                            )
                         )
-                    )
+                elif isinstance(query, DDL):
+                    for ddl in query.ddl:
+                        await connection.execute(ddl)
+
         except Exception as exception:
             await connection.execute("ROLLBACK")
             await connection.close()
@@ -513,6 +518,23 @@ async def run_querystring(
             table=querystring.table,
         )
 
+    async def run_ddl(self, ddl: str, in_pool: bool = False):
+        """
+        Connection pools aren't currently supported - the argument is there
+        for consistency with other engines.
+        """
+        # If running inside a transaction:
+        connection = self.transaction_connection.get()
+        if connection:
+            return await self._run_in_existing_connection(
+                connection=connection,
+                query=ddl,
+            )
+
+        return await self._run_in_new_connection(
+            query=ddl,
+        )
+
     def atomic(self) -> Atomic:
         return Atomic(engine=self)
 

piccolo/query/base.py

@@ -317,3 +317,92 @@ def __getattr__(self, name: str):
 
     def __str__(self) -> str:
         return self.query.__str__()
+
+
+class DDL:
+
+    __slots__ = ("table",)
+
+    def __init__(self, table: t.Type[Table], **kwargs):
+        self.table = table
+
+    @property
+    def engine_type(self) -> str:
+        engine = self.table._meta.db
+        if engine:
+            return engine.engine_type
+        else:
+            raise ValueError("Engine isn't defined.")
+
+    @property
+    def sqlite_ddl(self) -> t.Sequence[str]:
+        raise NotImplementedError
+
+    @property
+    def postgres_ddl(self) -> t.Sequence[str]:
+        raise NotImplementedError
+
+    @property
+    def default_ddl(self) -> t.Sequence[str]:
+        raise NotImplementedError
+
+    @property
+    def ddl(self) -> t.Sequence[str]:
+        """
+        Calls the correct underlying method, depending on the current engine.
+        """
+        engine_type = self.engine_type
+        if engine_type == "postgres":
+            try:
+                return self.postgres_ddl
+            except NotImplementedError:
+                return self.default_ddl
+        elif engine_type == "sqlite":
+            try:
+                return self.sqlite_ddl
+            except NotImplementedError:
+                return self.default_ddl
+        else:
+            raise Exception(
+                f"No querystring found for the {engine_type} engine."
+            )
+
+    def __await__(self):
+        """
+        If the user doesn't explicity call .run(), proxy to it as a
+        convenience.
+        """
+        return self.run().__await__()
+
+    async def run(self, in_pool=True):
+        engine = self.table._meta.db
+        if not engine:
+            raise ValueError(
+                f"Table {self.table._meta.tablename} has no db defined in "
+                "_meta"
+            )
+
+        if len(self.ddl) == 1:
+            return await engine.run_ddl(self.ddl[0], in_pool=in_pool)
+        else:
+            responses = []
+            # TODO - run in a transaction
+            for ddl in self.ddl:
+                response = await engine.run_ddl(ddl, in_pool=in_pool)
+                responses.append(response)
+            return responses
+
+    def run_sync(self, timed=False, *args, **kwargs):
+        """
+        A convenience method for running the coroutine synchronously.
+        """
+        coroutine = self.run(*args, **kwargs, in_pool=False)
+
+        if timed:
+            with Timer():
+                return run_sync(coroutine)
+        else:
+            return run_sync(coroutine)
+
+    def __str__(self) -> str:
+        return self.ddl.__str__()