bb-integrations-library 3.0.11__py3-none-any.whl

bb_integrations_lib/models/rita/email.py

@@ -0,0 +1,39 @@
+import base64
+from typing import Optional
+
+from pydantic import BaseModel, model_serializer, computed_field
+
+
+class EmailAttachment(BaseModel):
+    file_name: str
+    file_data_encoded: Optional[str] = None
+
+    def __init__(self, file_name: str, file_data: Optional[bytes] = None, file_data_encoded: Optional[str] = None):
+        super().__init__(file_name=file_name)
+        self.file_name = file_name
+        self.file_data_encoded = self._b64enc(file_data) if file_data is not None else file_data_encoded
+
+    def _b64enc(self, v) -> str:
+        return base64.b64encode(v).decode("utf-8")
+
+    @computed_field
+    def file_data(self) -> bytes:
+        return base64.b64decode(self.file_data_encoded)
+
+    @file_data.setter
+    def file_data(self, new_file_data: bytes) -> None:
+        self.file_data_encoded = self._b64enc(new_file_data)
+
+    @model_serializer()
+    def serialize_model(self):
+        return {
+            "file_name": self.file_name,
+            "file_data_encoded": self.file_data_encoded
+        }
+
+
+class EmailData(BaseModel):
+    to: str | list[str]
+    html_content: str
+    subject: str
+    attachments: Optional[list[EmailAttachment]] = None

bb_integrations_lib/models/rita/issue.py

@@ -0,0 +1,63 @@
+from datetime import datetime
+from enum import Enum
+from typing import Optional, List, Annotated
+
+from annotated_types import MaxLen
+from pydantic import BaseModel, Field
+
+
+class IssueCategory(Enum):
+    UNKNOWN = "unknown"
+    MISC = "misc"
+    ORDER = "order"
+    PAYROLL = "payroll"
+    TANK_READING = "tank"
+    PRICE = "price"
+    REFERENCE_DATA = "reference_data"
+    CROSSROADS = "crossroads"
+    CROSSROADS_MAPPINGS = "crossroads_mappings"
+
+
+class IssueBase(BaseModel):
+    key: str = Field(description="Unique key for this issue")
+    config_id: str = Field(description="Config object this issue is reported under.")
+    # group or config id?
+    name: str = Field(description="Human readable name")
+    category: IssueCategory = Field(description="Broad category of the issue")
+    problem_short: Optional[str] = Field(default=None,
+                                         description="Short description of the problem (e.g. an exception class name)")
+    problem_long: Optional[str] = Field(default=None, description="Long description of the problem (e.g. traceback)")
+    occurrences: Annotated[
+        List[datetime], Field(default=[], description="List of most recent datetimes that this issue was reported",
+                              max_length=50)]
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @property
+    def most_recent_occurrence(self) -> Optional[datetime]:
+        return max(self.occurrences, default=None)
+
+    def get_occurrence_limit(self) -> int:
+        """Get the annotated max_length of the occurrences list."""
+        for meta in self.model_fields["occurrences"].metadata:
+            if isinstance(meta, MaxLen):
+                return meta.max_length
+        raise AttributeError("IssueBase.occurrences does not have an annotated max_length")
+
+    def trim_occurrences(self):
+        """
+        Trims the occurrences array to fit within the annotated Field.max_length.
+        The array gets sorted as a side effect.
+        """
+        self.occurrences = sorted(self.occurrences)[-self.get_occurrence_limit():]
+
+
+class UpdateIssue(BaseModel):
+    key: str = Field(description="Unique key for this issue")
+    config_id: Optional[str] = Field(default=None, description="Config object this issue is reported under.")
+    name: Optional[str] = Field(default=None, description="Human readable name")
+    category: Optional[IssueCategory] = Field(default=None, description="Broad category of the issue")
+    problem_short: Optional[str] = Field(default=None,
+                                         description="Short description of the problem (e.g. an exception class name)")
+    problem_long: Optional[str] = Field(default=None, description="Long description of the problem (e.g. traceback)")

bb_integrations_lib/models/rita/mapping.py

