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

bb_integrations_lib/provider/api/warren_rogers/client.py

@@ -0,0 +1,207 @@
+import asyncio
+from datetime import datetime, timedelta
+from pprint import pprint
+from typing import Optional
+
+import httpx
+from loguru import logger
+
+from bb_integrations_lib.protocols.flat_file import TankReading, TankMonitorType
+
+
+class WarrenRogersClient(httpx.AsyncClient):
+    def __init__(
+            self,
+            client_id: str,
+            client_secret: str,
+            api_key: str,
+            company_id: str,
+            token_endpoint: str,
+            base_url: str,
+            timeout: float = 180.0
+    ):
+        super().__init__(base_url=base_url, timeout=timeout)
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.company_id = company_id
+        self.token_endpoint = token_endpoint
+        self.page_size = 100
+        self.oauth_token = None
+        # Expect to refresh token 30s before actual expiry, in case of time desync between us and server
+        self.pre_expiry_refresh = timedelta(seconds=30)
+
+    def __repr__(self):
+        return "Warren Rogers API Client"
+
+    @property
+    async def access_token(self):
+        await self.ensure_token()
+        return self.oauth_token["access_token"]
+
+    async def _refresh_token(self):
+        now = datetime.now()
+        res = await self.post(self.token_endpoint,
+                              auth=(self.client_id, self.client_secret),
+                              headers={"Content-Type": "application/x-www-form-urlencoded"},
+                              params="grant_type=client_credentials")
+        res.raise_for_status()
+        self.oauth_token = res.json()
+        self.oauth_token["expires_at"] = now + timedelta(seconds=self.oauth_token["expires_in"])
+
+    def is_token_expired(self):
+        return datetime.now() > self.oauth_token["expires_at"] + self.pre_expiry_refresh
+
+    async def ensure_token(self):
+        # First run or expired token? Then refresh
+        if not self.oauth_token or self.is_token_expired():
+            await self._refresh_token()
+
+    def _build_url(self, endpoint: str, api_ver: str = "v1") -> str:
+        return f"api/{api_ver}/companies/{self.company_id}/{endpoint}"
+
+    async def _get(self, url: str, params=None):
+        headers = {
+            "Accept": "application/json",
+            "x-api-key": self.api_key,
+            "wra-auth": await self.access_token
+        }
+        return await self.get(url, headers=headers, params=params)
+
+    async def _paginated_get(self, url: str):
+        pages = []
+        page_no = 1
+
+        first_page = (await self._get(url)).json()
+        total_results = first_page["totalResults"]
+        pages.append(first_page)
+        entries_retrieved = first_page["size"]
+
+        while entries_retrieved < total_results:
+            page_no += 1
+            page = (await self._get(url, params={"page": page_no})).json()
+            pages.append(page)
+            entries_retrieved += page["size"]
+
+        return pages
+
+    async def _continuation_get(self, url: str, starting_token: Optional[str] = None) -> tuple[str, list[dict]]:
+        """
+        Get data from a paginated stream that uses continuation tokens, like the /deliveries endpoint.
+        If a starting continuation token is provided, it will be passed to the API; otherwise, the default API behavior
+        will be used. In /deliveries case, this is retrieving the last hour of data, for example.
+
+        :param url: The URL of the endpoint.
+        :param starting_token: An optional continuation token to resume the stream from.
+        :return: A tuple; item 1 is the continuation token from the API, and item 2 is the list of pages of the API
+            response.
+        """
+        pages = []
+        page_n = 1
+        starting = True
+        con_token = starting_token
+        while starting or con_token:
+            logger.debug(f"Getting page {page_n}, continuation token {con_token}")
+            resp = await self._get(url=url, params={"continuationToken": con_token} if con_token else {})
+            resp.raise_for_status()
+            data = resp.json()
+            con_token = data.get("continuationToken")
+            pages.append(resp.json())
+            if not data.get("truncated"):
+                break
+            starting = False
+            page_n += 1
+        return con_token, pages
+
+    async def get_inventory_levels(self) -> list[dict]:
+        """
+        Get "inventory levels" (tank readings) for all tanks. Does not parse any values, but deserializes the json to a
+        dict.
+        """
+        raw = await self._paginated_get(self._build_url("inventory-levels"))
+        # The data this endpoint provides is a list of pages. Each page has some metadata plus an "inventoryLevels"
+        # field, which has the actual locations and their ATGs and readings. Flatten to one big list of inventoryLevels.
+        return [x for xs in raw for x in xs["inventoryLevels"]]
+
+    async def get_inventory_levels_at_loc(self, location: str) -> list[dict]:
+        """Like get_inventory_levels, but for a specific location."""
+        raw = await self._get(self._build_url(f"locations/{location}/inventory-levels"))
+        return raw.json()
+
+    def _parse_tankreadings(self, location: dict) -> list[TankReading]:
+        """Parse inventory levels from the WR API and convert them to TankReadings."""
+        trs: list[TankReading] = []
+        # For each ATG at location
+        for atg_level in location["atgLevels"]:
+            # For each of the actual inventory readings on this ATG
+            ils = atg_level["tankInventoryLevels"]
+            for il in ils:
+                # Specific inventory level reading - i.e., tank
+                # Note that this does not take the units reported by the WR API into account - this appears to be
+                # gallons, so far.
+                trs.append(
+                    TankReading(
+                        date=il["readingDateTime"],
+                        payload={},
+                        store=location["locationUniqueId"],
+                        tank=il["tankUniqueId"],
+                        timezone=None,  # Not needed, ISO date str has timezone
+                        volume=il["atgGrossVolume"],
+                        monitor_type=TankMonitorType.bbd
+                    )
+                )
+        return trs
+
+    async def get_unmapped_tank_readings(self) -> list[TankReading]:
+        """
+        Get all tank readings available to the client, like get_inventory_levels, but parse the tank readings with
+        _parse_tankreadings, returning them as a flat list. Note that the store and tank identifiers are returned
+        straight from the API, unmapped.
+        """
+        raw = await self._paginated_get(self._build_url("inventory-levels"))
+        # The data this endpoint provides is a list of pages. Each page has some metadata plus an "inventoryLevels"
+        # field, which has the actual locations and their ATGs and readings. Flatten to one big list of inventoryLevels.
+        inventory_levels = [x for xs in raw for x in xs["inventoryLevels"]]
+        trs = []
+        # For each store
+        for location in inventory_levels:
+            trs.extend(self._parse_tankreadings(location))
+        return trs
+
+    async def get_unmapped_tank_readings_loc(self, location: str):
+        """Like get_unmapped_tank_readings, but for a specific location."""
+        raw = await self._get(self._build_url(f"locations/{location}/inventory-levels"))
+        return self._parse_tankreadings(raw.json())
+
+    async def get_deliveries(self, starting_token: Optional[str] = None) -> tuple[str, list[dict]]:
+        """
+        Get a list of locations and their deliveries. If starting_token is provided, resume the data stream from that
+        continuation token. May return duplicates / repeated data.
+
+        :param starting_token: An optional continuation token to resume the stream from. If not provided, gets the last
+            hour of deliveries.
+        :return: A tuple; item 1 is the final continuation token from the API (for resuming the stream in the future),
+            and item 2 is the list of locations and their deliveries.
+        """
+        next_ct, raw = await self._continuation_get(
+            self._build_url("deliveries", api_ver="v3"),
+            starting_token=starting_token)
+        # Flatten the locationDeliveries lists from each page into one list of dicts of locations and their deliveries.
+        return next_ct, [x for xs in raw for x in xs["locationDeliveries"]]
+
+
+if __name__ == "__main__":
+    async def main():
+        client = WarrenRogersClient(
+            client_id="",
+            client_secret="",
+            api_key="",
+            company_id="LOVES_TRAVEL_STOPS_AND_COUNTRY_STORES",
+            token_endpoint="https://auth.api.wr-cloud.com/oauth2/token",
+            base_url="https://api.wr-cloud.com")
+
+        inventory_levels = await client.get_inventory_levels()
+        pprint(inventory_levels)
+
+    asyncio.run(main())

