workbench 0.8.162__py3-none-any.whl → 0.8.220__py3-none-any.whl

workbench/core/cloud_platform/aws/aws_df_store.py

@@ -1,404 +0,0 @@
-"""AWSDFStore: Fast/efficient storage of DataFrames using AWS S3/Parquet/Snappy"""
-
-from datetime import datetime
-from typing import Union
-import logging
-import awswrangler as wr
-import pandas as pd
-import re
-from urllib.parse import urlparse
-
-# Workbench Imports
-from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
-from workbench.utils.config_manager import ConfigManager
-from workbench.utils.aws_utils import not_found_returns_none
-
-
-class AWSDFStore:
-    """AWSDFStore: Fast/efficient storage of DataFrames using AWS S3/Parquet/Snappy
-
-    Common Usage:
-        ```python
-        df_store = AWSDFStore()
-
-        # List Data
-        df_store.list()
-
-        # Add DataFrame
-        df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
-        df_store.upsert("/test/my_data", df)
-
-        # Retrieve DataFrame
-        df = df_store.get("/test/my_data")
-        print(df)
-
-        # Delete Data
-        df_store.delete("/test/my_data")
-        ```
-    """
-
-    def __init__(self, path_prefix: Union[str, None] = None):
-        """AWSDFStore Init Method
-
-        Args:
-            path_prefix (Union[str, None], optional): Path prefix for storage locations (Defaults to None)
-        """
-        self.log = logging.getLogger("workbench")
-        self._base_prefix = "df_store/"
-        self.path_prefix = self._base_prefix + path_prefix if path_prefix else self._base_prefix
-        self.path_prefix = re.sub(r"/+", "/", self.path_prefix)  # Collapse slashes
-
-        # Get the Workbench Bucket
-        config = ConfigManager()
-        self.workbench_bucket = config.get_config("WORKBENCH_BUCKET")
-
-        # Get the S3 Client
-        self.boto3_session = AWSAccountClamp().boto3_session
-        self.s3_client = self.boto3_session.client("s3")
-
-    def list(self, include_cache: bool = False) -> list:
-        """List all objects in the data_store prefix
-
-        Args:
-            include_cache (bool, optional): Include cache objects in the list (Defaults to False)
-
-        Returns:
-            list: A list of all the objects in the data_store prefix.
-        """
-        df = self.summary(include_cache=include_cache)
-        return df["location"].tolist()
-
-    def last_modified(self, location: str) -> Union[datetime, None]:
-        """Return the last modified date of a graph.
-
-        Args:
-            location (str): Logical location of the graph.
-
-        Returns:
-            Union[datetime, None]: Last modified datetime or None if not found.
-        """
-        s3_uri = self._generate_s3_uri(location)
-        bucket, key = self._parse_s3_uri(s3_uri)
-
-        try:
-            response = self.s3_client.head_object(Bucket=bucket, Key=key)
-            return response["LastModified"]
-        except self.s3_client.exceptions.ClientError:
-            return None
-
-    def summary(self, include_cache: bool = False) -> pd.DataFrame:
-        """Return a nicely formatted summary of object locations, sizes (in MB), and modified dates.
-
-        Args:
-            include_cache (bool, optional): Include cache objects in the summary (Defaults to False)
-        """
-        df = self.details(include_cache=include_cache)
-
-        # Create a formatted DataFrame
-        formatted_df = pd.DataFrame(
-            {
-                "location": df["location"],
-                "size (MB)": (df["size"] / (1024 * 1024)).round(2),  # Convert size to MB
-                "modified": pd.to_datetime(df["modified"]).dt.strftime("%Y-%m-%d %H:%M:%S"),  # Format date
-            }
-        )
-        return formatted_df
-
-    def details(self, include_cache: bool = False) -> pd.DataFrame:
-        """Return detailed metadata for all objects, optionally excluding the specified prefix.
-
-        Args:
-            include_cache (bool, optional): Include cache objects in the details (Defaults to False)
-        """
-        try:
-            response = self.s3_client.list_objects_v2(Bucket=self.workbench_bucket, Prefix=self.path_prefix)
-            if "Contents" not in response:
-                return pd.DataFrame(columns=["location", "s3_file", "size", "modified"])
-
-            # Collect details for each object
-            data = []
-            for obj in response["Contents"]:
-                full_key = obj["Key"]
-
-                # Reverse logic: Strip the bucket/prefix in the front and .parquet in the end
-                location = full_key.replace(f"{self.path_prefix}", "/").split(".parquet")[0]
-                s3_file = f"s3://{self.workbench_bucket}/{full_key}"
-                size = obj["Size"]
-                modified = obj["LastModified"]
-                data.append([location, s3_file, size, modified])
-
-            # Create the DataFrame
-            df = pd.DataFrame(data, columns=["location", "s3_file", "size", "modified"])
-
-            # Apply the exclude_prefix filter if set
-            cache_prefix = "/workbench/dataframe_cache/"
-            if not include_cache:
-                df = df[~df["location"].str.startswith(cache_prefix)]
-
-            return df
-
-        except Exception as e:
-            self.log.error(f"Failed to get object details: {e}")
-            return pd.DataFrame(columns=["location", "s3_file", "size", "created", "modified"])
-
-    def check(self, location: str) -> bool:
-        """Check if a DataFrame exists at the specified location
-
-        Args:
-            location (str): The location of the data to check.
-
-        Returns:
-            bool: True if the data exists, False otherwise.
-        """
-        # Generate the specific S3 prefix for the target location
-        s3_prefix = f"{self.path_prefix}/{location}.parquet/"
-        s3_prefix = re.sub(r"/+", "/", s3_prefix)  # Collapse slashes
-
-        # Use list_objects_v2 to check if any objects exist under this specific prefix
-        response = self.s3_client.list_objects_v2(Bucket=self.workbench_bucket, Prefix=s3_prefix, MaxKeys=1)
-        return "Contents" in response
-
-    @not_found_returns_none
-    def get(self, location: str) -> Union[pd.DataFrame, None]:
-        """Retrieve a DataFrame from AWS S3.
-
-        Args:
-            location (str): The location of the data to retrieve.
-
-        Returns:
-            pd.DataFrame: The retrieved DataFrame or None if not found.
-        """
-        s3_uri = self._generate_s3_uri(location)
-        return wr.s3.read_parquet(s3_uri)
-
-    def upsert(self, location: str, data: Union[pd.DataFrame, pd.Series]):
-        """Insert or update a DataFrame or Series in the AWS S3.
-
-        Args:
-            location (str): The location of the data.
-            data (Union[pd.DataFrame, pd.Series]): The data to be stored.
-        """
-        # Check if the data is a Pandas Series, convert it to a DataFrame
-        if isinstance(data, pd.Series):
-            data = data.to_frame()
-
-        # Ensure data is a DataFrame
-        if not isinstance(data, pd.DataFrame):
-            raise ValueError("Only Pandas DataFrame or Series objects are supported.")
-
-        # Convert object columns to string type to avoid PyArrow type inference issues.
-        data = self.type_convert_before_parquet(data)
-
-        # Update/Insert the DataFrame to S3
-        s3_uri = self._generate_s3_uri(location)
-        try:
-            wr.s3.to_parquet(df=data, path=s3_uri, dataset=True, mode="overwrite", index=True)
-            self.log.info(f"Dataframe cached {s3_uri}...")
-        except Exception as e:
-            self.log.error(f"Failed to cache dataframe '{s3_uri}': {e}")
-            raise
-
-    @staticmethod
-    def type_convert_before_parquet(df: pd.DataFrame) -> pd.DataFrame:
-        # Convert object columns to string type to avoid PyArrow type inference issues.
-        df = df.copy()
-        object_cols = df.select_dtypes(include=["object"]).columns
-        df[object_cols] = df[object_cols].astype("str")
-        return df
-
-    def delete(self, location: str):
-        """Delete a DataFrame from the AWS S3.
-
-        Args:
-            location (str): The location of the data to delete.
-        """
-        s3_uri = self._generate_s3_uri(location)
-
-        # Check if the folder (prefix) exists in S3
-        if not wr.s3.list_objects(s3_uri):
-            self.log.info(f"Data '{location}' does not exist in S3...")
-            return
-
-        # Delete the data from S3
-        try:
-            wr.s3.delete_objects(s3_uri)
-            self.log.info(f"Data '{location}' deleted successfully from S3.")
-        except Exception as e:
-            self.log.error(f"Failed to delete data '{location}': {e}")
-
-    def delete_recursive(self, location: str):
-        """Recursively delete all data under the specified location in AWS S3.
-
-        Args:
-            location (str): The location prefix of the data to delete.
-        """
-        # Construct the full prefix for S3
-        s3_prefix = re.sub(r"/+", "/", f"{self.path_prefix}/{location}")  # Collapse slashes
-        s3_prefix = s3_prefix.rstrip("/") + "/"  # Ensure the prefix ends with a slash
-
-        # List all objects under the given prefix
-        try:
-            response = self.s3_client.list_objects_v2(Bucket=self.workbench_bucket, Prefix=s3_prefix)
-            if "Contents" not in response:
-                self.log.info(f"No data found under '{s3_prefix}' to delete.")
-                return
-
-            # Gather all keys to delete
-            keys = [{"Key": obj["Key"]} for obj in response["Contents"]]
-            response = self.s3_client.delete_objects(Bucket=self.workbench_bucket, Delete={"Objects": keys})
-            for response in response.get("Deleted", []):
-                self.log.info(f"Deleted: {response['Key']}")
-
-        except Exception as e:
-            self.log.error(f"Failed to delete data recursively at '{location}': {e}")
-
-    def list_subfiles(self, prefix: str) -> list:
-        """Return a list of file locations with the given prefix.
-
-        Args:
-            prefix (str, optional): Only include files with the given prefix
-
-        Returns:
-            list: List of file locations (paths)
-        """
-        try:
-            full_prefix = f"{self.path_prefix}{prefix.lstrip('/')}"
-            response = self.s3_client.list_objects_v2(Bucket=self.workbench_bucket, Prefix=full_prefix)
-            if "Contents" not in response:
-                return []
-
-            locations = []
-            for obj in response["Contents"]:
-                full_key = obj["Key"]
-                location = full_key.replace(f"{self.path_prefix}", "/").split(".parquet")[0]
-                locations.append(location)
-            return locations
-
-        except Exception as e:
-            self.log.error(f"Failed to list subfiles: {e}")
-            return []
-
-    def _generate_s3_uri(self, location: str) -> str:
-        """Generate the S3 URI for the given location."""
-        s3_path = f"{self.workbench_bucket}/{self.path_prefix}/{location}.parquet"
-        return f"s3://{re.sub(r'/+', '/', s3_path)}"
-
-    def _parse_s3_uri(self, s3_uri: str) -> tuple:
-        """Parse an S3 URI into bucket and key."""
-        parsed = urlparse(s3_uri)
-        if parsed.scheme != "s3":
-            raise ValueError(f"Invalid S3 URI: {s3_uri}")
-        return parsed.netloc, parsed.path.lstrip("/")
-
-    def __repr__(self):
-        """Return a string representation of the AWSDFStore object."""
-        # Use the summary() method and format it to align columns for printing
-        summary_df = self.summary()
-
-        # Sanity check: If there are no objects, return a message
-        if summary_df.empty:
-            return "AWSDFStore: No data objects found in the store."
-
-        # Dynamically compute the max length of the 'location' column and add 5 spaces for padding
-        max_location_len = summary_df["location"].str.len().max() + 2
-        summary_df["location"] = summary_df["location"].str.ljust(max_location_len)
-
-        # Format the size column to include (MB) and ensure 3 spaces between size and date
-        summary_df["size (MB)"] = summary_df["size (MB)"].apply(lambda x: f"{x:.2f} MB")
-
-        # Enclose the modified date in parentheses and ensure 3 spaces between size and date
-        summary_df["modified"] = summary_df["modified"].apply(lambda x: f" ({x})")
-
-        # Convert the DataFrame to a string, remove headers, and return
-        return summary_df.to_string(index=False, header=False)
-
-
-if __name__ == "__main__":
-    """Exercise the AWSDFStore Class"""
-    import time
-
-    # Create a AWSDFStore manager
-    df_store = AWSDFStore()
-
-    # Details of the Dataframe Store
-    print("Detailed Data...")
-    print(df_store.details())
-
-    # List all objects
-    print("List Data...")
-    print(df_store.list())
-
-    # Add a new DataFrame
-    my_df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
-    df_store.upsert("/testing/test_data", my_df)
-
-    # Check the last modified date
-    print("Last Modified Date:")
-    print(df_store.last_modified("/testing/test_data"))
-
-    # Get the DataFrame
-    print(f"Getting data 'test_data':\n{df_store.get('/testing/test_data')}")
-
-    # Now let's test adding a Series
-    series = pd.Series([1, 2, 3, 4], name="Series")
-    df_store.upsert("/testing/test_series", series)
-    print(f"Getting data 'test_series':\n{df_store.get('/testing/test_series')}")
-
-    # Summary of the data
-    print("Summary Data...")
-    print(df_store.summary())
-
-    # Repr of the AWSDFStore object
-    print("AWSDFStore Object:")
-    print(df_store)
-
-    # Check if the data exists
-    print("Check if data exists...")
-    print(df_store.check("/testing/test_data"))
-    print(df_store.check("/testing/test_series"))
-
-    # Time the check
-    start_time = time.time()
-    print(df_store.check("/testing/test_data"))
-    print("--- Check %s seconds ---" % (time.time() - start_time))
-
-    # Test list_subfiles
-    print("List Subfiles:")
-    print(df_store.list_subfiles("/testing"))
-
-    # Now delete the test data
-    df_store.delete("/testing/test_data")
-    df_store.delete("/testing/test_series")
-
-    # Check if the data exists
-    print("Check if data exists...")
-    print(df_store.check("/testing/test_data"))
-    print(df_store.check("/testing/test_series"))
-
-    # Add a bunch of dataframes and then test recursive delete
-    for i in range(10):
-        df_store.upsert(f"/testing/data_{i}", pd.DataFrame({"A": [1, 2], "B": [3, 4]}))
-    print("Before Recursive Delete:")
-    print(df_store.summary())
-    df_store.delete_recursive("/testing")
-    print("After Recursive Delete:")
-    print(df_store.summary())
-
-    # Get a non-existent DataFrame
-    print("Getting non-existent data...")
-    print(df_store.get("/testing/no_where"))
-
-    # Test path_prefix
-    df_store = AWSDFStore(path_prefix="/super/test")
-    print(df_store.path_prefix)
-    df_store.upsert("test_data", my_df)
-    print(df_store.get("test_data"))
-    print(df_store.summary())
-    df_store.delete("test_data")
-    print(df_store.summary())
-
-    # Test columns with Spaces in them
-    my_df = pd.DataFrame({"My A": [1, 2], "My B": [3, 4]})
-    df_store.upsert("/testing/test_data", my_df)
-    my_df = df_store.get("/testing/test_data")
-    print(my_df)

