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

bb_integrations_lib/pipelines/wrappers/import_tank_readings.py

@@ -0,0 +1,192 @@
+from datetime import datetime, timedelta
+from typing import Iterable, Self
+
+import pytz
+from loguru import logger
+
+from bb_integrations_lib.gravitate.rita_api import GravitateRitaAPI
+from bb_integrations_lib.mappers.rita_mapper import RitaAPICachedMappingProvider
+from bb_integrations_lib.pipelines.parsers.tank_reading_parser import TankReadingParser
+from bb_integrations_lib.pipelines.steps.importing.load_imap_attachment_step import LoadIMAPAttachmentStep
+from bb_integrations_lib.pipelines.steps.importing.sftp_file_config_step import SFTPFileConfigStep
+from bb_integrations_lib.pipelines.steps.null_step import NullStep
+from bb_integrations_lib.pipelines.steps.processing.archive_gcs_step import ArchiveGCSStep
+from bb_integrations_lib.pipelines.steps.processing.archive_sftp_step import ArchiveSFTPStep
+from bb_integrations_lib.pipelines.steps.processing.bbd_upload_tankreading_step import BBDUploadTankReadingStep
+from bb_integrations_lib.pipelines.steps.processing.delete_sftp_step import DeleteSFTPStep
+from bb_integrations_lib.pipelines.steps.processing.file_config_parser_step_v2 import FileConfigParserV2
+from bb_integrations_lib.pipelines.wrappers.wrapper import PipelineWrapper
+from bb_integrations_lib.protocols.pipelines import JobPipeline
+from bb_integrations_lib.secrets import SecretProvider
+from bb_integrations_lib.secrets.factory import APIFactory
+from bb_integrations_lib.shared.model import ImportTankReadings, ATGConfig, ConfigMode
+from bb_integrations_lib.util.config.manager import GlobalConfigManager
+from bb_integrations_lib.util.config.model import GlobalConfig
+
+
+class BBDImportTankReadingsPipeline(JobPipeline):
+    @classmethod
+    async def create(
+            cls,
+            config: ATGConfig,
+            secret_provider: SecretProvider,
+            rita_client: GravitateRitaAPI,
+            bucket_name: str,
+            config_name: str,
+            tenant_name: str,
+            send_reports: bool = True
+    ):
+        api_factory = APIFactory(secret_provider)
+
+        steps = [{
+            "id": "start",
+            "parent_id": None,
+            "step": NullStep({})
+        }]
+        if config.minutes_back is not None and config.timezone is not None:
+            min_date = datetime.now(pytz.timezone(config.timezone)) - timedelta(minutes=config.minutes_back)
+        else:
+            min_date = None
+
+        ftp_client = None
+        if config.mode == "ftp":
+            ftp_client = await api_factory.ftp(config.ftp_credentials)
+            pull_step = SFTPFileConfigStep(
+                rita_client=rita_client,
+                ftp_client=ftp_client,
+                bucket_name=bucket_name,
+                config_name=config_name,
+                mode=ConfigMode.ByName,
+                match_mode=config.file_match_type,
+                min_date=min_date or datetime.min
+            )
+        elif config.mode == "email":
+            pull_step = LoadIMAPAttachmentStep(
+                rita_client=rita_client,
+                imap_client=await api_factory.imap(config.email_credentials),
+                from_email_address=config.from_email_address,
+                delivered_to_email_address=config.delivered_to_email_address,
+                to_email_address=config.to_email_address,
+                attachment_extension=config.attachment_extension or ".csv",
+                bucket_name=bucket_name,
+                config_names=[config_name],
+                email_subject=config.email_subject
+            )
+        else:
+            raise ValueError(f"Invalid mode: {config.mode}")
+
+        mapping_provider = RitaAPICachedMappingProvider(rita_client)
+
+        steps += [
+            {
+                "id": "pull",
+                "parent_id": "start",
+                "step": pull_step
+            },
+            {
+                "id": f"parse_readings",
+                "parent_id": "pull",
+                "step": FileConfigParserV2(
+                    rita_client=rita_client,
+                    mapping_type=config.mapping_type,
+                    parser=TankReadingParser,
+                    parser_kwargs={
+                        "included_payload": {
+                            "source": f"{tenant_name} flat file tank readings"
+                        },
+                        "mapping_provider": mapping_provider,
+                        "sd_client": await api_factory.sd(config.sd_credentials),
+                    }
+                )
+            },
+            {
+                "id": f"bbd_upload",
+                "parent_id": f"parse_readings",
+                "step": BBDUploadTankReadingStep(
+                    sd_client=await api_factory.sd(config.sd_credentials),
+                )
+            },
+        ]
+        if config.archive_gcs_bucket_path is not None and config.gcs_credentials is not None:
+            steps.append({
+                "id": f"archive_gcs",
+                "parent_id": f"bbd_upload",
+                "alt_input": "pull",
+                "step": ArchiveGCSStep(
+                    gcloud_storage=await api_factory.gcloud_storage(config.gcs_credentials),
+                    bucket_path=config.archive_gcs_bucket_path,
+                )
+            })
+        if config.archive_files:
+            if ftp_client:
+                steps.append({
+                    "id": f"archive_ftp",
+                    "parent_id": f"bbd_upload",
+                    "alt_input": "pull",
+                    "step": ArchiveSFTPStep(
+                        ftp_client=ftp_client,
+                    )
+                })
+            else:
+                logger.warning("No FTP credentials provided, skipping FTP archive step")
+        if config.delete_files:
+            if ftp_client:
+                steps.append({
+                    "id": "delete_ftp",
+                    "parent_id": f"bbd_upload",
+                    "alt_input": "pull",
+                    "step": DeleteSFTPStep(
+                        ftp_client=ftp_client
+                    )
+                })
+            else:
+                logger.warning("No FTP credentials provided, skipping FTP delete step")
+        config = await rita_client.get_config_by_name(bucket_path="/" + bucket_name, config_name=config_name)
+        config_id = config[config_name].config_id
+        return cls(
+            steps,
+            rita_client=rita_client,
+            pipeline_name=f"Import tank readings - {config_name}",
+            pipeline_config_id=config_id,
+            secret_provider=secret_provider,
+            send_reports=send_reports
+        )
+
+
+class MultiImportTankReadingsPipeline:
+    def __init__(self, pipelines: Iterable[BBDImportTankReadingsPipeline]):
+        self.pipelines = pipelines
+
+    @classmethod
+    async def create(cls, config_id: str, tenant_name: str, rita_client: GravitateRitaAPI,
+                     secret_provider: SecretProvider, send_reports: bool = True) -> Self:
+        main_config = await rita_client.get_config_by_id(config_id)
+        [bucket] = [b for b in await rita_client.get_all_buckets() if b["_id"] == main_config.owning_bucket_id]
+
+        import_config = ImportTankReadings.model_validate(main_config.config)
+        pipelines = []
+        for atg_config in import_config.configs:
+            for config_name in atg_config.config_names:
+                logger.info(f"Creating pipeline for config {config_name}")
+                pipelines.append(await BBDImportTankReadingsPipeline.create(
+                    config=atg_config,
+                    secret_provider=secret_provider,
+                    rita_client=rita_client,
+                    bucket_name=bucket["name"],
+                    config_name=config_name,
+                    tenant_name=tenant_name,
+                    send_reports=send_reports
+                ))
+        return cls(pipelines)
+
+    async def run(self) -> None:
+        if not self.pipelines:
+            raise Exception("No pipelines to run - check config and try again")
+        for pipeline in self.pipelines:
+            try:
+                logger.info(f"Starting pipeline '{pipeline.pipeline_name}'")
+                await pipeline.execute()
+                logger.info(f"Finished pipeline '{pipeline.pipeline_name}'")
+                logger.info("-" * 40)
+            except Exception as e:
+                logger.exception(f"Failed to run pipeline '{pipeline.pipeline_name}'")

