abi_ds_utils 1.2.4__py3-none-any.whl

abi_ds_utils/__init__.py

@@ -0,0 +1,14 @@
+from abi_ds_utils.airflow import write_to_xcom
+from abi_ds_utils.aws import get_parameter, get_secret
+from abi_ds_utils.spark import getSpark, getGlue
+
+get_spark = getSpark
+get_glue = getGlue
+
+__all__ = [
+    "get_glue",
+    "get_secret",
+    "get_spark",
+    "get_parameter",
+    "write_to_xcom"
+]

abi_ds_utils/airflow.py

@@ -0,0 +1,29 @@
+import json
+from typing import Dict
+from pathlib import Path
+
+
+def write_to_xcom(dict_value: Dict) -> Dict:
+    """Write `dict_value` for XCOM to pass it to the next task
+
+    ref: https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/stable/operators.html#how-does-xcom-work
+
+    """
+    file_save = Path('/airflow/xcom/return.json')
+    file_save.parent.mkdir(parents=True, exist_ok=True)
+
+    # Get data from the file if there is some
+    if file_save.is_file():
+        with file_save.open('r') as fin:
+            dict_restored = json.load(fin)
+
+    else:
+        dict_restored = dict()
+
+    # Update dict with new data
+    dict_restored.update(dict_value)
+
+    # Overwrite file with updated dict
+    with file_save.open('w') as fout:
+        json.dump(dict_restored, fout)
+    return dict_restored

abi_ds_utils/aws.py

@@ -0,0 +1,51 @@
+import json
+import base64
+from typing import Dict, Optional
+import boto3
+from botocore.exceptions import ClientError
+
+
+def get_secret(secret_name: str, region_name: str) -> Dict:
+    # Create a Secrets Manager client
+    session = boto3.session.Session()
+    client = session.client(
+        service_name='secretsmanager',
+        region_name=region_name
+    )
+
+    # In this sample we only handle the specific exceptions for the 'GetSecretValue' API.
+    # See https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html
+    # We rethrow the exception by default.
+    try:
+        get_secret_value_response = client.get_secret_value(
+            SecretId=secret_name
+        )
+    except ClientError as e:
+        raise e
+    else:
+        # Decrypts secret using the associated KMS key.
+        # Depending on whether the secret is a string or binary, one of these fields will be populated.
+        if 'SecretString' in get_secret_value_response:
+            secret = get_secret_value_response['SecretString']
+        else:
+            secret = base64.b64decode(get_secret_value_response['SecretBinary'])
+        return json.loads(secret)
+
+
+def get_parameter(name: str, region: Optional[str] = None) -> Optional[str]:
+    """Retrieve secret from Parameter Store.
+
+    :param name: Name of the parameter
+    :param region: AWS region otherwise env AWS_DEFAULT_REGION
+    :return: Value of parameter
+    """
+    if region:
+        client = boto3.client("ssm", region_name=region)
+    else:
+        client = boto3.client("ssm")
+
+    try:
+        parameter = client.get_parameter(Name=name, WithDecryption=True)
+    except ClientError as e:
+        raise e
+    return parameter.get('Parameter', {}).get('Value')

abi_ds_utils/spark.py

