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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
|
from abc import ABCMeta, abstractmethod
from typing import Any, Dict, Optional, Sequence, Tuple, Union, cast
import pymongo
from odmantic.field import FieldProxy
from odmantic.query import FieldProxyAny, SortExpression, asc
class ODMBaseIndex(metaclass=ABCMeta):
def __init__(self, unique: bool, index_name: Optional[str]) -> None:
self.unique = unique
self.index_name = index_name
@abstractmethod
def get_index_specifier(self) -> Sequence[Tuple[str, int]]: ...
def get_pymongo_index(self) -> pymongo.IndexModel:
kwargs: Dict[str, Any] = {"keys": self.get_index_specifier()}
if self.index_name is not None:
kwargs["name"] = self.index_name
if self.unique:
kwargs["unique"] = True
return pymongo.IndexModel(**kwargs)
class ODMSingleFieldIndex(ODMBaseIndex):
def __init__(self, key_name: str, unique: bool, index_name: Optional[str] = None):
super().__init__(unique, index_name)
self.key_name = key_name
def get_index_specifier(self) -> Sequence[Tuple[str, int]]:
return [
(self.key_name, pymongo.ASCENDING),
]
class ODMCompoundIndex(ODMBaseIndex):
def __init__(
self,
fields: Tuple[SortExpression, ...],
unique: bool,
index_name: Optional[str],
):
super().__init__(unique, index_name)
self.fields = fields
def get_index_specifier(self) -> Sequence[Tuple[str, int]]:
return [
(
list(f.keys())[0],
pymongo.ASCENDING if list(f.values())[0] == 1 else pymongo.DESCENDING,
)
for f in self.fields
]
class Index:
def __init__(
self,
*fields: Union[FieldProxyAny, SortExpression],
unique: bool = False,
name: Optional[str] = None,
) -> None:
"""Declare an ODM index in the Model.Config.indexes generator.
Example usage:
```python
from odmantic import Model, Index
from odmantic.query import desc
class Player(Model):
name: str
score: int
model_config = {
"indexes": lambda: [Index(Player.name, desc(Player.score))],
}
```
Args:
*fields (Any | SortExpression | str): fields to build the index with
unique: build a unique index
name: specify an optional custom index name
"""
self.fields = cast(Tuple[Union[SortExpression, FieldProxy], ...], fields)
self.unique = unique
self.name = name
def to_odm_index(self) -> "ODMBaseIndex":
if len(self.fields) == 1:
field = self.fields[0]
if isinstance(field, SortExpression):
key_name = list(field.keys())[0]
else:
key_name = object.__getattribute__(field, "_get_key_name")()
return ODMSingleFieldIndex(
key_name, unique=self.unique, index_name=self.name
)
else:
fields = tuple(
(f if isinstance(f, SortExpression) else asc(f) for f in self.fields)
)
return ODMCompoundIndex(fields, unique=self.unique, index_name=self.name)
|
