build.db.models

Complexity

Total Complexity

9

Size/Duplication

Total Lines	133
Duplicated Lines	0 %

Importance

Changes

0

4 Methods

Rating	Name	Duplication	Size	Complexity
A	SwitchDoc.projection()	0	27	1
A	SwitchDoc.preset_interfaces()	0	6	2
A	LinkDoc.projection()	0	13	1
A	DocumentBaseModel.dict()	0	7	5

"""DB models."""

from typing import List, Optional
from datetime import datetime

from pydantic import BaseModel
from pydantic import conlist
from pydantic import Field
from pydantic import validator


class DocumentBaseModel(BaseModel):
    """DocumentBaseModel."""

    id: str = Field(None, alias="_id")
    inserted_at: Optional[datetime]
    updated_at: Optional[datetime]

    def dict(self, **kwargs) -> dict:
        values = super().dict(**kwargs)
        if "id" in values and values["id"]:
            values["_id"] = values["id"]
        if "exclude" in kwargs and "_id" in kwargs["exclude"]:
            values.pop("_id")
        return values


class InterfaceSubDoc(BaseModel):
    """Interface DB SubDocument Model."""

    id: str
    enabled: bool
    active: bool
    mac: str
    speed: float
    port_number: int
    name: str
    nni: bool = False
    lldp: bool
    switch: str
    link: Optional[str]
    link_side: Optional[str]
    metadata: dict = {}
    updated_at: Optional[datetime]


class SwitchDoc(DocumentBaseModel):
    """Switch DB Document Model."""

    enabled: bool
    active: bool
    data_path: Optional[str]
    hardware: Optional[str]
    manufacturer: Optional[str]
    software: Optional[str]
    connection: Optional[str]
    ofp_version: Optional[str]
    serial: Optional[str]
    metadata: dict = {}
    interfaces: List[InterfaceSubDoc] = []

    @validator("interfaces", pre=True)
    def preset_interfaces(cls, v, values, **kwargs) -> List[InterfaceSubDoc]:
        """Preset interfaces."""
        if isinstance(v, dict):
            return list(v.values())
        return v

    @staticmethod
    def projection() -> dict:
        """Base projection of this model."""
        return {
            "_id": 0,
            "id": 1,
            "enabled": 1,
            "active": 1,
            "data_path": 1,
            "hardware": 1,
            "manufacturer": 1,
            "software": 1,
            "connection": 1,
            "ofp_version": 1,
            "serial": 1,
            "metadata": 1,
            "interfaces": {
                "$arrayToObject": {
                    "$map": {
                        "input": "$interfaces",
                        "as": "intf",
                        "in": {"k": "$$intf.id", "v": "$$intf"},
                    }
                }
            },
            "updated_at": 1,
            "inserted_at": 1,
        }


class InterfaceIdSubDoc(BaseModel):
    """InterfaceId DB SubDocument Model."""

    id: str


class LinkDoc(DocumentBaseModel):
    """Link DB Document Model."""

    enabled: bool
    active: bool
    metadata: dict = {}
    endpoints: conlist(InterfaceIdSubDoc, min_items=2, max_items=2)

    @staticmethod
    def projection() -> dict:
        """Base projection of this model."""
        return {
            "_id": 0,
            "id": 1,
            "enabled": 1,
            "active": 1,
            "metadata": 1,
            "endpoint_a": {"$first": "$endpoints"},
            "endpoint_b": {"$last": "$endpoints"},
            "updated_at": 1,
            "inserted_at": 1,
        }


class InterfaceDetailDoc(DocumentBaseModel):
    """InterfaceDetail DB Document Model."""

    available_vlans: List[int]