data-factory-utils 1.0.0__py3-none-any.whl

data_factory_utils/__init__.py

@@ -0,0 +1 @@
+"""Init file."""

data_factory_utils/environment.py

@@ -0,0 +1,334 @@
+"""Class to help with AWS environment inference and configuration."""
+
+import logging
+import os
+import secrets
+import string
+from pathlib import Path
+from typing import Self
+
+import boto3
+import botocore.exceptions
+from cloudpathlib import S3Path
+from mypy_boto3_s3.type_defs import BucketTypeDef
+from mypy_boto3_sts.type_defs import CredentialsTypeDef
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+alphabet = string.ascii_lowercase
+
+
+def _raise_missing_token() -> None:
+    msg = "AWS_WEB_IDENTITY_TOKEN_FILE environment variable not set."
+    logger.exception(msg)
+    raise RuntimeError(msg)
+
+
+def _raise_missing_role_arn() -> None:
+    raise MissingEnvVarRoleArnError
+
+
+class MissingEnvVarRoleArnError(Exception):
+    """Exception for a missing role arn env var."""
+
+    def __init__(self) -> None:
+        """Raise error with message."""
+        self.message = "Add AWS_ROLE_ARN to environment variables."
+        super().__init__()
+
+
+class Environment:
+    """AWS environment inference and configuration.
+
+    This class helps determine the environment (prod, preprod, test, dev),
+    manage AWS credentials (via web identity or default),
+    and construct environment-specific S3 bucket URLs.
+    """
+
+    _instance: Self | None = None
+
+    def __new__(cls, *_args: object, **_kwargs: object) -> Self:
+        """Singleton instantiation."""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(
+        self,
+        job_name: str | None = "",
+        bucket_prefix: str | None = "emds",
+        *,
+        use_web_identity: bool,
+    ) -> None:
+        """Initialize the environment context.
+
+        Parameters
+        ----------
+        job_name : str, optional
+            Name used when assuming a role session, by default empty string.
+        bucket_prefix : str, optional
+            Base prefix used to identify datahub buckets, by default "emds".
+        use_web_identity : bool, optional
+            Whether to try using web identity credentials first.
+
+        """
+        if hasattr(self, "_initialized") and self._initialized:
+            return
+
+        self.job_name = job_name
+        self.bucket_prefix = bucket_prefix
+        self.use_web_identity = use_web_identity
+
+        self.session = self._init_session()
+        self.alias = self._fetch_account_alias()
+        self.account_no = self._get_account_number()
+        self.bucket_list = self._list_buckets()
+
+        self._initialized: bool = True
+
+    def _init_session(self) -> boto3.Session:
+        """Initialise a boto3 session, optionally using web identity credentials.
+
+        Returns
+        -------
+        boto3.Session
+            Configured boto3 session.
+
+        """
+        if self.use_web_identity:
+            try:
+                token_path = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
+                if token_path is None:
+                    _raise_missing_token()
+                else:
+                    path = Path(token_path)
+                with path.open() as f:
+                    web_identity_token = f.read()
+
+                role_arn = os.environ.get("AWS_ROLE_ARN")
+                if role_arn is None:
+                    _raise_missing_role_arn()
+                else:
+                    role_arn_arg = role_arn
+
+                sts_client = boto3.client("sts")
+                response = sts_client.assume_role_with_web_identity(
+                    RoleArn=role_arn_arg,
+                    RoleSessionName=f"session-{self.job_name}",
+                    WebIdentityToken=web_identity_token,
+                    DurationSeconds=900,
+                )
+                return boto3.session.Session(
+                    aws_access_key_id=response["Credentials"]["AccessKeyId"],
+                    aws_secret_access_key=response["Credentials"]["SecretAccessKey"],
+                    aws_session_token=response["Credentials"]["SessionToken"],
+                )
+            except Exception as e:
+                logger.warning("Web identity failed: %s. Falling back to default session.", e)
+                raise
+
+        return boto3.session.Session()
+
+    def _fetch_account_alias(self) -> str:
+        """Fetch the AWS account alias.
+
+        Returns
+        -------
+        str
+            Account alias or default.
+
+        Notes
+        -----
+        Falls back to 'preproduction' alias if none found.
+
+        """
+        try:
+            aliases = boto3.client("iam").list_account_aliases().get("AccountAliases", [])
+            return aliases[0] if aliases else "electronic-monitoring-data-preproduction"
+        except botocore.exceptions.ClientError:
+            logger.warning("Failed to fetch account alias, assuming preproduction.")
+            return "electronic-monitoring-data-preproduction"
+
+    def _get_account_number(self) -> str:
+        """Return the AWS account number."""
+        try:
+            return boto3.client("sts").get_caller_identity()["Account"]
+        except botocore.exceptions.NoCredentialsError:
+            msg = "AWS credentials not found."
+            logger.exception(msg)
+            raise RuntimeError(msg) from None
+
+    def _list_buckets(self) -> list[BucketTypeDef]:
+        """List all available S3 buckets."""
+        try:
+            return boto3.client("s3").list_buckets()["Buckets"]
+        except Exception as e:
+            logger.warning("Could not list buckets: %s", e)
+            raise
+
+    @property
+    def account_number(self) -> str:
+        """Return the AWS account number."""
+        return self.account_no
+
+    @property
+    def environment_name(self) -> str:
+        """Infer environment name from account alias.
+
+        Returns
+        -------
+        str
+            One of: prod, dev, preprod, test, or fallback to raw alias suffix.
+
+        """
+        full_env_name = self.alias.split("-")[-1]
+        return {
+            "production": "prod",
+            "development": "dev",
+            "preproduction": "preprod",
+            "test": "test",
+        }.get(full_env_name, full_env_name)
+
+    @property
+    def is_prod(self) -> bool:
+        """Check if the environment is production."""
+        return self.environment_name == "prod"
+
+    def get_full_bucket_url(
+        self,
+        bucket_prefix: str | None = None,
+        *,
+        full_prefix: bool,
+    ) -> S3Path | None:
+        """Get S3Path to bucket matching environment and prefix.
+
+        Parameters
+        ----------
+        bucket_prefix : str, optional
+            Prefix to search for (overrides default prefix).
+        full_prefix : bool, optional
+            Whether to match full bucket name exactly.
+
+        Returns
+        -------
+        Optional[S3Path]
+            S3Path to the matched bucket or None if not found.
+
+        """
+        search_prefix = bucket_prefix or self.bucket_prefix
+        expected_name = f"{self.bucket_prefix}-{self.environment_name}-{search_prefix}"
+
+        for bucket in self.bucket_list:
+            bucket_name = bucket["Name"]
+            if full_prefix:
+                if expected_name == "-".join(bucket_name.split("-")[:-1]) or expected_name == bucket_name:
+                    return S3Path(f"s3://{bucket_name}")
+            elif expected_name in bucket_name:
+                return S3Path(f"s3://{bucket_name}")
+        return None
+
+    def get_api_invoke_url(self, api_name: str, region: str) -> str:
+        """Get API invoke url from env."""
+        client = boto3.client("apigateway", region)
+        rest_api_response = client.get_rest_apis()
+        matches = [it for it in rest_api_response["items"] if it["name"] == api_name]
+        if len(matches) > 1:
+            raise ValueError
+        api_details = matches[0]
+        return f"https://{api_details['id']}.execute-api.{region}.amazonaws.com/"
+
+    def refresh_credentials(self) -> CredentialsTypeDef:
+        """Refresh credentials via STS.
+
+        Returns
+        -------
+        dict
+            New credentials.
+
+        """
+        try:
+            token_path = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
+            if token_path is None:
+                _raise_missing_token()
+            else:
+                path = Path(token_path)
+            with path.open() as f:
+                web_identity_token = f.read()
+
+            role_arn = os.environ.get("AWS_ROLE_ARN")
+            if role_arn is None:
+                _raise_missing_role_arn()
+            else:
+                roel_arn_arg = role_arn
+
+            sts_client = boto3.client("sts")
+            rand_suffix = "".join(secrets.choice(alphabet) for _ in range(10))
+            response_assume_role = sts_client.assume_role_with_web_identity(
+                RoleArn=roel_arn_arg,
+                RoleSessionName=f"session-{self.job_name}-{rand_suffix}",
+                WebIdentityToken=web_identity_token,
+                DurationSeconds=900,
+            )
+            return response_assume_role["Credentials"]
+
+        except Exception:
+            logger.exception("Web identity failed. Falling back to get_session_token.")
+
+        sts_client = boto3.client("sts")
+        response_session_token = sts_client.get_session_token(DurationSeconds=900)
+        return response_session_token["Credentials"]
+
+    def export_dbt_variables(self, *, actions: bool = False, airflow: bool = False) -> None:
+        """Export dbt variables for the environment."""
+        s3_data_bucket_name = self.get_full_bucket_url("cadt", full_prefix=True)
+        dbt_test_profile_workgroup = f"{self.account_number}-default"
+        dbt_suffix = "" if self.is_prod else f"_{self.environment_name}_dbt"
+        h3_lambda_arn = f"arn:aws:lambda:eu-west-2:{self.account_no}:function:h3-udf"
+
+        if actions:
+            export_suffix = f'echo "DBT_SUFFIX={dbt_suffix}" \
+                >> $GITHUB_ENV\n'
+            export_bucket = f'echo "S3_DATA_BUCKET_NAME={s3_data_bucket_name}" \
+                >> $GITHUB_ENV\n'
+            export_dbt_profile = f'echo \
+                "DBT_TEST_PROFILE_WORKGROUP={dbt_test_profile_workgroup}"\
+                 >> $GITHUB_ENV\n'
+            export_dbt_profile_location = ""
+            export_h3_lambda_arn = f"""echo \
+                export H3_LAMBDA_ARN='{h3_lambda_arn}'
+                 >> $GITHUB_ENV\n
+            """
+            export_dbt_suffix = f'echo "DBT_SUFFIX={dbt_suffix}" \
+                >> $GITHUB_ENV\n'
+        else:
+            export_suffix = f"export DBT_SUFFIX='{dbt_suffix}'\n"
+            export_bucket = f"export S3_DATA_BUCKET_NAME='{s3_data_bucket_name}'\n"
+            export_dbt_profile = f"""
+                export DBT_TEST_PROFILE_WORKGROUP='{dbt_test_profile_workgroup}'\n
+                """
+            export_dbt_profile_location = 'export DBT_PROFILES_DIR="../.dbt/"\n'
+            export_h3_lambda_arn = f"export H3_LAMBDA_ARN='{h3_lambda_arn}'"
+            export_dbt_suffix = f"export DBT_SUFFIX='{dbt_suffix}'\n"
+
+        with Path("set_env.sh").open("w") as f:
+            f.write(export_suffix)
+            f.write(export_bucket)
+            f.write(export_dbt_profile)
+            if not airflow:
+                f.write(export_dbt_profile_location)
+            f.write(export_h3_lambda_arn)
+            f.write(export_dbt_suffix)
+
+    @classmethod
+    def clear(cls) -> None:
+        """Reset the singleton instance.
+
+        Use this to force restart of the class. Mainly for testing.
+        """
+        cls._instance = None
+
+    @classmethod
+    def instance(cls) -> "Environment | None":
+        """Return the current singleton instance, if any. Mainly for testing."""
+        return cls._instance

