Merge pull request piccolo-orm#72 from piccolo-orm/expose_index_options

Expose index options

Author: dantownsend

piccolo/apps/migrations/auto/migration_manager.py

@@ -404,12 +404,32 @@ async def _run_alter_columns(self, backwards=False):
                     ).run()
 
                 index = params.get("index")
-                if index is not None:
+                index_method = params.get("index_method")
+                if index is None:
+                    if index_method is not None:
+                        # If the index value hasn't changed, but the
+                        # index_method value has, this indicates we need
+                        # to change the index type.
+                        column = Column()
+                        column._meta._table = _Table
+                        column._meta._name = column_name
+                        await _Table.drop_index([column]).run()
+                        await _Table.create_index(
+                            [column], method=index_method, if_not_exists=True
+                        ).run()
+                else:
+                    # If the index value has changed, then we are either
+                    # dropping, or creating an index.
                     column = Column()
                     column._meta._table = _Table
                     column._meta._name = column_name
-                    if index:
-                        await _Table.create_index([column]).run()
+                    if index is True:
+                        kwargs = (
+                            {"method": index_method} if index_method else {}
+                        )
+                        await _Table.create_index(
+                            [column], if_not_exists=True, **kwargs
+                        ).run()
                     else:
                         await _Table.drop_index([column]).run()
 

piccolo/columns/base.py

@@ -27,12 +27,13 @@
 from piccolo.columns.combination import Where
 from piccolo.columns.defaults.base import Default
 from piccolo.columns.reference import LazyTableReference
+from piccolo.columns.indexes import IndexMethod
 from piccolo.querystring import QueryString
 from piccolo.utils.warnings import colored_warning
 
 if t.TYPE_CHECKING:  # pragma: no cover
-    from piccolo.table import Table
     from piccolo.columns.column_types import ForeignKey
+    from piccolo.table import Table
 
 
 class OnDelete(str, Enum):
@@ -120,6 +121,7 @@ class ColumnMeta:
     key: bool = False
     unique: bool = False
     index: bool = False
+    index_method: IndexMethod = IndexMethod.btree
     required: bool = False
 
     # Used for representing the table in migrations and the playground.
@@ -231,9 +233,12 @@ class Column(Selectable):
         If set, a unique contraint will be added to the column.
 
     :param index:
-        If set, the an index is created for the column, which can improve
+        Whether an index is created for the column, which can improve
         the speed of selects, but can slow down inserts.
 
+    :param index_method:
+        If index is set to True, this specifies what type of index is created.
+
     :param required:
         This isn't used by the database - it's to indicate to other tools that
         the user must provide this value. Example uses are in serialisers for
@@ -250,6 +255,7 @@ def __init__(
         key: bool = False,
         unique: bool = False,
         index: bool = False,
+        index_method: IndexMethod = IndexMethod.btree,
         required: bool = False,
         **kwargs,
     ) -> None:
@@ -263,6 +269,7 @@ def __init__(
                 "key": key,
                 "unique": unique,
                 "index": index,
+                "index_method": index_method,
             }
         )
 
@@ -278,6 +285,7 @@ def __init__(
             key=key,
             unique=unique,
             index=index,
+            index_method=index_method,
             params=kwargs,
             required=required,
         )

piccolo/columns/indexes.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class IndexMethod(str, Enum):
+    btree = "btree"
+    hash = "hash"
+    gist = "gist"
+    gin = "gin"
+
+    def __str__(self):
+        return f"{self.__class__.__name__}.{self.name}"
+
+    def __repr__(self):
+        return self.__str__()

piccolo/query/methods/create.py

@@ -44,11 +44,12 @@ def querystrings(self) -> t.Sequence[QueryString]:
 
         create_indexes: t.List[QueryString] = []
         for column in columns:
-            if column._meta.index:
+            if column._meta.index is True:
                 create_indexes.extend(
                     CreateIndex(
                         table=self.table,
                         columns=[column],
+                        method=column._meta.index_method,
                         if_not_exists=self.if_not_exists,
                     ).querystrings
                 )