@@ -0,0 +1,227 @@
+from datetime import datetime, UTC
+from enum import Enum
+from typing import Optional, Dict, Self
+import json
+import pandas as pd
+from functools import cached_property
+from bson.objectid import ObjectId
+from pydantic import BaseModel, field_validator, Field
+
+
+class MappingType(str, Enum):
+    site = "site"
+    tank = "tank"
+    counterparty = "counterparty"
+    credential = "credential"
+    product = "product"
+    terminal = "terminal"
+    driver = "driver"
+    depot = "depot"
+    trailer = "trailer"
+    tractor = "tractor"
+    other = "other"
+    composite = "composite"
+
+
+class MapType(str, Enum):
+    parent = "parent"
+    child = "child"
+
+
+class View(str, Enum):
+    grid = "grid"
+    gallery = "gallery"
+
+
+class CompositeMapKey(BaseModel):
+    """Multi-field composite key for mapping lookups.
+
+    A composite key allows mappings to be identified by multiple fields
+    instead of a single source_id. For example, a mapping might be uniquely
+    identified by (tenant, product, terminal) rather than just source_id.
+
+    Example:
+        >>> key = CompositeMapKey(key={"tenant": "acme", "product": "fuel"})
+        >>> key.to_cache_key()
+        '{"product":"fuel","tenant":"acme"}'
+        >>> key.matches({"tenant": "acme"})
+        True
+    """
+    model_config = {"frozen": True}
+
+    key: dict[str, str] = Field(..., description="Field-value pairs forming the composite key")
+
+    @field_validator('key', mode='before')
+    @classmethod
+    def validate_key(cls, v: dict) -> dict[str, str]:
+        if not v:
+            raise ValueError("Composite key cannot be empty")
+        validated = {}
+        for k, val in v.items():
+            if not isinstance(k, str) or not k.strip():
+                raise ValueError(f"Key must be non-empty string, got: {k}")
+            if not isinstance(val, str):
+                raise ValueError(f"Value for '{k}' must be string, got {type(val).__name__}")
+            validated[k.strip()] = val
+        return validated
+
+    @cached_property
+    def _cache_key_str(self) -> str:
+        """Cached canonical string representation for hashing."""
+        return json.dumps(self.key, sort_keys=True, separators=(',', ':'))
+
+    def to_cache_key(self) -> str:
+        """Generate a canonical string representation for caching/hashing (JSON format)."""
+        return self._cache_key_str
+
+    @classmethod
+    def from_string(cls, s: str) -> Self:
+        """Parse a composite key from its JSON string representation."""
+        if not s:
+            raise ValueError("Cannot parse empty string")
+        try:
+            key_dict = json.loads(s)
+            if not isinstance(key_dict, dict):
+                raise ValueError("Invalid format: expected JSON object")
+            return cls(key=key_dict)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Failed to parse composite key: {e}")
+
+    @cached_property
+    def _cached_hash(self) -> int:
+        return hash(self._cache_key_str)
+
+    def __hash__(self) -> int:
+        return self._cached_hash
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, CompositeMapKey):
+            return False
+        return self.key == other.key
+
+    def __str__(self) -> str:
+        return self.to_cache_key()
+
+    def __repr__(self) -> str:
+        return f"CompositeMapKey({self.to_cache_key()})"
+
+    def __lt__(self, other: 'CompositeMapKey') -> bool:
+        if not isinstance(other, CompositeMapKey):
+            return NotImplemented
+        return self.to_cache_key() < other.to_cache_key()
+
+    def matches(self, partial: dict[str, str]) -> bool:
+        """
+        Check if this key contains all fields from the partial key.
+
+        Returns True if all key-value pairs in partial exist in this key.
+        Short-circuits on first mismatch for performance.
+        """
+        if len(partial) > len(self.key):
+            return False
+        for k, v in partial.items():
+            if self.key.get(k) != v:
+                return False
+        return True
+
+    def get(self, field: str, default: str | None = None) -> str | None:
+        """Get a specific component from the composite key."""
+        return self.key.get(field, default)
+
+    @property
+    def fields(self) -> list[str]:
+        """Return the list of field names in this composite key."""
+        return list(self.key.keys())
+
+
+class MapBase(BaseModel):
+    updated_by: str | None = Field(default="admin", description="User who last updated the mapping")
+    updated_on: datetime | None = Field(default_factory=lambda: datetime.now(UTC), description="Timestamp of last update")
+    type: Optional[MappingType | None] = Field(default=None, description="Type of mapping (site, tank, counterparty, etc.)")
+    source_system: str | None = Field(default=None, description="External system the mapping originates from")
+    is_active: Optional[bool] = Field(default=True, description="Whether the mapping is currently active")
+    source_name: str | None = Field(default=None, description="Display name in the source system")
+    gravitate_name: str | None = Field(default=None, description="Display name in Gravitate")
+    extra_data: Optional[Dict] = Field(default=None, description="Additional metadata for the mapping")
+
+    @field_validator('type', 'source_system', mode='before')
+    @classmethod
+    def validate_type_source_system(cls, v):
+        if isinstance(v, str) and (v == 'nan' or v == ''):
+            return None
+        elif isinstance(v, float) and pd.isna(v):  # Check for actual NaN (float nan)
+            return None
+        return v
+
+
+class Children(MapBase):
+    source_id: str | None = Field(
+        default=None,
+        description="Source system identifier"
+    )
+    gravitate_id: str | None = Field(
+        default=None,
+        description="Gravitate identifier"
+    )
+    id: Optional[str] = Field(
+        default_factory=lambda: str(ObjectId()),
+        description="Unique identifier for this child mapping"
+    )
+
+    @field_validator('source_id', 'gravitate_id', mode='before')
+    @classmethod
+    def validate_identifiers(cls, v):
+        if v is None or v == '' or v == 'nan':
+            return None
+        if isinstance(v, float) and pd.isna(v):
+            return None
+        return str(v)
+
+
+
+class Map(MapBase):
+    source_id: str | CompositeMapKey = Field(
+        ...,
+        description="Source system identifier (string or composite key)"
+    )
+    gravitate_id: str | CompositeMapKey = Field(
+        ...,
+        description="Gravitate identifier (string or composite key)"
+    )
+    children: Optional[list[Children]] = Field(
+        default_factory=list,
+        description="List of child mappings"
+    )
+    children_type: MappingType | None = Field(
+        default=None,
+        description="Type of child mappings"
+    )
+    owning_bucket_id: Optional[str] = Field(
+        default=None,
+        description="Bucket that owns this mapping"
+    )
+    group_id: str | None = Field(
+        default=None,
+        description="Group that owns this mapping"
+    )
+
+
+    def is_composite(self) -> bool:
+        """Check if this mapping uses composite keys."""
+        return isinstance(self.source_id, CompositeMapKey)
+
+    def source_id_str(self) -> str | None:
+        """Get source_id as string (for simple IDs) or None (for composite)."""
+        return self.source_id if isinstance(self.source_id, str) else None
+
+    def source_id_composite(self) -> CompositeMapKey | None:
+        """Get source_id as CompositeMapKey or None (for simple IDs)."""
+        return self.source_id if isinstance(self.source_id, CompositeMapKey) else None
+
+    def gravitate_id_str(self) -> str | None:
+        """Get gravitate_id as string (for simple IDs) or None (for composite)."""
+        return self.gravitate_id if isinstance(self.gravitate_id, str) else None
+
+    def gravitate_id_composite(self) -> CompositeMapKey | None:
+        """Get gravitate_id as CompositeMapKey or None (for simple IDs)."""
+        return self.gravitate_id if isinstance(self.gravitate_id, CompositeMapKey) else None

