feat : add serverless data pipeline scripts

Author: Sanjay-dev-ds

.idea/.gitignore

@@ -0,0 +1,40 @@
+import json
+import boto3
+import urllib3
+import datetime
+import os
+
+FIREHOSE_NAME = os.environ['FIREHOSE_NAME']
+
+
+def lambda_handler(event, context):
+    http = urllib3.PoolManager()
+
+    r = http.request("POST", "https://www.cse.lk/api/aspiData")
+
+    # turn it into a dictionary
+    r_dict = json.loads(r.data.decode(encoding='utf-8', errors='strict'))
+
+    # extract pieces of the dictionary
+    processed_dict = {}
+    processed_dict['id'] = r_dict['id']
+    processed_dict['aspi_value'] = r_dict['value']
+    processed_dict['lowValue'] = r_dict['lowValue']
+    processed_dict['highValue'] = r_dict['highValue']
+    processed_dict['change'] = r_dict['change']
+    processed_dict['sectorId'] = r_dict['sectorId']
+    processed_dict['timestamp'] = r_dict['timestamp']
+
+    # turn it into a string and add a newline
+    msg = str(processed_dict) + '\n'
+
+    fh = boto3.client('firehose')
+
+    reply = fh.put_record(
+        DeliveryStreamName=FIREHOSE_NAME,
+        Record={
+            'Data': msg
+        }
+    )
+
+    return msg

data-extractor.py

@@ -0,0 +1,78 @@
+import sys
+import boto3
+
+
+ATHENA_CLIENT = boto3.client('athena')
+GLUE_CLIENT = boto3.client('glue')
+MY_DATABASE = 'cse-aspi'
+
+SOURCE_TABLE_NAME = 'cse_aspi_index_raw_data_bucket'
+
+STAGE_TABLE_NAME = 'stage_cse_aspi_index_parquet_data_table'
+STAGE_TABLE_S3_BUCKET = 's3://cse-aspi-index-parquet-data-bucket/'
+
+
+QUERY_RESULTS_S3_BUCKET = 's3://athena-query-results-output-data-bucket/'
+
+QUERY_STATUS_CODES = ["FAILED", "SUCCEEDED", "CANCELLED"]
+
+
+def create_stage_table():
+    queryStart = ATHENA_CLIENT.start_query_execution(
+        QueryString=f"""
+        CREATE TABLE "{STAGE_TABLE_NAME}" WITH
+        (external_location='{STAGE_TABLE_S3_BUCKET}',
+        format='PARQUET',
+        write_compression='SNAPPY',
+        partitioned_by = ARRAY['year', 'month', 'day']
+        )
+        AS
+        SELECT
+            id AS cse_aspi_index_id
+            ,aspi_value  as cse_aspi_value
+            ,lowValue as cse_aspi_low_value
+            ,highValue as cse_aspi_high_value
+            ,change as cse_aspi_change
+            ,sectorId as cse_aspi_sector_id
+            ,timestamp as cse_aspi_timestamp
+            ,cast(from_unixtime(timestamp / 1000) AT TIME ZONE 'Asia/Kolkata' as timestamp) AS ist_timestamp
+            ,year(CAST(from_unixtime(timestamp / 1000) AT TIME ZONE 'Asia/Kolkata' AS timestamp)) AS year
+            ,month(CAST(from_unixtime(timestamp / 1000) AT TIME ZONE 'Asia/Kolkata' AS timestamp)) AS month
+            ,day(CAST(from_unixtime(timestamp / 1000) AT TIME ZONE 'Asia/Kolkata' AS timestamp)) AS day
+        FROM "{MY_DATABASE}"."{SOURCE_TABLE_NAME}"
+        ;
+        """,
+        QueryExecutionContext={
+            'Database': f'{MY_DATABASE}'
+        },
+        ResultConfiguration={'OutputLocation': f'{QUERY_RESULTS_S3_BUCKET}'}
+    )
+
+    res = ATHENA_CLIENT.get_query_execution(QueryExecutionId=queryStart["QueryExecutionId"])
+
+    while res["QueryExecution"]["Status"]["State"] not in QUERY_STATUS_CODES:
+        res = ATHENA_CLIENT.get_query_execution(QueryExecutionId=queryStart["QueryExecutionId"])
+
+    return res
+
+
+def drop_table_if_exists():
+    try:
+        GLUE_CLIENT.get_table(DatabaseName=MY_DATABASE, Name=STAGE_TABLE_NAME)
+        GLUE_CLIENT.delete_table(DatabaseName=MY_DATABASE, Name=STAGE_TABLE_NAME)
+        print(f"Dropped existing table: {STAGE_TABLE_NAME}")
+    except GLUE_CLIENT.exceptions.EntityNotFoundException:
+        print(f"No existing table found: {STAGE_TABLE_NAME}")
+
+
+def response_handler(response):
+    if response["QueryExecution"]["Status"]["State"] == 'FAILED':
+        print(f"Error creating stage table: {response}")
+        sys.exit(response["QueryExecution"]["Status"]["StateChangeReason"])
+    else:
+        print(f"Stage table {STAGE_TABLE_NAME} created successfully")
+
+
+if __name__ == '__main__':
+    drop_table_if_exists()
+    response_handler(create_stage_table())

