Merge pull request piccolo-orm#37 from piccolo-orm/duplicate_index_error

fixing bug with duplicate Table.create(if_not_exists=True) calls and indexes

Author: dantownsend

piccolo/query/methods/create.py

@@ -47,7 +47,9 @@ def querystrings(self) -> t.Sequence[QueryString]:
             if column._meta.index:
                 create_indexes.extend(
                     CreateIndex(
-                        table=self.table, columns=[column]
+                        table=self.table,
+                        columns=[column],
+                        if_not_exists=self.if_not_exists,
                     ).querystrings
                 )
 

piccolo/query/methods/create_index.py

@@ -23,9 +23,11 @@ def __init__(
         table: t.Type[Table],
         columns: t.List[t.Union[Column, str]],
         method: IndexMethod = IndexMethod.btree,
+        if_not_exists: bool = False,
     ):
         self.columns = columns
         self.method = method
+        self.if_not_exists = if_not_exists
         super().__init__(table)
 
     @property
@@ -34,6 +36,13 @@ def column_names(self) -> t.List[str]:
             i._meta.name if isinstance(i, Column) else i for i in self.columns
         ]
 
+    @property
+    def prefix(self) -> str:
+        prefix = "CREATE INDEX"
+        if self.if_not_exists:
+            prefix += " IF NOT EXISTS"
+        return prefix
+
     @property
     def postgres_querystrings(self) -> t.Sequence[QueryString]:
         column_names = self.column_names
@@ -43,8 +52,8 @@ def postgres_querystrings(self) -> t.Sequence[QueryString]:
         column_names_str = ", ".join(column_names)
         return [
             QueryString(
-                f"CREATE INDEX {index_name} ON {tablename} USING {method_name}"
-                f" ({column_names_str})"
+                f"{self.prefix} {index_name} ON {tablename} USING "
+                f"{method_name} ({column_names_str})"
             )
         ]
 
@@ -61,7 +70,7 @@ def sqlite_querystrings(self) -> t.Sequence[QueryString]:
         column_names_str = ", ".join(column_names)
         return [
             QueryString(
-                f"CREATE INDEX {index_name} ON {tablename} "
+                f"{self.prefix} {index_name} ON {tablename} "
                 f"({column_names_str})"
             )
         ]

tests/table/test_create.py

@@ -29,3 +29,23 @@ def test_create_table_with_indexes(self):
         index_names = BandMember.indexes().run_sync()
         index_name = BandMember._get_index_name(["name"])
         self.assertTrue(index_name in index_names)
+
+    def test_create_if_not_exists_with_indexes(self):
+        """
+        Make sure that if the same table is created again, with the
+        `if_not_exists` flag, then no errors are raised for duplicate indexes
+        (i.e. the indexes should also be created with IF NOT EXISTS).
+        """
+        query = BandMember.create_table(if_not_exists=True)
+
+        # Shouldn't raise any errors:
+        query.run_sync()
+
+        self.assertTrue(
+            query.querystrings[0]
+            .__str__()
+            .startswith("CREATE TABLE IF NOT EXISTS"),
+            query.querystrings[1]
+            .__str__()
+            .startswith("CREATE INDEX IF NOT EXISTS"),
+        )