sift_py.data_import.status

  1import time
  2from datetime import datetime
  3from enum import Enum
  4from typing import Optional, Union
  5from urllib.parse import urljoin
  6
  7from pydantic import BaseModel, ConfigDict, field_validator
  8from pydantic.alias_generators import to_camel
  9from pydantic_core import PydanticCustomError
 10from typing_extensions import Self
 11
 12from sift_py.rest import SiftRestConfig, _RestService
 13
 14
 15class DataImportStatusType(Enum):
 16    """Status of the data import."""
 17
 18    SUCCEEDED = "DATA_IMPORT_STATUS_SUCCEEDED"
 19    PENDING = "DATA_IMPORT_STATUS_PENDING"
 20    IN_PROGRESS = "DATA_IMPORT_STATUS_IN_PROGRESS"
 21    FAILED = "DATA_IMPORT_STATUS_FAILED"
 22
 23    @classmethod
 24    def from_str(cls, val: str) -> Optional[Self]:
 25        try:
 26            return cls(val)
 27        except ValueError:
 28            return None
 29
 30    def as_human_str(self) -> str:
 31        return self.value
 32
 33
 34class DataImport(BaseModel):
 35    """Metadata regarding the data import."""
 36
 37    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
 38
 39    data_import_id: str
 40    created_date: datetime
 41    modified_date: datetime
 42    source_url: str = ""
 43    status: Union[str, DataImportStatusType]
 44    error_message: str = ""
 45    csv_config: dict
 46
 47    @field_validator("status", mode="before")
 48    @classmethod
 49    def convert_status(cls, raw: Union[str, DataImportStatusType]) -> DataImportStatusType:
 50        if isinstance(raw, DataImportStatusType):
 51            return raw
 52        elif isinstance(raw, str):
 53            value = DataImportStatusType.from_str(raw)
 54            if value is not None:
 55                return value
 56
 57        raise PydanticCustomError(
 58            "invalid_data_import_error", f"Invalid data import status: {raw}."
 59        )
 60
 61
 62class DataImportService(_RestService):
 63    """
 64    Service used to retrieve information about a particular data import.
 65    """
 66
 67    STATUS_PATH = "/api/v1/data-imports"
 68    _data_import_id: str
 69
 70    # TODO: rename restconf to rest_conf for consistency between services
 71    def __init__(self, restconf: SiftRestConfig, data_import_id: str):
 72        super().__init__(rest_conf=restconf)
 73        self._data_import_id = data_import_id
 74        self._status_uri = urljoin(self._base_uri, self.STATUS_PATH)
 75
 76    def get_data_import(self) -> DataImport:
 77        """
 78        Returns information about the data import.
 79        """
 80        response = self._session.get(
 81            url=f"{self._status_uri}/{self._data_import_id}",
 82        )
 83        response.raise_for_status()
 84        data = response.json().get("dataImport")
 85        data_import = DataImport(**data)
 86        return data_import
 87
 88    def wait_until_complete(self) -> DataImport:
 89        """
 90        Blocks until the data import is completed. Check the status to determine
 91        if the import was successful or not.
 92        """
 93        polling_interval = 1
 94        while True:
 95            data_import = self.get_data_import()
 96            status: DataImportStatusType = data_import.status  # type: ignore
 97            if status in [
 98                DataImportStatusType.SUCCEEDED,
 99                DataImportStatusType.FAILED,
100            ]:
101                return data_import
102            elif status in [
103                DataImportStatusType.PENDING,
104                DataImportStatusType.IN_PROGRESS,
105            ]:
106                pass
107            else:
108                raise Exception(f"Unknown status: {status}")
109            time.sleep(polling_interval)
110            polling_interval = min(polling_interval * 2, 60)