glue-scripts/01_create_or_replace_stage_table.py

@@ -0,0 +1,41 @@
+import sys
+import boto3
+import awswrangler as wr
+
+ATHENA_CLIENT = wr.athena
+GLUE_CLIENT = boto3.client('glue')
+MY_DATABASE = 'cse-aspi'
+
+STAGE_TABLE_NAME = 'stage_cse_aspi_index_parquet_data_table'
+
+QUERY_RESULTS_S3_BUCKET = 's3://athena-query-results-output-data-bucket/'
+
+QUERY_STATUS_CODES = ["FAILED", "SUCCEEDED", "CANCELLED"]
+
+
+def data_quality_evaluator() -> None:
+    NULL_DQ_CHECK = f"""
+    SELECT 
+      COUNT(CASE WHEN cse_aspi_index_id IS NULL THEN 1 END) +
+      COUNT(CASE WHEN cse_aspi_value IS NULL THEN 1 END) +
+      COUNT(CASE WHEN cse_aspi_low_value IS NULL THEN 1 END) +
+      COUNT(CASE WHEN cse_aspi_high_value IS NULL THEN 1 END) +
+      COUNT(CASE WHEN cse_aspi_change IS NULL THEN 1 END) +
+      COUNT(CASE WHEN cse_aspi_sector_id IS NULL THEN 1 END) +
+      COUNT(CASE WHEN cse_aspi_timestamp IS NULL THEN 1 END) +
+      COUNT(CASE WHEN ist_timestamp IS NULL THEN 1 END) AS total_null_values
+    FROM {STAGE_TABLE_NAME};
+    """
+
+    df = ATHENA_CLIENT.read_sql_query(sql=NULL_DQ_CHECK, database=MY_DATABASE)
+
+    # exit if we get a result > 0
+    # else, the check was successful
+    if df['total_null_values'][0] > 0:
+        sys.exit('Results returned. Quality check failed.')
+    else:
+        print('Quality check passed.')
+
+
+if __name__ == '__main__':
+    data_quality_evaluator()

glue-scripts/02_data_quality_check_on_stage_table.py

@@ -0,0 +1,69 @@
+import boto3
+from botocore.exceptions import ClientError
+import sys
+
+ATHENA_CLIENT = boto3.client('athena')
+GLUE_CLIENT = boto3.client('glue')
+MY_DATABASE = 'cse-aspi'
+PROD_TABLE_NAME = 'prod_cse_aspi_index_data_table'
+PROD_TABLE_S3_BUCKET = 's3://cse-aspi-index-prod-table-data-bucket/'
+QUERY_RESULTS_S3_BUCKET = 's3://athena-query-results-output-data-bucket/'
+QUERY_STATUS_CODES = ["FAILED", "SUCCEEDED", "CANCELLED"]
+
+
+def check_table_exists():
+    try:
+        GLUE_CLIENT.get_table(DatabaseName=MY_DATABASE, Name=PROD_TABLE_NAME)
+        return True
+    except ClientError as e:
+        if e.response['Error']['Code'] == 'EntityNotFoundException':
+            return False
+        else:
+            raise e  # Re-raise unexpected errors for handling
+
+
+def create_production_table():
+    query = f"""
+    CREATE TABLE {PROD_TABLE_NAME} (
+        cse_aspi_index_id int,
+        cse_aspi_value float,
+        cse_aspi_low_value float,
+        cse_aspi_high_value float,
+        cse_aspi_change float,
+        cse_aspi_sector_id int,
+        ist_timestamp timestamp
+    )
+    LOCATION '{PROD_TABLE_S3_BUCKET}'
+    TBLPROPERTIES ( 'table_type' = 'ICEBERG' );
+    """
+
+    try:
+        query_start = ATHENA_CLIENT.start_query_execution(
+            QueryString=query,
+            QueryExecutionContext={
+                'Database': MY_DATABASE
+            },
+            ResultConfiguration={'OutputLocation': QUERY_RESULTS_S3_BUCKET}
+        )
+        response = ATHENA_CLIENT.get_query_execution(QueryExecutionId=query_start["QueryExecutionId"])
+
+        while response["QueryExecution"]["Status"]["State"] not in QUERY_STATUS_CODES:
+            response = ATHENA_CLIENT.get_query_execution(QueryExecutionId=query_start["QueryExecutionId"])
+
+        return response
+    except ClientError as e:
+        print(f"Error creating production table: {e}")
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    if check_table_exists():
+        print(f"Table {PROD_TABLE_NAME} already exists")
+    else:
+        print(f"Creating table {PROD_TABLE_NAME}")
+        response = create_production_table()
+
+        if response["QueryExecution"]["Status"]["State"] == 'FAILED':
+            sys.exit(response["QueryExecution"]["Status"]["StateChangeReason"])
+
+        print("Table creation successful")

