migrations are using the index_method value

Author: dantownsend

piccolo/apps/migrations/auto/migration_manager.py

@@ -404,12 +404,30 @@ 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
+                        ).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], **kwargs).run()
                     else:
                         await _Table.drop_index([column]).run()
 

piccolo/columns/base.py

@@ -27,12 +27,12 @@
 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.columns.column_types import ForeignKey
-    from piccolo.query.methods.create_index import IndexMethod
     from piccolo.table import Table
 
 
@@ -120,7 +120,8 @@ class ColumnMeta:
     primary: bool = False
     key: bool = False
     unique: bool = False
-    index: t.Union[IndexMethod, bool] = False
+    index: bool = False
+    index_method: IndexMethod = IndexMethod.btree
     required: bool = False
 
     # Used for representing the table in migrations and the playground.
@@ -233,9 +234,10 @@ class Column(Selectable):
 
     :param index:
         Whether an index is created for the column, which can improve
-        the speed of selects, but can slow down inserts. If set to True, the
-        default index type is created (b-tree). If False, no index is created.
-        If an IndexMethod value is provided, that type of index is created.
+        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
@@ -252,7 +254,8 @@ def __init__(
         primary: bool = False,
         key: bool = False,
         unique: bool = False,
-        index: t.Union[IndexMethod, bool] = False,
+        index: bool = False,
+        index_method: IndexMethod = IndexMethod.btree,
         required: bool = False,
         **kwargs,
     ) -> None:
@@ -266,6 +269,7 @@ def __init__(
                 "key": key,
                 "unique": unique,
                 "index": index,
+                "index_method": index_method,
             }
         )
 
@@ -281,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

@@ -3,7 +3,7 @@
 
 from piccolo.query.base import Query
 from piccolo.querystring import QueryString
-from piccolo.query.methods.create_index import CreateIndex, IndexMethod
+from piccolo.query.methods.create_index import CreateIndex
 
 if t.TYPE_CHECKING:  # pragma: no cover
     from piccolo.table import Table
@@ -44,22 +44,14 @@ def querystrings(self) -> t.Sequence[QueryString]:
 
         create_indexes: t.List[QueryString] = []
         for column in columns:
-            if column._meta.index is False:
-                continue
-            elif column._meta.index is True:
-                method = IndexMethod.btree
-            elif isinstance(column._meta.index, IndexMethod):
-                method = column._meta.index
-            else:
-                raise ValueError("Unrecognised value for column index")
-
-            create_indexes.extend(
-                CreateIndex(
-                    table=self.table,
-                    columns=[column],
-                    method=method,
-                    if_not_exists=self.if_not_exists,
-                ).querystrings
-            )
+            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
+                )
 
         return [create_table] + create_indexes

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
 

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
             ),
         )