-
Notifications
You must be signed in to change notification settings - Fork 86
/
create_index.py
76 lines (64 loc) · 2.12 KB
/
create_index.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
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: # 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,
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
def column_names(self) -> t.List[str]:
return [
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
index_name = self.table._get_index_name(column_names)
tablename = self.table._meta.tablename
method_name = self.method.value
column_names_str = ", ".join(column_names)
return [
QueryString(
f"{self.prefix} {index_name} ON {tablename} USING "
f"{method_name} ({column_names_str})"
)
]
@property
def sqlite_querystrings(self) -> t.Sequence[QueryString]:
column_names = self.column_names
index_name = self.table._get_index_name(column_names)
tablename = self.table._meta.tablename
method_name = self.method.value
if method_name != "btree":
raise ValueError("SQLite only support btree indexes.")
column_names_str = ", ".join(column_names)
return [
QueryString(
f"{self.prefix} {index_name} ON {tablename} "
f"({column_names_str})"
)
]