glue-scripts/03_create_prod_table.py

@@ -0,0 +1,103 @@
+import boto3
+from botocore.exceptions import ClientError
+import sys
+
+ATHENA_CLIENT = boto3.client('athena')
+GLUE_CLIENT = boto3.client('glue')
+MY_DATABASE = 'cse-aspi'
+
+PROD_TABLE_NAME = 'prod_cse_aspi_index_data_table'
+PROD_TABLE_S3_BUCKET = 's3://cse-aspi-index-prod-table-data-bucket/'
+
+STAGE_TABLE_NAME = 'stage_cse_aspi_index_parquet_data_table'
+STAGE_TABLE_S3_BUCKET = 's3://cse-aspi-index-parquet-data-bucket/'
+
+
+QUERY_RESULTS_S3_BUCKET = 's3://athena-query-results-output-data-bucket/'
+QUERY_STATUS_CODES = ["FAILED", "SUCCEEDED", "CANCELLED"]
+
+
+def merge_data_operation():
+    query = f"""
+        MERGE INTO {PROD_TABLE_NAME} AS target
+        USING (
+            -- Select the unique records with the latest timestamp
+            WITH LatestRecords AS (
+                SELECT
+                    cse_aspi_index_id,
+                    cse_aspi_value,
+                    cse_aspi_low_value,
+                    cse_aspi_high_value,
+                    cse_aspi_change,
+                    cse_aspi_sector_id,
+                    CAST(ist_timestamp AS TIMESTAMP(6)) AS ist_timestamp,
+                    ROW_NUMBER() OVER (PARTITION BY cse_aspi_index_id ORDER BY cse_aspi_timestamp DESC) AS rn
+                FROM {STAGE_TABLE_NAME}
+            )
+            SELECT
+                cse_aspi_index_id,
+                cse_aspi_value,
+                cse_aspi_low_value,
+                cse_aspi_high_value,
+                cse_aspi_change,
+                cse_aspi_sector_id,
+                ist_timestamp
+            FROM LatestRecords
+            WHERE rn = 1
+        ) AS source
+        ON target.cse_aspi_index_id = source.cse_aspi_index_id
+        WHEN MATCHED THEN
+            UPDATE SET
+                cse_aspi_value = source.cse_aspi_value,
+                cse_aspi_low_value = source.cse_aspi_low_value,
+                cse_aspi_high_value = source.cse_aspi_high_value,
+                cse_aspi_change = source.cse_aspi_change,
+                cse_aspi_sector_id = source.cse_aspi_sector_id,
+                ist_timestamp = source.ist_timestamp
+        WHEN NOT MATCHED THEN
+            INSERT (
+                cse_aspi_index_id,
+                cse_aspi_value,
+                cse_aspi_low_value,
+                cse_aspi_high_value,
+                cse_aspi_change,
+                cse_aspi_sector_id,
+                ist_timestamp
+            ) VALUES (
+                source.cse_aspi_index_id,
+                source.cse_aspi_value,
+                source.cse_aspi_low_value,
+                source.cse_aspi_high_value,
+                source.cse_aspi_change,
+                source.cse_aspi_sector_id,
+                source.ist_timestamp
+            );
+        ;
+        """
+    try:
+        query_start = ATHENA_CLIENT.start_query_execution(
+            QueryString=query,
+            QueryExecutionContext={
+                'Database': MY_DATABASE
+            },
+            ResultConfiguration={'OutputLocation': QUERY_RESULTS_S3_BUCKET}
+        )
+        response = ATHENA_CLIENT.get_query_execution(QueryExecutionId=query_start["QueryExecutionId"])
+
+        while response["QueryExecution"]["Status"]["State"] not in QUERY_STATUS_CODES:
+            response = ATHENA_CLIENT.get_query_execution(QueryExecutionId=query_start["QueryExecutionId"])
+
+        return response
+    except ClientError as e:
+        print(f"Error creating production table: {e}")
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+
+    response = merge_data_operation()
+
+    if response["QueryExecution"]["Status"]["State"] == 'FAILED':
+        sys.exit(response["QueryExecution"]["Status"]["StateChangeReason"])
+
+    print(f"Table {PROD_TABLE_NAME} Refreshed successfully")