bb_integrations_lib/pipelines/wrappers/wrapper.py

@@ -0,0 +1,81 @@
+from typing import Tuple, TypeVar, Generic, Type, Optional
+from bb_integrations_lib.gravitate.rita_api import GravitateRitaAPI
+from bb_integrations_lib.models.rita.config import GenericConfig
+from bb_integrations_lib.shared.exceptions import ConfigNotFoundError, ConfigValidationError
+from bb_integrations_lib.util.config.manager import GlobalConfigManager
+from bb_integrations_lib.util.config.model import GlobalConfig, ClientConstructor
+
+T = TypeVar("T")
+
+
+
+
+class PipelineWrapper(Generic[T]):
+    def __init__(self,
+                 tenant_name: str,
+                 bucket_name: str,
+                 config_class: Type[T],
+                 mode: Optional[str] = "production",
+                 ):
+        self.tenant_name = tenant_name
+        self.bucket_name = bucket_name
+        self.config_class = config_class
+        self.config_name: Optional[str] = None
+        self.config_manager = GlobalConfigManager()
+        self.mode = mode
+
+    async def load_config(self,
+                          config_name: str,
+                          rita_client: GravitateRitaAPI = None,
+                          ) -> Tuple[T, str, str]:
+        """
+        Load and validate configuration from Rita API.
+
+        Returns:
+            Tuple of (parsed_config, config_id, config_name)
+
+        Raises:
+            ConfigNotFoundError: If the configuration is not found
+            ConfigValidationError: If the configuration fails validation
+        """
+        if not rita_client:
+            rita_client = GravitateRitaAPI(
+                client_id=self.secret_data.prod.rita.client_id,
+                client_secret=self.secret_data.prod.rita.client_secret,
+                tenant=self.tenant_name,
+            )
+        try:
+            configs = await rita_client.get_config_by_name(
+                bucket_path=self.bucket_name,
+                config_name=config_name
+            )
+        except Exception as e:
+            msg = f"Failed to retrieve config '{config_name}' from bucket '{self.bucket_name}': {e}"
+            raise ConfigNotFoundError(
+                msg
+            ) from e
+
+        job_config: GenericConfig = configs[config_name]
+
+        try:
+            pipeline_config: T = self.config_class.model_validate(job_config.config)
+        except Exception as e:
+            raise ConfigValidationError(
+                f"Failed to validate config '{config_name}' as {self.config_class.__name__}: {e}"
+            ) from e
+
+        return pipeline_config, job_config.config_id, config_name
+
+    @property
+    def secret_data(self) -> GlobalConfig:
+        return self.config_manager.get_environment(self.tenant_name)
+
+    @property
+    def api_clients(self) -> ClientConstructor:
+        return self.config_manager.environment_from_name(self.tenant_name, self.mode)
+
+    def __repr__(self) -> str:
+        return (f"PipelineWrapper(tenant='{self.tenant_name}', "
+                f"bucket='{self.bucket_name}', "
+                f"type={self.config_class.__name__})")
+

