added create_index and drop_index methods to Table

Author: dantownsend

dev-requirements.txt

@@ -1,5 +1,5 @@
 ipdb==0.12.2
 ipython==7.8.0
 twine==3.1.1
-mypy==0.750
+mypy==0.782
 pip-upgrader==1.4.15

piccolo/columns/base.py

@@ -121,6 +121,13 @@ def engine_type(self) -> str:
         else:
             raise ValueError("The table has no engine defined.")
 
+    @property
+    def index_name(self) -> str:
+        """
+        If an index is added for this column, this is the name which is used.
+        """
+        return f"{self.table._meta.tablename}_{self.name}"
+
     def get_full_name(self, just_alias=False) -> str:
         """
         Returns the full column name, taking into account joins.

piccolo/query/__init__.py

@@ -1,4 +1,4 @@
-from .base import Query
+from .base import Query  # noqa
 from .methods import (
     Alter,
     Select,
@@ -11,4 +11,6 @@
     TableExists,
     Exists,
     Count,
+    CreateIndex,
+    DropIndex,
 )

piccolo/query/methods/__init__.py

@@ -1,11 +1,13 @@
-from .alter import Alter
-from .count import Count
-from .create import Create
-from .delete import Delete
-from .exists import Exists
-from .insert import Insert
-from .objects import Objects
-from .raw import Raw
-from .select import Select
-from .table_exists import TableExists
-from .update import Update
+from .alter import Alter  # noqa
+from .count import Count  # noqa
+from .create import Create  # noqa
+from .create_index import CreateIndex  # noqa
+from .delete import Delete  # noqa
+from .drop_index import DropIndex  # noqa
+from .exists import Exists  # noqa
+from .insert import Insert  # noqa
+from .objects import Objects  # noqa
+from .raw import Raw  # noqa
+from .select import Select  # noqa
+from .table_exists import TableExists  # noqa
+from .update import Update  # noqa

piccolo/query/methods/create_index.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+from enum import Enum
+import typing as t
+
+from piccolo.columns import Column
+from piccolo.query.base import Query
+from piccolo.querystring import QueryString
+
+if t.TYPE_CHECKING:
+    from piccolo.table import Table
+
+
+class IndexMethod(str, Enum):
+    btree = "btree"
+    hash = "hash"
+    gist = "gist"
+    gin = "gin"
+
+
+class CreateIndex(Query):
+    def __init__(
+        self,
+        table: t.Type[Table],
+        column: Column,
+        method: IndexMethod = IndexMethod.btree,
+    ):
+        self.column = column
+        self.method = method
+        super().__init__(table)
+
+    @property
+    def querystrings(self) -> t.Sequence[QueryString]:
+        index_name = self.column._meta.index_name
+        tablename = self.table._meta.tablename
+        method_name = self.method.value
+        return [
+            QueryString(
+                f"CREATE INDEX {index_name} ON {tablename} USING {method_name}"
+                f" ({self.column._meta.name})"
+            )
+        ]

piccolo/query/methods/drop_index.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+import typing as t
+
+from piccolo.query.base import Query
+from piccolo.querystring import QueryString
+
+if t.TYPE_CHECKING:
+    from piccolo.columns.base import Column
+    from piccolo.table import Table
+
+
+class DropIndex(Query):
+    def __init__(self, table: t.Type[Table], column: Column):
+        self.column = column
+        super().__init__(table)
+
+    @property
+    def querystrings(self) -> t.Sequence[QueryString]:
+        index_name = self.column._meta.index_name
+        return [QueryString(f"DROP INDEX {index_name}")]

piccolo/table.py

@@ -12,9 +12,10 @@
 from piccolo.columns.defaults.base import Default
 from piccolo.query import (
     Alter,
-    Create,
     Count,
+    Create,
     Delete,
+    DropIndex,
     Exists,
     Insert,
     Objects,
@@ -23,6 +24,7 @@
     TableExists,
     Update,
 )
+from piccolo.query.methods.create_index import CreateIndex, IndexMethod
 from piccolo.querystring import QueryString, Unquoted
 from piccolo.utils import _camel_to_snake
 
@@ -399,7 +401,7 @@ def raw(cls, sql: str, *args: t.Any) -> Raw:
         """
         Execute raw SQL queries on the underlying engine - use with caution!
 
-        await Band.raw('select * from band')
+        await Band.raw('select * from band').run()
 
         Or passing in parameters:
 
@@ -478,7 +480,7 @@ def alter(cls) -> Alter:
         """
         Used to modify existing tables and columns.
 
-        await Band.alter().rename_column(Band.popularity, 'rating')
+        await Band.alter().rename_column(Band.popularity, 'rating').run()
         """
         return Alter(table=cls)
 
@@ -535,10 +537,32 @@ def update(cls, values: t.Dict[Column, t.Any] = {}) -> Update:
 
         await Band.update().values(
             {Band.name: "Spamalot"}
-        ).where(Band.name=="Pythonistas")
+        ).where(
+            Band.name=="Pythonistas"
+        ).run()
         """
         return Update(table=cls).values(values)
 
+    @classmethod
+    def create_index(
+        cls, column: Column, method: IndexMethod = IndexMethod.btree,
+    ) -> CreateIndex:
+        """
+        Create a table index.
+
+        await Band.create_index(Band.name).run()
+        """
+        return CreateIndex(table=cls, column=column, method=method)
+
+    @classmethod
+    def drop_index(cls, column: Column) -> DropIndex:
+        """
+        Drop a table index.
+
+        await Band.drop_index(Band.name).run()
+        """
+        return DropIndex(table=cls, column=column)
+
     ###########################################################################
 
     @classmethod

tests/table/test_alter.py

@@ -1,4 +1,3 @@
-import decimal
 from unittest import TestCase
 
 from piccolo.columns import Integer, Numeric

tests/table/test_batch.py

@@ -1,9 +1,6 @@
 import asyncio
 import math
 
-from piccolo.engine.postgres import AsyncBatch, PostgresEngine
-from piccolo.querystring import QueryString
-
 from ..base import DBTestCase
 from ..example_project.tables import Manager
 

tests/table/test_create.py

@@ -1,4 +1,5 @@
 from unittest import TestCase
+
 from ..example_project.tables import Manager