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

bb_integrations_lib/pipelines/steps/importing/sftp_file_config_step.py

@@ -0,0 +1,124 @@
+import asyncio
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, Any, AsyncIterator
+
+from loguru import logger
+
+from bb_integrations_lib.gravitate.rita_api import GravitateRitaAPI
+from bb_integrations_lib.protocols.pipelines import GeneratorStep
+from bb_integrations_lib.provider.ftp.client import FTPIntegrationClient
+from bb_integrations_lib.secrets.credential_models import FTPCredential
+from bb_integrations_lib.provider.ftp.model import FTPType
+from bb_integrations_lib.shared.model import RawData, FileConfigRawData, ConfigMode, ConfigMatchMode
+from bb_integrations_lib.util.utils import check_if_file_greater_than_date, file_exact_match
+
+
+class SFTPFileConfigStep(GeneratorStep):
+
+    def __init__(self, rita_client: GravitateRitaAPI,
+                 ftp_client: FTPIntegrationClient | dict[str, FTPIntegrationClient], mode: ConfigMode,
+                 match_mode: ConfigMatchMode = ConfigMatchMode.Partial, bucket_name: str | None = None,
+                 config_name: str | None = None, min_date: datetime = datetime.min, strip_trailing_digits:bool=False, *args,
+                 **kwargs) -> None:
+        """
+        Imports SFTP files based on the provided or discovered fileconfigs.
+
+        :param rita_client: The RITA client to use to retrieve fileconfigs.
+        :param ftp_client: The FTP client, or a dict of FTP clients with keys matching confignames, to use to retrieve
+          data.
+        :param mode: How the step should discover fileconfigs.
+        :param match_mode: How the step should match fileconfigs to various properties of the files being scanned.
+        :param bucket_name: The bucket name which holds fileconfigs, for FromBucket and ByName modes.
+        :param config_name: The fileconfig name, if using ByName mode.
+        :param min_date: Filter out files with a date before this.
+        """
+        super().__init__(*args, **kwargs)
+        self.rita_client = rita_client
+        self.ftp_client = ftp_client
+        self.mode = mode
+
+        self.match_mode = match_mode
+        self.strip_trailing_digits = strip_trailing_digits
+        self.bucket_name = bucket_name
+        self.config_name = config_name
+        self.min_date = min_date
+        self.file_configs: dict[str, Any] = {}
+
+        if self.mode == ConfigMode.FromBucket and not self.bucket_name:
+            raise ValueError("Cannot use FromBucket mode without setting a bucket_name")
+        if self.mode == ConfigMode.ByName and not self.bucket_name:
+            raise ValueError("Cannot use ByName mode without setting a bucket_name")
+        if self.mode == ConfigMode.ByName and not self.config_name:
+            raise ValueError("Cannot use ByName mode without setting a config_name")
+
+    async def load_file_configs(self):
+        if self.mode == ConfigMode.AllFiltered:
+            self.file_configs = await self.rita_client.get_file_configs()
+        elif self.mode == ConfigMode.FromBucket:
+            self.file_configs = await self.rita_client.get_fileconfigs_from_bucket(self.bucket_name)
+        elif self.mode == ConfigMode.ByName:
+            self.file_configs = await self.rita_client.get_fileconfig_by_name(self.bucket_name, self.config_name)
+        logger.info(f"Loaded {len(self.config_name)} fileconfigs: {self.config_name}")
+
+    def describe(self) -> str:
+        return "Importing SFTP files based on file configs"
+
+    async def generator(self, i: Any) -> AsyncIterator[RawData]:
+        await self.load_file_configs()
+
+        for config_name, file_config in self.file_configs.items():
+            if isinstance(self.ftp_client, dict):
+                selected_ftp_client = self.ftp_client[config_name]
+            else:
+                selected_ftp_client = self.ftp_client
+            logger.info(f"Scanning with fileconfig '{config_name}' in directory {file_config.inbound_directory}")
+            file_names = list(selected_ftp_client.list_files(file_config.inbound_directory))
+            for idx, file_name in enumerate(file_names):
+                if self.match_mode == ConfigMatchMode.Exact:
+                    logger.info(f"Exact Matching file {file_name}")
+                    if not file_exact_match(file_name, file_config.file_name):
+                        logger.debug(f"Skipping file {file_name} due to not matching exactly to {file_config.file_name}")
+                        continue
+                elif self.match_mode == ConfigMatchMode.Partial:
+                    if not file_config.file_name in file_name:
+                        continue
+                elif self.match_mode == ConfigMatchMode.ByExtension:
+                    if not file_name.endswith(file_config.file_extension):
+                        continue
+                if file_config.date_format != "" and \
+                    not check_if_file_greater_than_date(file_name, file_config.file_name, file_config.date_format,
+                                                        self.min_date, self.strip_trailing_digits):
+                    logger.debug(f"Skipping file {file_name} due to having date > {self.min_date}")
+                    continue
+                logger.info(f"fetching {idx+1}/{len(file_names)}: {file_name}")
+                rd = selected_ftp_client.download_file(str(Path(file_config.inbound_directory) / Path(file_name)))
+                self.pipeline_context.file_config = file_config
+                yield FileConfigRawData(data=rd.data, file_name=rd.file_name, file_config=file_config)
+
+
+if __name__ == "__main__":
+    async def main():
+        s = SFTPFileConfigStep(
+            rita_client=GravitateRitaAPI(
+                base_url="",
+                username="",
+                password=""
+            ),
+            ftp_client=FTPIntegrationClient(
+                credentials=FTPCredential(
+                    host="",
+                    username="",
+                    password="",
+                    port=22,
+                    ftp_type=FTPType.sftp
+                ),
+            ),
+            mode=ConfigMode.ByName,
+            config_name="my_config"
+        )
+        async for r in s.generator(None):
+            print(r)
+
+    asyncio.run(main())