@@ -0,0 +1,62 @@
+import os
+from pyspark.sql import SparkSession
+
+
+def getSpark(driver_memory: str = "21g") -> SparkSession:
+    spark = (
+        SparkSession.builder
+        # General
+        .master('local[*]')
+        .config("spark.driver.maxResultSize", 0)
+
+        # Get 80% of free memory (this might be a bad idea)
+        .config("spark.driver.memory", driver_memory)
+        .config("spark.dynamicAllocation.enabled", "true")
+
+        # PyArrow for dtypes conversions
+        .config("spark.sql.execution.arrow.pyspark.enabled", "true")
+
+        # Spark 3.5.x requires a Delta 3.x line; Delta 2.2 only supports Spark 3.3.x.
+        .config('spark.jars.packages', 'org.apache.hadoop:hadoop-aws:3.3.2,io.delta:delta-spark_2.12:3.3.2')
+
+        # Delta Lake setup
+        .config("spark.hadoop.fs.s3a.connection.maximum", 128)
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+        .config("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog")
+        .config("spark.delta.logStore.class", "org.apache.spark.sql.delta.storage.S3SingleDriverLogStore")
+    )
+    
+    spark = spark.config(
+        "spark.hadoop.fs.s3a.aws.credentials.provider",
+        "com.amazonaws.auth.DefaultAWSCredentialsProviderChain"
+    )
+    return spark.getOrCreate()
+
+
+def getGlue(driver_memory: str = "21g") -> SparkSession:
+    spark = (
+        SparkSession.builder
+        # General
+        .config("spark.driver.maxResultSize", 0)
+
+        # Get 80% of free memory (this might be a bad idea)
+        .config("spark.driver.memory", driver_memory)
+        .config("spark.dynamicAllocation.enabled", "true")
+
+        # PyArrow for dtypes conversions
+        .config("spark.sql.execution.arrow.pyspark.enabled", "true")
+
+        # Delta Lake setup
+        .config("spark.hadoop.fs.s3a.connection.maximum", 128)
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+        .config("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog")
+        .config("spark.delta.logStore.class", "org.apache.spark.sql.delta.storage.S3SingleDriverLogStore")
+
+        # Glue setup
+        .config("hive.metastore.client.factory.class", "com.amazonaws.glue.catalog.metastore.AWSGlueDataCatalogHiveClientFactory")
+        .config("aws.region", os.environ.get('AWS_DEFAULT_REGION'))
+        .config("hive.metastore.glue.catalogid", os.environ.get('GLUE_CATALOG_ID'))
+        .enableHiveSupport()
+    )
+
+    return spark.getOrCreate()

abi_ds_utils-1.2.4.dist-info/METADATA

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: abi_ds_utils
+Version: 1.2.4
+Summary: Utility modules for working with spark, containers, aws and more.
+Author: Martin Matousek, Ioannis Chios
+License: Private
+Requires-Python: >=3.8
+Requires-Dist: boto3<2.0.0,>=1.21.14
+Requires-Dist: pyspark==3.5.8
+Requires-Dist: pyarrow<13.0.0,>=7.0.0; python_version < "3.12"
+Requires-Dist: pyarrow<18.0.0,>=15.0.2; python_version >= "3.12"

abi_ds_utils-1.2.4.dist-info/RECORD

@@ -0,0 +1,8 @@
+abi_ds_utils/__init__.py,sha256=uXMLKzJ-mkzoylygp1eS00zt8OV7jeTeab1EhlaPdMc,299
+abi_ds_utils/airflow.py,sha256=8cxCGZmDl2thQO_eUGl1XB8FTaolgWv3x4s6Xou0Qxw,823
+abi_ds_utils/aws.py,sha256=P8e0lgcmy1GDQWzyxeQYo2b5mE84bSzrAGIwodC4coo,1737
+abi_ds_utils/spark.py,sha256=iwk68vJ8uV_Fyi8Q7DoFNT4yay3fVxa59f3GnExxl2s,2512
+abi_ds_utils-1.2.4.dist-info/METADATA,sha256=rVogXdRb7LkFqG0jd0xuCf9ZaDm-lD4UZYHgMgXrunU,405
+abi_ds_utils-1.2.4.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+abi_ds_utils-1.2.4.dist-info/top_level.txt,sha256=SsHKDlGfPNApOYPYVuSKuXkHYHD697V8A4X865QjwVo,13
+abi_ds_utils-1.2.4.dist-info/RECORD,,

abi_ds_utils-1.2.4.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

abi_ds_utils-1.2.4.dist-info/top_level.txt

@@ -0,0 +1 @@
+abi_ds_utils