bb_integrations_lib/provider/aws/s3/client.py

@@ -0,0 +1,126 @@
+import json
+from dataclasses import dataclass
+from typing import Iterable
+from json import JSONDecodeError
+import boto3
+from bb_integrations_lib.shared.model import RawData
+from loguru import logger
+
+
+@dataclass
+class FileData:
+    """
+    Data class representing a file and its contents.
+
+    Attributes:
+        file_name (str): The name of the file.
+        data (dict): The content of the file parsed as a dictionary.
+    """
+    file_name: str
+    data: dict
+
+
+class S3Client:
+    """
+    Client for interacting with an AWS S3 bucket, providing methods for fetching,
+    archiving, and processing files.
+
+    Attributes:
+        bucket_name (str): Name of the AWS S3 bucket.
+        access_key_id (str): AWS access key ID for authentication.
+        secret_access_key (str): AWS secret access key for authentication.
+        archive_dir (str): Directory in the bucket where files are archived.
+        file_prefix (list[str]): List of file prefixes used for filtering files in the bucket.
+        files_per_chunk (int): Maximum number of files to process in a single iteration.
+        s3 (boto3.resource): Boto3 resource for S3 operations.
+        bucket (boto3.Bucket): Boto3 bucket object for interacting with the specified S3 bucket.
+    """
+
+    def __init__(
+        self,
+        *,
+        aws_bucket_name: str,
+        aws_access_key_id: str,
+        aws_secret_access_key: str,
+        aws_archive_dir: str,
+        aws_file_prefix: list[str],
+        files_per_chunk: int = 500,
+    ):
+        """
+        Initialize the S3Client with AWS credentials and configuration.
+
+        Args:
+            aws_bucket_name (str): Name of the S3 bucket to interact with.
+            aws_access_key_id (str): AWS access key ID.
+            aws_secret_access_key (str): AWS secret access key.
+            aws_archive_dir (str): Directory in the S3 bucket for archiving files.
+            aws_file_prefix (list[str]): List of file prefixes for filtering files.
+            files_per_chunk (int, optional): Number of files to fetch per query. Defaults to 500.
+        """
+        self.bucket_name = aws_bucket_name
+        self.access_key_id = aws_access_key_id
+        self.secret_access_key = aws_secret_access_key
+        self.archive_dir = aws_archive_dir
+        self.file_prefix = aws_file_prefix
+        self.files_per_chunk = files_per_chunk
+
+        self.s3 = boto3.resource(
+            "s3",
+            aws_access_key_id=self.access_key_id,
+            aws_secret_access_key=self.secret_access_key,
+        )
+        self.bucket = self.s3.Bucket(name=self.bucket_name)
+
+    def get_raw_data(self) -> Iterable[RawData]:
+        """
+        Fetch raw data from the S3 bucket.
+
+        Iterates through files matching the specified prefixes, loads their contents,
+        and returns them as a collection of `RawData` objects. Malformed files are
+        archived and deleted from the bucket.
+
+        Returns:
+            Iterable[RawData]: An iterable containing raw data from the files.
+        """
+        ret = []
+        for file in self._files():
+            try:
+                ret.append(
+                    RawData(file_name=file.key, data=(json.load(file.get()["Body"])))
+                )
+            except JSONDecodeError:
+                logger.error(f"Archiving bad file {file.key}")
+                copy_source = {"Bucket": self.bucket_name, "Key": file.key}
+                destination = f"{self.archive_dir}/{file.key}"
+                self.bucket.copy(copy_source, destination)
+                self.s3.Object(self.bucket_name, file.key).delete()
+        return ret
+
+    def archive_data(self, raw_data: RawData):
+        """
+        Archive a processed file in the designated archive directory of the S3 bucket.
+
+        Args:
+            raw_data (RawData): The raw data object representing the file to be archived.
+        """
+        copy_source = {"Bucket": self.bucket_name, "Key": raw_data.file_name}
+        destination = f"{self.archive_dir}/{raw_data.file_name}"
+        self.bucket.copy(copy_source, destination)
+        self.s3.Object(self.bucket_name, raw_data.file_name).delete()
+
+    def _files(self) -> Iterable:
+        """
+        Generate an iterable of file objects from the S3 bucket.
+
+        Queries the S3 bucket for files matching the configured prefixes, yielding
+        them in chunks defined by `files_per_chunk`.
+
+        Returns:
+            Iterable: An iterable of S3 file objects.
+        """
+        file_queries = [
+            self.bucket.objects.filter(Prefix=prefix) for prefix in self.file_prefix
+        ]
+        return (
+            obj for query in file_queries for obj in query.limit(self.files_per_chunk)
+        )