bb_integrations_lib/protocols/flat_file.py

@@ -0,0 +1,210 @@
+from enum import Enum
+from typing import Protocol, Iterable, Optional, List, Literal
+from datetime import datetime
+from pydantic import BaseModel, Field, field_validator
+from bb_integrations_lib.shared.model import RawData, File
+
+
+class FileStorageClient(Protocol):
+    """Protocol for classes that can upload and download files from some source, when credentials are provided."""
+
+    def list_files(self, directory: str, credential_file_name: str = None) -> list[str]:
+        """Integration method to list files in a remote directory"""
+        pass
+
+    def get_raw_data(self, min_date: datetime, credential_file_name: str = None, check_date: bool = True) -> \
+            Iterable[RawData]:
+        """Integration method to download a file from a remote directory"""
+        pass
+
+    def rename_file(self, old_name: str, new_name: str, credential_file_name: str = None) -> None:
+        """Integration method to rename a remote file"""
+        pass
+
+    def delete_file(self, path: str, credential_file_name: str = None) -> None:
+        """Integration method to delete a remote file"""
+        pass
+
+    def upload_file(self, file: File, path: str, credential_file_name: str = None) -> None:
+        """Integration method to upload a file to a remote directory"""
+        pass
+
+
+class Integration(Protocol):
+    """
+    An integration retrieves data from a source and returns an iterable of RawData objects.
+    """
+
+    def get_raw_data(self, min_date: datetime = datetime.min) -> Iterable[RawData]:
+        """Method for fetching the raw data from the integration"""
+
+    def archive_data(self, raw_data: RawData):
+        """Method that archives the data after it has been processed"""
+
+
+class Translator(Protocol):
+    def translate_row(self, row: dict) -> dict:
+        """Translate the data in the provided row. Returns a new dictionary that has been translated."""
+
+
+class Parser(Protocol):
+    def parse_raw_data(self, raw_data: RawData) -> Iterable:
+        """Method to process the raw data from the get_raw_data Integration method."""
+
+
+class TankMonitorType(str, Enum):
+    bbd = "bbd"
+    aws = "aws"
+
+
+class TankReading(BaseModel):
+    date: str  # ISO Datetime
+    payload: dict = {}
+    store: str
+    tank: str
+    timezone: str | None
+    volume: float
+    monitor_type: TankMonitorType
+
+
+class SalesAdjustedDeliveryReading(BaseModel):
+    source: str
+    number: str
+    store_id: str
+    tank_id: str
+    product_id: str
+    volume: float
+    date: datetime
+
+class TankSales(BaseModel):
+    store_number: str
+    tank_id: str
+    date: str  # ISO Datetime
+    sales: float
+
+class PELookup(BaseModel):
+    """Lookup model for common source identifiers"""
+    SourceSystemId: Optional[int] = None
+    SourceId: Optional[int] = None
+    SourceId2: Optional[int] = None
+    SourceId3: Optional[int] = None
+    SourceIdString: Optional[str] = None
+
+
+class PriceInstrumentDTO(BaseModel):
+    """Price Instrument Data Transfer Object"""
+    Name: str
+    Abbreviation: str
+    UnitOfMeasureLookup: Optional[PELookup] = None
+    CurrencyLookup: Optional[PELookup] = None
+    BookLookup: Optional[PELookup] = None
+    ProductLookup: Optional[PELookup] = None
+    LocationLookup: Optional[PELookup] = None
+    ToCurrencyLookup: Optional[PELookup] = None
+    CounterPartyLookup: Optional[PELookup] = None
+    IsActive: bool = True
+    ExternalReferenceNumber: Optional[str] = None
+    SourceId: Optional[int] = None
+    SourceIdString: Optional[str] = None
+
+
+class PriceTypeDTO(BaseModel):
+    """Price Type Data Transfer Object"""
+    PriceTypeMeaning: str
+    ExtractPrices: bool = True
+
+
+class BulkSyncIntegrationDTO(BaseModel):
+    """Integration Data Transfer Object"""
+    Name: Optional[str] = None
+    Abbreviation: Optional[str] = None
+    PricePublisherTypeMeaning: Optional[str] = None
+    IsActive: bool = True
+    MatchByType: Literal["SourceId", "SourceIdString"] = None
+    PriceInstrumentDTOs: List[PriceInstrumentDTO] = Field(default_factory=list)
+    PriceTypeDTOs: List[PriceTypeDTO] = Field(default_factory=list)
+    SourceId: Optional[int] = None
+    SourceIdString: Optional[str] = None
+
+
+
+class PeBulkSyncIntegrationOptions(BaseModel):
+    """Options configuration"""
+    IsPartialDataSet: bool = True
+
+
+class PeBulkSyncIntegration(BaseModel):
+    """Root model for price integration data"""
+    Options: PeBulkSyncIntegrationOptions = PeBulkSyncIntegrationOptions()
+    IntegrationDtos: List[BulkSyncIntegrationDTO] = Field(default_factory=list)
+    SourceSystemId: int = None
+
+class PriceMergeValue(BaseModel):
+    Value: float
+    PriceTypeMeaning: Literal["High", "Low", "Posting", "Average"] = "Posting"
+
+class PriceMergeIntegrationDTO(BaseModel):
+    PriceInstrumentLookup: Optional[PELookup] = None
+    EstimateActual: Literal["A", "S"] = "A"
+    EffectiveFromDateTime: str = None
+    EffectiveToDateTime: Optional[str] = None
+    TradePeriodFromDateTime: Optional[str] = None
+    TradePeriodToDateTime: Optional[str] = None
+    UnitOfMeasureLookup: Optional[PELookup] = None
+    IsActive: Optional[bool] = True
+    PriceValues: List[PriceMergeValue]
+
+class PePriceMergeIntegration(BaseModel):
+    IntegrationDtos: List
+    SourceSystemId: Optional[int] = None
+
+class DriverCredential(BaseModel):
+    driver_id: dict
+    expiration_date: str | None  # ISO Datetime
+    certification_date: str | None  # ISO Datetime
+    credential_id: str | None
+
+
+class PriceRow(BaseModel):
+    effective_from: str  # ISO Datetime
+    effective_to: str  # ISO Datetime
+    price: float
+    price_type: str
+    contract: Optional[str] = ""  # This is to make sure we don't break contracts
+    timezone: Optional[str] = "UTC"
+    terminal_id: Optional[str] = None
+    terminal_source_id: Optional[str] = None
+    terminal_source_system_id: Optional[str] = None
+    terminal: Optional[str] = None
+    product_id: Optional[str] = None
+    product_source_id: Optional[str] = None
+    product_source_system_id: Optional[str] = None
+    product: Optional[str] = None
+    supplier_id: Optional[str] = None
+    supplier_source_id: Optional[str] = None
+    supplier_source_system_id: Optional[str] = None
+    counterparty_source_id: Optional[str] = None
+    counterparty_source_system_id: Optional[str] = None
+    supplier: Optional[str] = None
+    enabled: Optional[bool] = None
+    disabled_until: Optional[str] = None
+    expire: Optional[str] = None
+    min_quantity: Optional[str] = None
+    max_quantity: Optional[str] = None
+    curve_id: Optional[str] = None
+    row: Optional[str] = None
+
+
+class TankSales(BaseModel):
+    store_number: str
+    tank_id: str
+    date: str  # ISO Datetime
+    sales: float
+
+
+class UniqueTank(BaseModel):
+    store: str
+    tank_id: str
+    monitor_type: TankMonitorType
+
+