bb_integrations_lib/models/rita/probe.py

@@ -0,0 +1,58 @@
+from datetime import datetime
+from enum import Enum
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class AuthType(str, Enum):
+    no_auth = "no auth"
+    bearer_with_username = "bearer with username and password"
+    bearer_with_clientid = "bearer with client id"
+
+class ProbeSubscriber(BaseModel):
+    id: str | None = None
+    url: str | None = None
+    auth_url: str | None = None
+    username: str | None = None
+    password: str | None = None
+    psk: str | None = None
+    auth_type: AuthType = AuthType.no_auth
+    is_active: bool = True
+
+class WorkerTarget(BaseModel):
+    worker_name: str | None = None
+    kwargs: dict | None = None
+
+class ProbeConfig(BaseModel):
+    probe_id: str | None = None
+    is_active: bool = True
+    report_create: bool = False
+    report_delete: bool = False
+    report_update: bool = False
+    conn_str: str | None = None
+    database: str | None = None
+    collection: str | None = None
+    query: dict = {} # Events will only be reported if the object passes this filter.
+    entity_field: str | None = None
+    resume_token: str | None = None
+    update_check_interval_minutes: int = 0  # 0 implies a keep-alive probe that will always run.
+    external_subscribers: list[ProbeSubscriber] = []
+    worker_targets: list[WorkerTarget] = []  # List of RITA-owned workers to notify.
+    sending_entity_type: str | None = None
+    probe_type: str = "StandardProbe"  # The type of probe to be run on this config. Usually "StandardProbe"
+    resume_token_write_interval_seconds: int = 0  # At least this many seconds must pass before the resume token is written.
+    args: dict = {}  # arguments for custom probes.
+    probe_args: dict = {}
+    mode: Literal["test", "prod"] = "test"
+    output_request_data: str | None = None
+
+
+class ProbeStats(BaseModel):
+    time_started: str | None = None
+    changes_seen: int = 0
+    changes_processed: int = 0
+    num_times_updated_resume_token: int = 0
+    resume_token_time: str | None = None  # The timestamp that was associated with the last resume token updated. If the probe fell behind and is catching up, this may be significantly in the past.
+    resume_token_update_time: str | None = None  # The actual wall clock time when the resume token was last updated.
+    extra_data: dict = {}