data_factory_utils-1.0.0.dist-info/METADATA

@@ -0,0 +1,52 @@
+Metadata-Version: 2.3
+Name: data-factory-utils
+Version: 1.0.0
+Summary: Utility functions for interacting with data factories.
+Requires-Dist: boto3>=1.42.8
+Requires-Dist: boto3-stubs>=1.42.89
+Requires-Dist: botocore>=1.42.7
+Requires-Dist: cloudpathlib>=0.23.0
+Requires-Dist: mypy-boto3>=1.42.3
+Requires-Dist: mypy-boto3-s3>=1.42.85
+Requires-Dist: mypy-boto3-sts>=1.42.3
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+# data-factory-utils
+A package for random utils for data factories.
+
+## Installation
+This is a published package. Install using your favourite installation method.
+
+```bash
+uv add data-factory-utils
+pip install data-factory-utils
+```
+
+## Usage
+### Environment functions
+This set of functions reads from your data factory dynamically. It should infer the environment you are in as well.
+
+No matter how many times you initiate the class, it will re-use old variables. To do so...
+```python
+from data_factory_utils.environment import Environment
+env = Environment()
+```
+
+To return information about the environment (if we are in development with account number 0101010101):
+```python
+env.account_no
+# 0101010101
+env.environment_name
+# dev
+env.is_prod
+# False
+```
+
+To get an S3 bucket name (outputted as `cloudpathlib`'s `S3Path`) (let us imagine here that the name is `emds-dev-random-name-202512161154001309058001`):
+
+```python
+s3_random_name_bucket = env.get_full_bucket_url("random-name", full_prefix=True)
+print(str(s3_random_name_bucket.bucket))
+# emds-dev-random-name-202512161154001309058001
+```

data_factory_utils-1.0.0.dist-info/RECORD

@@ -0,0 +1,5 @@
+data_factory_utils/__init__.py,sha256=2DpcGqE-C0G785KJpkdZBg_em4jPuutiiIXgIVtLDvc,17
+data_factory_utils/environment.py,sha256=hvj25smOkmy3bgfdTqNtBs_KVhVQ4kqwaCDMw7AfFkQ,11639
+data_factory_utils-1.0.0.dist-info/WHEEL,sha256=s_zqWxHFEH8b58BCtf46hFCqPaISurdB9R1XJ8za6XI,80
+data_factory_utils-1.0.0.dist-info/METADATA,sha256=Ezvz6ewe99Vi-1h5Jjx6Q6X0RqS4tWDSsS_vFAFmH-k,1534
+data_factory_utils-1.0.0.dist-info/RECORD,,

data_factory_utils-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.11.6
+Root-Is-Purelib: true
+Tag: py3-none-any