bb_integrations_lib/protocols/gravitate_client.py

@@ -0,0 +1,104 @@
+from typing import Protocol, Optional, Dict, Any
+from pymongo.database import Database
+
+
+class System(Protocol):
+    sd: str
+    pe: str
+    rita: str
+
+
+class GravitateConfig(Protocol):
+    system_psk: Optional[str]
+    url: Optional[str]
+    conn_str: Optional[str]
+    dbs: Optional[Dict[str, str]]
+    system: Optional[System]
+    short_name: Optional[str]
+    admin_username: Optional[str]
+    admin_password: Optional[str]
+
+
+class SecureWebAPI(Protocol):
+    system_name: str
+    username: Optional[str]
+    password: Optional[str]
+    client_id: Optional[str]
+    client_secret: Optional[str]
+    psk: str
+    base_url: str
+    token: Optional[str]
+    system: Any
+    extra_headers: Dict[str, str]
+
+    def post(self, url: str, params: Optional[Dict] = None, headers: Optional[Dict] = None,
+             data: Optional[Dict] = None, json: Optional[Dict] = None, files: Optional[Dict] = None,
+             timeout: int = 180, raise_error: bool = True) -> Any: ...
+
+    def get(self, url: str, params: Optional[Dict] = None, headers: Optional[Dict] = None,
+            timeout: int = 180, raise_error: bool = True) -> Any: ...
+
+    def token_post(self, url: str, params: Optional[Dict] = None, data: Optional[Dict] = None,
+                   json: Optional[Dict] = None, files: Optional[Dict] = None, timeout: int = 180,
+                   raise_error: bool = True, headers: Optional[Dict] = None) -> Any: ...
+
+    def token_get(self, url: str, params: Optional[Dict] = None, timeout: int = 180,
+                  raise_error: bool = True) -> Any: ...
+
+    def psk_post(self, url: str, params: Optional[Dict] = None, data: Optional[Dict] = None,
+                 json: Optional[Dict] = None, timeout: int = 180, psk_param: str = "system_psk",
+                 raise_error: bool = True) -> Any: ...
+
+    def get_service_api(self, service_name: str) -> 'SecureWebAPI': ...
+
+
+class DatabaseAPI(Protocol):
+    dbs: Dict[str, str]
+    conn_string: str
+    client: Any
+    backend: Optional[Database]
+    price: Optional[Database]
+    payroll: Optional[Database]
+    auth: Optional[Database]
+    forecast: Optional[Database]
+    ims: Optional[Database]
+    valuation: Optional[Database]
+
+    def get_db(self, db_name: str) -> Optional[Database]: ...
+
+
+class WebAPIWrapper(Protocol):
+    pass
+
+
+class GravitateAPIClient(Protocol):
+    system: System
+    web: SecureWebAPI
+    db: DatabaseAPI
+    client_id: Optional[str]
+    client_secret: Optional[str]
+    base_url: str
+    url: str
+    system_name: Optional[str]
+    short_name: Optional[str]
+    username: Optional[str]
+    password: Optional[str]
+    psk: Optional[str]
+    connection_string: Optional[str]
+    dbs: Optional[Dict[str, str]]
+    api_map: Dict[Any, WebAPIWrapper]
+    apis: WebAPIWrapper
+
+
+class SDWebAPIWrapper(WebAPIWrapper):
+    backend: GravitateAPIClient
+    forecast: GravitateAPIClient
+    valuation: GravitateAPIClient
+
+
+class PEWebAPIWrapper(WebAPIWrapper):
+    integration: GravitateAPIClient
+
+
+class RitaWebAPIWrapper(WebAPIWrapper):
+    backend: GravitateAPIClient