piccolo/query/methods/create_index.py

@@ -1,22 +1,15 @@
 from __future__ import annotations
-from enum import Enum
 import typing as t
 
 from piccolo.columns import Column
+from piccolo.columns.indexes import IndexMethod
 from piccolo.query.base import Query
 from piccolo.querystring import QueryString
 
 if t.TYPE_CHECKING:  # pragma: no cover
     from piccolo.table import Table
 
 
-class IndexMethod(str, Enum):
-    btree = "btree"
-    hash = "hash"
-    gist = "gist"
-    gin = "gin"
-
-
 class CreateIndex(Query):
     def __init__(
         self,

piccolo/table.py

@@ -16,6 +16,7 @@
     LAZY_COLUMN_REFERENCES,
 )
 from piccolo.columns.defaults.base import Default
+from piccolo.columns.indexes import IndexMethod
 from piccolo.query import (
     Alter,
     Count,
@@ -31,7 +32,7 @@
     Update,
 )
 from piccolo.query.methods.indexes import Indexes
-from piccolo.query.methods.create_index import CreateIndex, IndexMethod
+from piccolo.query.methods.create_index import CreateIndex
 from piccolo.querystring import QueryString, Unquoted
 from piccolo.utils import _camel_to_snake
 
@@ -622,14 +623,20 @@ def create_index(
         cls,
         columns: t.List[t.Union[Column, str]],
         method: IndexMethod = IndexMethod.btree,
+        if_not_exists: bool = False,
     ) -> CreateIndex:
         """
         Create a table index. If multiple columns are specified, this refers
         to a multicolumn index, rather than multiple single column indexes.
 
         await Band.create_index([Band.name]).run()
         """
-        return CreateIndex(table=cls, columns=columns, method=method)
+        return CreateIndex(
+            table=cls,
+            columns=columns,
+            method=method,
+            if_not_exists=if_not_exists,
+        )
 
     @classmethod
     def drop_index(

tests/apps/migrations/auto/test_schema_differ.py

@@ -37,7 +37,7 @@ def test_add_table(self):
         self.assertTrue(len(new_table_columns.statements) == 1)
         self.assertEqual(
             new_table_columns.statements[0],
-            "manager.add_column(table_class_name='Band', tablename='band', column_name='name', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary': False, 'key': False, 'unique': False, 'index': False})",  # noqa
+            "manager.add_column(table_class_name='Band', tablename='band', column_name='name', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary': False, 'key': False, 'unique': False, 'index': False, 'index_method': IndexMethod.btree})",  # noqa
         )
 
     def test_drop_table(self):
@@ -119,7 +119,7 @@ def test_add_column(self):
         self.assertTrue(len(schema_differ.add_columns.statements) == 1)
         self.assertEqual(
             schema_differ.add_columns.statements[0],
-            "manager.add_column(table_class_name='Band', tablename='band', column_name='genre', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary': False, 'key': False, 'unique': False, 'index': False})",  # noqa
+            "manager.add_column(table_class_name='Band', tablename='band', column_name='genre', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary': False, 'key': False, 'unique': False, 'index': False, 'index_method': IndexMethod.btree})",  # noqa
         )
 
     def test_drop_column(self):

tests/table/test_str.py

@@ -9,8 +9,8 @@ def test_str(self):
             Manager._table_str(),
             (
                 "class Manager(Table, tablename='manager'):\n"
-                "    id = PrimaryKey(null=False, primary=True, key=True, unique=False, index=False)\n"  # noqa
-                "    name = Varchar(length=50, default='', null=False, primary=False, key=False, unique=False, index=False)\n"  # noqa
+                "    id = PrimaryKey(null=False, primary=True, key=True, unique=False, index=False, index_method=IndexMethod.btree)\n"  # noqa
+                "    name = Varchar(length=50, default='', null=False, primary=False, key=False, unique=False, index=False, index_method=IndexMethod.btree)\n"  # noqa
             ),
         )