bb_integrations_lib/pipelines/steps/importing/test_exact_file_match.py

@@ -0,0 +1,57 @@
+import unittest
+import re
+
+
+def file_exact_match(file_name: str, sub_str: str, sep='_') -> bool:
+    if not file_name or not sub_str:
+        return False
+    if not file_name.startswith(sub_str):
+        return False
+    remainder = file_name[len(sub_str):]
+    if not remainder:
+        return True
+    if remainder.startswith('.'):
+        return True
+    if remainder.startswith(sep):
+        date_part = remainder[1:]
+        date_pattern = r'^(\d{4}-\d{2}-\d{2}|\d{2}-\d{2}-\d{4}|\d{8}|\d{4}_\d{2}_\d{2}|' + \
+                       r'\d{2}_\d{2}_\d{4}|\d{4}\d{2}\d{2})(\.[a-zA-Z0-9]+)?$'
+        return bool(re.match(date_pattern, date_part))
+
+    return False
+
+
+class TestFileExactMatch(unittest.TestCase):
+    def test_examples_from_requirements(self):
+        self.assertFalse(file_exact_match("Gross_contractrack.csv", "contractrack"))
+        self.assertTrue(file_exact_match("contractrack.csv", "contractrack"))
+        self.assertTrue(file_exact_match("contractrack_2025-04-16.csv", "contractrack"))
+        self.assertFalse(file_exact_match("Gross_contractrack_2025-04-16.csv", "contractrack"))
+
+    def test_valid_date_formats(self):
+        self.assertTrue(file_exact_match("contractrack_20250416.csv", "contractrack"))
+        self.assertTrue(file_exact_match("contractrack_2025-04-16.csv", "contractrack"))
+        self.assertTrue(file_exact_match("contractrack_2025_04_16.csv", "contractrack"))
+
+    def test_invalid_formats(self):
+        # Invalid because there's additional text after the date
+        self.assertFalse(file_exact_match("contractrack_2025-04-16_final.csv", "contractrack"))
+
+        # Invalid because there's text instead of a date
+        self.assertFalse(file_exact_match("contractrack_version2.csv", "contractrack"))
+
+        # Invalid because substring isn't at the start
+        self.assertFalse(file_exact_match("prefix_contractrack_2025-04-16.csv", "contractrack"))
+
+    def test_different_separator(self):
+        self.assertTrue(file_exact_match("contractrack-2025-04-16.csv", "contractrack", sep='-'))
+        self.assertFalse(file_exact_match("Gross-contractrack-2025-04-16.csv", "contractrack", sep='-'))
+
+    def test_edge_cases(self):
+        self.assertTrue(file_exact_match("contractrack", "contractrack"))
+        self.assertFalse(file_exact_match("contractrack_", "contractrack"))
+        self.assertFalse(file_exact_match("contractrack_notadate.csv", "contractrack"))
+
+
+if __name__ == "__main__":
+    unittest.main()