bb_integrations_lib/models/rita/reference_data.py

@@ -0,0 +1,110 @@
+from typing import Any, Dict, Optional, Self
+from enum import Enum
+from pydantic import BaseModel
+from datetime import datetime, UTC
+
+class ReferenceDataType(str, Enum):
+    terminal = "terminal"
+    store = "store"
+    tank = "tank"
+    location = "location"
+    product = "product"
+    counterparty = "counterparty"
+
+
+class CoreMasterReferenceData(BaseModel):
+    """
+    A single 'item' in the industry. This could point to a store, a product, a company, or anything else.
+    The data here lives in the RITA master tenant only, and works in conjunction with the data stored in RITA client
+    tenant's tables to enable conversion from one client's representation to another.
+    """
+    data_type: ReferenceDataType
+    """The data type of this item. Used for lookups."""
+
+    mrid: str
+    """The unique ID of this item. MasterReferenceLink objects will link on this field."""
+
+    name: str
+    """Display name of this item."""
+
+    source_tenant: str | None = None
+    """The original tenant of this item, if it has one."""
+
+    matching_info: Dict[str, Any] = {}
+    """A dictionary describing the info that makes this object unique. E.g. for terminals this could be the federal TCN."""
+    # We plan to engineer some way to throw a client's data at an LLM for matching. Matching_info would be the place to store
+    # data that's similar in structure.
+
+    mrd_extra_data: Dict[str, Any] = {}
+    """Any additional data about the object that should be tied to it during the conversion process."""
+
+    updated_by: Optional[str] = None
+    updated_on: Optional[datetime] = datetime.now(UTC)
+
+
+class MasterReferenceData(CoreMasterReferenceData):
+    children: Dict[str, CoreMasterReferenceData] = {}
+    """The children of this master reference data"""
+
+
+class CoreMasterReferenceLink(BaseModel):
+    """
+    Links 'items' in this tenant's gravitate instance to the MasterReferenceData. The data here lives only in RITA client
+    tenants, not the RITA master tenant, and matches some record in the MasterReferenceData. It is possible for one
+    MasterReferenceLink to match one MasterReferenceData.
+    """
+    data_type: ReferenceDataType
+    """The data type of this item. Used for lookups."""
+
+    mrid: Optional[str] = None
+    """Linked MasterReferenceData item. This may be empty and indicates that there is no associated item."""
+
+    display_name: str
+    """
+    Display name of this item, in the context of the tenant's gravitate. This is not used for code lookups and can be
+    set to anything. It may be set automatically by the sync module, if the source_system is "Gravitate"
+    """
+
+    source_id: str
+    """
+    Either (but not both): 
+      - Mongodb ID of the item in the tenant's gravitate (if source_system is "Gravitate").
+      - ID of the item in the specified source system.
+    """
+
+    matching_info: Dict[str, Any] = {}
+    """A dictionary describing the info that makes this object unique. E.g. for terminals this could be the federal TCN, or an address."""
+
+    mrl_extra_data: Dict[str, Any] = {}
+    """Any additional data about the object that should be tied to it during the conversion process."""
+
+
+    updated_by: Optional[str] = None
+    updated_on: Optional[datetime] = datetime.now(UTC)
+
+
+class MasterReferenceLink(CoreMasterReferenceLink):
+    source_system: str = "Gravitate"
+    """The source system of this item. Defaults to 'Gravitate'. Used in mapping requests to add detail in mappings."""
+
+    children: Dict[str, CoreMasterReferenceLink] = {}
+    """Child references keyed by their Source ID. These can only match child references of the parent's matched master
+    reference data."""
+
+
+class ReferenceDataMappingExtraData(BaseModel):
+    origin_mrl_extra_data: dict[str, Any] = {}
+    mrd_extra_data: dict[str, Any] = {}
+
+
+class ReferenceDataMapping(BaseModel):
+    origin_source_id: str = ""
+    origin_source_system: str = ""
+    origin_mrl_id: str | None = None
+    origin_tenant: str = ""
+    target_tenant: str | None = None
+    target_mrls: list[MasterReferenceLink] = []
+    matched_mrid: str | None = None
+    matched_child_mrid: str | None = None
+    extra_data: ReferenceDataMappingExtraData = ReferenceDataMappingExtraData()
+    milliseconds_taken: int | None = None