bb_integrations_lib/provider/ftp/client.py

@@ -0,0 +1,140 @@
+from typing import Iterable, Callable
+
+from bb_integrations_lib.provider.ftp.interface import FTPClient, SFTPClient, FTPClientInterface
+from bb_integrations_lib.provider.ftp.model import FTPFileInfo, FTPType
+from bb_integrations_lib.secrets.credential_models import FTPCredential
+from bb_integrations_lib.shared.model import RawData, File
+from bb_integrations_lib.util.utils import load_credentials
+
+
+class FTPIntegrationClient:
+    """
+    A client for interacting with an FTP or SFTP server, offering methods for listing, uploading, downloading,
+    renaming, and deleting files.
+    Gracefully manages the FTP connection for you, handling timeouts and retries.
+
+    Note that most commands are simply sent to the server, not double-checked - for example, a delete may not report an
+    error even if the file doesn't exist.
+    """
+
+    def __init__(self, credentials: FTPCredential):
+        self.credentials = credentials
+        self.interface: FTPClientInterface
+        match self.credentials.ftp_type:
+            case FTPType.ftp:
+                self.interface = FTPClient(self.credentials)
+            case FTPType.ftps:
+                self.interface = FTPClient(self.credentials)
+            case FTPType.sftp:
+                self.interface = SFTPClient(self.credentials)
+            case FTPType.ftpes:
+                self.interface = FTPClient(self.credentials)
+            case _:
+                raise Exception(f"Unknown FTP credential type {self.credentials.ftp_type}, please implement an adapter")
+        self.cur_dir = "/"
+
+    def __enter__(self):
+        self.interface.connect()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.interface.disconnect()
+
+    def list_files(self, directory: str) -> Iterable[str]:
+        """
+        List all files in the given directory.
+        :returns: An iterable of file names in the directory.
+        """
+        return self.interface.list_files(directory)
+
+    def rename_file(self, old_name: str, new_name: str) -> None:
+        """
+        Rename a single file.
+        :param old_name: The current name of the file.
+        :param new_name: The name to change to.
+        """
+        return self.rename_files([(old_name, new_name)])
+
+    def rename_files(self, files: Iterable[tuple[str, str]]) -> None:
+        """
+        Rename a batch of files.
+        :param files: An iterable of tuples, each containing the old/new names for each files like (old, new).
+        """
+        return self.interface.rename_files(files)
+
+    def delete_file(self, path: str) -> None:
+        """
+        Delete a single file.
+        :param path: The path to the file to delete.
+        """
+        return self.delete_files([path])
+
+    def delete_files(self, paths: Iterable[str]) -> None:
+        """
+        Delete a batch of files.
+        :param paths: The paths to files to delete.
+        """
+        return self.interface.delete_files(paths)
+
+    def upload_file(self, file: File, path: str) -> None:
+        """
+        Upload a single File to a given parent directory.
+        :param file: The file to upload. Needs at least ``file_name`` and data of a type that ``file.to_bytes()``
+          supports.
+        :param path: The parent directory to upload the file to.
+        """
+        self.upload_files([file], path)
+
+    def upload_files(self, files: Iterable[File], path: str) -> None:
+        """
+        Like ``upload_file``, but for a batch of files.
+        :param files: An iterable of File objects to upload. See docs on ``upload_file`` for how these are handled.
+        :param path: The parent directory to upload the files to.
+        """
+        self.interface.upload_files(files, path)
+
+    def download_file(self, path: str) -> RawData:
+        """
+        Download a single file into memory.
+        :param path: The remote path of the file to download.
+        :return: A RawData object representing the downloaded file. ``file_name`` and ``data`` attributes will be set.
+          ``data`` will be a BytesIO object pre-seeked to 0.
+        :raises FileNotFoundError: If the file doesn't exist.
+        """
+        try:
+            return next(iter(self.download_files([path])))
+        except StopIteration:
+            raise FileNotFoundError(f"Remote file {path} not found")
+
+    def download_files(self, paths: Iterable[str]) -> Iterable[RawData]:
+        """
+        Like ``download_file``, but for a batch of files. Not all downloads may succeed; if any fail they will not be
+        present in the result list. If none succeed, an empty list will be returned.
+        :param paths: The remote paths of the files to download.
+        :return:
+        """
+        return self.interface.download_files(paths)
+
+    def get_file_info(self, path: str) -> FTPFileInfo:
+        """
+        Get details about a single file.
+        :param path: The remote path of the file to get info for.
+        :return: A FTPFileInfo object with details about the file. If this is an FTP connection, only the ``size``
+          attribute is guaranteed. If this is an SFTP connection, all attributes but the permissions are guaranteed.
+        """
+        return self.interface.get_file_info(path)
+
+    def download_dir(self, directory: str, filt: Callable[[str], bool] = None) -> Iterable[RawData]:
+        """
+        Download files from an entire directory, optionally filtering with the ``filt`` callable.
+
+        :param directory: The directory to download.
+        :param filt: An optional filter callable. Should take one argument (the file name) and return a boolean
+          indicating whether to include the file..
+
+        :returns: An iterable of `RawData` objects representing the downloaded files.
+        """
+        file_list = self.list_files(directory)
+        if filt:
+            file_list = filter(filt, file_list)
+        return self.download_files(map(lambda file_name: f"{directory}/{file_name}", file_list))