bb_integrations_lib/pipelines/steps/null_step.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+from bb_integrations_lib.protocols.pipelines import Step, Input, Output
+
+
+class NullStep(Step[Any, None]):
+    """A step that performs no action. May be useful as the first step in a job pipeline."""
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def describe(self) -> str:
+        return "Null Step"
+
+    async def execute(self, i: Any) -> None:
+        return None

bb_integrations_lib/pipelines/steps/pe_integration_job_step.py

@@ -0,0 +1,32 @@
+from typing import Dict, Any, Awaitable
+
+from bb_integrations_lib.gravitate.pe_api import GravitatePEAPI
+from bb_integrations_lib.mappers.prices.model import Action
+from bb_integrations_lib.protocols.pipelines import Step
+
+
+class PEIntegrationJobActionStep(Step):
+    def __init__(self, pe_client: GravitatePEAPI, action: Action, integration_name: str, source_system_id: str,
+                 *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.pe_client = pe_client
+        self.action = action
+        self.integration_name = integration_name
+        self.source_system_id = source_system_id
+
+    def describe(self) -> str:
+        return "Start, end or error pricing engine integration job"
+
+    async def execute(self, _: Any):
+        return await self.match_action()
+
+    async def match_action(self) -> Awaitable:
+        match self.action:
+            case Action.start:
+                return self.pe_client.integration_start
+            case Action.stop:
+                return self.pe_client.integration_stop
+            case Action.error:
+                return self.pe_client.integration_error
+            case _:
+                raise ValueError(f"Unexpected action {self.action}")

bb_integrations_lib/pipelines/steps/processing/archive_gcs_step.py

@@ -0,0 +1,76 @@
+import os
+import uuid
+
+import pandas as pd
+from gcloud.aio.storage import Storage
+from loguru import logger
+
+from bb_integrations_lib.models.pipeline_structs import BolExportResults
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.shared.model import RawData, FileConfigRawData
+
+
+class ArchiveGCSStep(Step):
+    def __init__(
+            self,
+            gcloud_storage: Storage,
+            bucket_path: str,
+            field_sep: str = ",",
+            content_type: str = "",
+            error_on_exists: bool = False,
+            *args,
+            **kwargs
+    ) -> None:
+        """
+        Archive a file (RawData or BolExportResults) to Google Cloud Storage.
+        :param gcloud_storage: A gcloud-aio-storage client.
+        :param bucket_path: The bucket and optional directory to upload the file to.
+          Example: "my-bucket" or "my-bucket/my-dir/my-dir-2". If a FileConfigRawData is passed to the step, the
+          source_system will be appended to the directory, before the file name. (bucket/prefix/source_system/file_name)
+        :param field_sep: Field separator when uploading BolExportResults.
+        :param content_type: Optional - explicitly specify the content type of the file.
+        :param error_on_exists: Whether to raise an exception if the file exists.
+        """
+        super().__init__(*args, **kwargs)
+
+        self.storage = gcloud_storage
+        self.bucket = bucket_path
+        self.prefix = ""
+        self.field_sep = field_sep
+        self.content_type = content_type
+        self.error_on_exists = error_on_exists
+
+        if "/" in bucket_path:
+            [self.bucket, self.prefix] = bucket_path.split("/", maxsplit=1)
+
+    def describe(self):
+        return "Archiving file in GCS"
+
+    async def execute(self, i: RawData | FileConfigRawData | BolExportResults) -> RawData:
+        def file_path():
+            if isinstance(i, FileConfigRawData):
+                return os.path.join(self.prefix, i.file_config.source_system, i.file_name)
+            else:
+                return os.path.join(self.prefix, i.file_name)
+
+        if isinstance(i, BolExportResults):
+            df = pd.DataFrame.from_records(i.orders)
+            csv_text = df.to_csv(index=False, sep=self.field_sep)
+            # Default to text/csv for BolExportResults if we don't have an explicit content type
+            self.content_type = "text/csv" if not self.content_type else self.content_type
+            contents = csv_text.encode("utf-8")
+        else:
+            contents = i.data
+
+        if await self.storage.get_bucket(self.bucket).blob_exists(file_path()):
+            if self.error_on_exists:
+                raise Exception(f"File '{self.bucket}/{file_path()}' already exists")
+            old_file_name = i.file_name
+            i.file_name = f"DUPLICATE_{uuid.uuid4()}_{i.file_name}"
+            logger.debug(
+                f"Blob named '{old_file_name}' already exists; archiving '{i.file_name}' to GCS (backup attempt)")
+        else:
+            logger.debug(f"Archiving '{i.file_name}' to GCS")
+        await self.storage.upload(self.bucket, file_path(), contents or '', content_type=self.content_type)
+
+        return i

bb_integrations_lib/pipelines/steps/processing/archive_sftp_step.py

@@ -0,0 +1,48 @@
+import os
+import uuid
+from typing import Any
+
+import loguru
+
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.provider.ftp.client import FTPIntegrationClient
+from bb_integrations_lib.shared.model import FileConfigRawData, RawData
+
+
+class ArchiveSFTPStep(Step):
+    def __init__(self, ftp_client: FTPIntegrationClient, src_directory: str | None = None,
+                 archive_directory: str | None = None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.ftp_client = ftp_client
+        self.src_directory = src_directory
+        self.archive_directory = archive_directory
+
+    def describe(self) -> str:
+        return "SFTP Rename Step"
+
+    async def execute(self, i: Any) -> Any:
+        if isinstance(i, FileConfigRawData):
+            old_name = os.path.join(i.file_config.inbound_directory, i.file_name)
+            new_name = os.path.join(i.file_config.archive_directory, i.file_name)
+        elif isinstance(i, RawData):
+            if self.src_directory is None or self.archive_directory is None:
+                raise RuntimeError("Attempted to archive with a RawData object but src_directory or archive_directory was not provided.")
+            old_name = os.path.join(self.src_directory, i.file_name)
+            new_name = os.path.join(self.archive_directory, i.file_name)
+        else:
+            raise NotImplementedError(f"Unsupported input type: {type(i)}")
+        try:
+            loguru.logger.debug(f"Archiving file {old_name} -> {new_name}")
+            self.ftp_client.rename_file(old_name, new_name)
+        except Exception as e:
+            try:
+                loguru.logger.debug(f"Archiving file failed...")
+                # this file may already exist. Give the file a randomized name and then archive.
+                new_name = new_name.replace(i.file_name, f"DUPLICATE_{uuid.uuid4()}_{i.file_name}")
+                loguru.logger.debug(f"Archiving file (backup attempt) {old_name} -> {new_name}")
+                self.ftp_client.rename_file(old_name, new_name)
+            except:
+                loguru.logger.warning(f"Archiving backup file failed. Deleting source file to prevent duplicate readings.")
+                self.ftp_client.delete_file(old_name)
+        loguru.logger.debug(f"Archived file {old_name} -> {new_name}")
+        return i