bb_integrations_lib/models/rita/source_system.py

@@ -0,0 +1,9 @@
+from datetime import datetime, UTC
+
+from pydantic import BaseModel
+
+class SourceSystem(BaseModel):
+    updated_by: str | None = None,
+    updated_on: datetime | None = datetime.now(UTC)
+    name: str
+    description: str

bb_integrations_lib/models/rita/workers.py

@@ -0,0 +1,76 @@
+from datetime import datetime, UTC
+from enum import Enum
+from typing import Optional, Union, Literal, Annotated, Any
+
+from pydantic import BaseModel, Field
+
+
+class UnknownRunnableException(Exception):
+    pass
+
+WorkerTaskState = Literal["sent", "received", "working", "failed", "completed"]
+
+class WorkerRequest(BaseModel):
+    """A request for a specific runnable to be executed on an available worker."""
+    runnable_name: str
+    runnable_kwargs: dict[str, Any] = {}
+    tenant_name: str
+    originator: Literal["backend", "probe", "other"] = "other"
+
+class WorkerContentFile(BaseModel):
+    """
+    A file included in a worker's response. Supports binary (non-text) files by automatically encoding/decoding from
+    base64 when serialized.
+    """
+    file_name: str
+    mime_type: str | None = None
+    content: bytes | None = None
+
+    class Config:
+        ser_json_bytes = "base64"
+        val_json_bytes = "base64"
+
+class WorkerResponseStatus(str, Enum):
+    success = "success"
+    error = "error"
+
+class WorkerSuccessResponse(BaseModel):
+    """A response from a worker indicating success and providing results."""
+    # This is a pretty weird field definition. Basically, to support using this field as a discriminator for the
+    # WorkerResponse union type, we need it to only ever be status="success", (or "error", in the error class).
+    # These type annotations mean it is a WorkerResponseStatus that can only ever be .success, and that it cannot be set
+    # in the init method (init=False) or anywhere else (init_var=True).
+    # You can't use a frozen field here because it won't play nice with Beanie serialization.
+    status: Annotated[Literal[WorkerResponseStatus.success], Field(init_var=True, init=False)] = WorkerResponseStatus.success
+    runnable_name: str
+    content: str | None = None
+    content_file: WorkerContentFile | None = None
+    extra_data: dict = {}
+
+class WorkerErrorResponse(BaseModel):
+    """A response from a worker indicating an error occurred and providing error details."""
+    status: Annotated[Literal[WorkerResponseStatus.error], Field(init_var=True, init=False)] = WorkerResponseStatus.error
+    runnable_name: str
+    exception_type_name: str
+    error_message: str
+    reference_code: Optional[str] = None
+
+WorkerResponse = Annotated[Union[WorkerSuccessResponse, WorkerErrorResponse], Field(discriminator="status")]
+
+
+class WorkerUpdate(BaseModel):
+    """Provides status updates to the requester for backgrounded tasks."""
+    tenant_name: str
+    originator: str
+    sent_at: datetime = Field(default_factory=lambda: datetime.now(tz=UTC))
+    state: WorkerTaskState
+    response: Optional[WorkerResponse] = None
+
+class WorkerTask(BaseModel):
+    correlation_id: str
+    created_at: datetime
+    updated_at: datetime
+    mode: Literal["immediate", "background"]
+    state: WorkerTaskState
+    request: WorkerRequest
+    response: Optional[WorkerResponse] = None