workbench/core/cloud_platform/aws/aws_parameter_store.py

@@ -1,280 +0,0 @@
-"""AWSParameterStore: Manages Workbench parameters in AWS Systems Manager Parameter Store."""
-
-from typing import Union
-import logging
-import json
-import zlib
-import base64
-from botocore.exceptions import ClientError
-
-# Workbench Imports
-from workbench.core.cloud_platform.aws.aws_session import AWSSession
-from workbench.utils.json_utils import CustomEncoder
-
-
-class AWSParameterStore:
-    """AWSParameterStore: Manages Workbench parameters in AWS Systems Manager Parameter Store.
-
-    Common Usage:
-        ```python
-        params = AWSParameterStore()
-
-        # List Parameters
-        params.list()
-
-        ['/workbench/abalone_info',
-         '/workbench/my_data',
-         '/workbench/test',
-         '/workbench/pipelines/my_pipeline']
-
-        # Add Key
-        params.upsert("key", "value")
-        value = params.get("key")
-
-        # Add any data (lists, dictionaries, etc..)
-        my_data = {"key": "value", "number": 4.2, "list": [1,2,3]}
-        params.upsert("my_data", my_data)
-
-        # Retrieve data
-        return_value = params.get("my_data")
-        pprint(return_value)
-
-        {'key': 'value', 'list': [1, 2, 3], 'number': 4.2}
-
-        # Delete parameters
-        param_store.delete("my_data")
-        ```
-    """
-
-    def __init__(self):
-        """AWSParameterStore Init Method"""
-        self.log = logging.getLogger("workbench")
-
-        # Initialize a Workbench Session (to assume the Workbench ExecutionRole)
-        self.boto3_session = AWSSession().boto3_session
-
-        # Create a Systems Manager (SSM) client for Parameter Store operations
-        self.ssm_client = self.boto3_session.client("ssm")
-
-    def list(self, prefix: str = None) -> list:
-        """List all parameters in the AWS Parameter Store, optionally filtering by a prefix.
-
-        Args:
-            prefix (str, optional): A prefix to filter the parameters by. Defaults to None.
-
-        Returns:
-            list: A list of parameter names and details.
-        """
-        try:
-            # Set up parameters for the query
-            params = {"MaxResults": 50}
-
-            # If a prefix is provided, add the 'ParameterFilters' for optimization
-            if prefix:
-                params["ParameterFilters"] = [{"Key": "Name", "Option": "BeginsWith", "Values": [prefix]}]
-
-            # Initialize the list to collect parameter names
-            all_parameters = []
-
-            # Make the initial call to describe parameters
-            response = self.ssm_client.describe_parameters(**params)
-
-            # Aggregate the names from the initial response
-            all_parameters.extend(param["Name"] for param in response["Parameters"])
-
-            # Continue to paginate if there's a NextToken
-            while "NextToken" in response:
-                # Update the parameters with the NextToken for subsequent calls
-                params["NextToken"] = response["NextToken"]
-                response = self.ssm_client.describe_parameters(**params)
-
-                # Aggregate the names from the subsequent responses
-                all_parameters.extend(param["Name"] for param in response["Parameters"])
-
-        except Exception as e:
-            self.log.error(f"Failed to list parameters: {e}")
-            return []
-
-        # Return the aggregated list of parameter names
-        return all_parameters
-
-    def get(self, name: str, warn: bool = True, decrypt: bool = True) -> Union[str, list, dict, None]:
-        """Retrieve a parameter value from the AWS Parameter Store.
-
-        Args:
-            name (str): The name of the parameter to retrieve.
-            warn (bool): Whether to log a warning if the parameter is not found.
-            decrypt (bool): Whether to decrypt secure string parameters.
-
-        Returns:
-            Union[str, list, dict, None]: The value of the parameter or None if not found.
-        """
-        try:
-            # Retrieve the parameter from Parameter Store
-            response = self.ssm_client.get_parameter(Name=name, WithDecryption=decrypt)
-            value = response["Parameter"]["Value"]
-
-            # Auto-detect and decompress if needed
-            if value.startswith("COMPRESSED:"):
-                # Base64 decode and decompress
-                self.log.important(f"Decompressing parameter '{name}'...")
-                compressed_value = base64.b64decode(value[len("COMPRESSED:") :])
-                value = zlib.decompress(compressed_value).decode("utf-8")
-
-            # Attempt to parse the value back to its original type
-            try:
-                parsed_value = json.loads(value)
-                return parsed_value
-            except (json.JSONDecodeError, TypeError):
-                # If parsing fails, return the value as is "hope for the best"
-                return value
-
-        except ClientError as e:
-            if e.response["Error"]["Code"] == "ParameterNotFound":
-                if warn:
-                    self.log.warning(f"Parameter '{name}' not found")
-            else:
-                self.log.error(f"Failed to get parameter '{name}': {e}")
-            return None
-
-    def upsert(self, name: str, value, precision: int = 3):
-        """Insert or update a parameter in the AWS Parameter Store.
-
-        Args:
-            name (str): The name of the parameter.
-            value (str | list | dict): The value of the parameter.
-            precision (int): The precision for float values in the JSON encoding.
-        """
-        try:
-            # Convert to JSON and check if compression is needed
-            json_value = json.dumps(value, cls=CustomEncoder, precision=precision)
-            if len(json_value) <= 4096:
-                # Store normally if under 4KB
-                self._store_parameter(name, json_value)
-                return
-
-            # Need compression - log warning
-            self.log.important(
-                f"Parameter {name} exceeds 4KB ({len(json_value)} bytes): compressing and reducing precision..."
-            )
-
-            # Try compression with precision reduction
-            compressed_value = self._compress_value(value)
-
-            if len(compressed_value) <= 4096:
-                self._store_parameter(name, compressed_value)
-                return
-
-            # Try clipping the data
-            clipped_value = self._clip_data(value)
-            compressed_clipped = self._compress_value(clipped_value)
-
-            if len(compressed_clipped) <= 4096:
-                self.log.warning(
-                    f"Parameter {name} data clipped to 100 items/elements: ({len(compressed_clipped)} bytes)"
-                )
-                self._store_parameter(name, compressed_clipped)
-                return
-
-            # Still too large - give up
-            self._handle_oversized_data(name, len(compressed_clipped))
-
-        except Exception as e:
-            self.log.critical(f"Failed to add/update parameter '{name}': {e}")
-            raise
-
-    @staticmethod
-    def _compress_value(value) -> str:
-        """Compress a value with precision reduction."""
-        json_value = json.dumps(value, cls=CustomEncoder, precision=3)
-        compressed = zlib.compress(json_value.encode("utf-8"), level=9)
-        return "COMPRESSED:" + base64.b64encode(compressed).decode("utf-8")
-
-    @staticmethod
-    def _clip_data(value):
-        """Clip data to reduce size, clip to first 100 items/elements."""
-        if isinstance(value, dict):
-            return dict(list(value.items())[:100])
-        elif isinstance(value, list):
-            return value[:100]
-        return value
-
-    def _store_parameter(self, name: str, value: str):
-        """Store parameter in AWS Parameter Store."""
-        self.ssm_client.put_parameter(Name=name, Value=value, Type="String", Overwrite=True)
-        self.log.info(f"Parameter '{name}' added/updated successfully.")
-
-    def _handle_oversized_data(self, name: str, size: int):
-        """Handle data that's too large even after compression and clipping."""
-        doc_link = "https://supercowpowers.github.io/workbench/api_classes/df_store"
-        self.log.error(f"Compressed size {size} bytes, cannot store > 4KB")
-        self.log.error(f"For larger data use the DFStore() class ({doc_link})")
-
-    def delete(self, name: str):
-        """Delete a parameter from the AWS Parameter Store.
-
-        Args:
-            name (str): The name of the parameter to delete.
-        """
-        try:
-            # Delete the parameter from Parameter Store
-            self.ssm_client.delete_parameter(Name=name)
-            self.log.info(f"Parameter '{name}' deleted successfully.")
-        except Exception as e:
-            self.log.error(f"Failed to delete parameter '{name}': {e}")
-
-    def delete_recursive(self, prefix: str):
-        """Delete all parameters with a given prefix from the AWS Parameter Store.
-
-        Args:
-            prefix (str): The prefix of the parameters to delete.
-        """
-        # List all parameters with the given prefix
-        parameters = self.list(prefix=prefix)
-        for param in parameters:
-            self.delete(param)
-
-    def __repr__(self):
-        """Return a string representation of the AWSParameterStore object."""
-        return "\n".join(self.list())
-
-
-if __name__ == "__main__":
-    """Exercise the AWSParameterStore Class"""
-
-    # Create a AWSParameterStore manager
-    param_store = AWSParameterStore()
-
-    # List the parameters
-    print("Listing Parameters...")
-    print(param_store.list())
-
-    # Add a new parameter
-    param_store.upsert("/workbench/test", "value")
-
-    # Get the parameter
-    print(f"Getting parameter 'test': {param_store.get('/workbench/test')}")
-
-    # Add a dictionary as a parameter
-    sample_dict = {"key": "str_value", "awesome_value": 4.2}
-    param_store.upsert("/workbench/my_data", sample_dict)
-
-    # Retrieve the parameter as a dictionary
-    retrieved_value = param_store.get("/workbench/my_data")
-    print("Retrieved value:", retrieved_value)
-
-    # List the parameters
-    print("Listing Parameters...")
-    print(param_store.list())
-
-    # List the parameters with a prefix
-    print("Listing Parameters with prefix '/workbench':")
-    print(param_store.list("/workbench"))
-
-    # Delete the parameters
-    param_store.delete("/workbench/test")
-    param_store.delete("/workbench/my_data")
-
-    # Out of scope tests
-    param_store.upsert("test", "value")
-    param_store.delete("test")