mongomantic.core.index.Index.to_pymongo()   B
last analyzed

Complexity

Conditions 5

Size

Total Lines 23
Code Lines 19

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
eloc 19
dl 0
loc 23
rs 8.9833
c 0
b 0
f 0
cc 5
nop 1
1
from typing import List, Optional
2
3
from pydantic import BaseModel, Field
4
from pymongo import ASCENDING, DESCENDING, TEXT, IndexModel
5
6
7
class Index(BaseModel):
8
    name: Optional[str] = Field(
9
        default=None,
10
        description="Custom name to use for this index - if none is given, a name will be generated.",
11
    )
12
13
    fields: List[str] = Field(
14
        min_items=1,
15
        description=(
16
            "Fields to index. Can be prefixed with '+' or '-' to specify index direction as ascending"
17
            "or descending. Prefix with '$' to specify text index."
18
        ),
19
    )
20
21
    unique: Optional[bool] = Field(default=False, descrition="If True, creates a uniqueness constraint on the index.")
22
23
    sparse: Optional[bool] = Field(
24
        default=False, description="If True, omit from the index any documents that lack the indexed field."
25
    )
26
27
    background: Optional[bool] = Field(
28
        default=False, description="If True, this index should be created in the background."
29
    )
30
31
    # Used to create an expiring (TTL) collection.
32
    # MongoDB will automatically delete documents from this collection after <int> seconds.
33
    # The indexed field must be a UTC datetime or the data will not expire.
34
    expire_after_seconds: Optional[int] = Field(
35
        default=0,
36
        description="Used to create an expiring (TTL) collection. Documents automatically deleted after <int> seconds.",
37
    )
38
39
    def to_pymongo(self):
40
        # Create pymongo index models
41
        pymongo_fields = []
42
        for field in self.fields:
43
            # Process prefix
44
            direction = ASCENDING
45
            if field.startswith("-"):
46
                direction = DESCENDING
47
                field = field[1:]
48
            elif field.startswith("+"):
49
                field = field[1:]
50
            elif field.startswith("$"):
51
                direction = TEXT
52
                field = field[1:]
53
54
            pymongo_fields.append((field, direction))
55
56
        return IndexModel(
57
            pymongo_fields,
58
            unique=self.unique,
59
            background=self.background,
60
            sparse=self.sparse,
61
            expireAfterSeconds=self.expire_after_seconds,
62
        )
63