tenant1/service_request_metrics/main.py.notebook

import marimo

__generated_with = "0.13.15"
app = marimo.App()


@app.cell
def init():

    import sys
    import time
    from pyspark.sql.utils import AnalysisException
    sys.path.append('/opt/spark/work-dir/')
    from workflow_templates.spark.udf_manager import bootstrap_udfs 
    from util import get_logger, observe_metrics, collect_metrics, log_info, log_error, forgiving_serializer
    from pyspark.sql.functions import udf 
    from pyspark.sql.functions import count, expr, lit, input_file_name
    from pyspark.sql.types import StringType, IntegerType, MapType, StructType,StructField
    from postal.parser import parse_address
    import uuid
    from pathlib import Path
    from pyspark import SparkConf, Row
    from pyspark.sql import SparkSession
    from pyspark.sql.observation import Observation
    from pyspark import StorageLevel
    import os
    import pandas as pd
    import polars as pl
    import pyarrow as pa
    from pyspark.sql.functions import approx_count_distinct, avg, collect_list, collect_set, corr, count, countDistinct, covar_pop, covar_samp, first, kurtosis, last, max, mean, min, skewness, stddev, stddev_pop, stddev_samp, sum, var_pop, var_samp, variance,expr,to_json,struct, date_format, col, lit, when, regexp_replace, ltrim, lpad, format_number
    from functools import reduce
    from handle_structs_or_arrays import preprocess_then_expand
    import requests
    from requests.adapters import HTTPAdapter
    from urllib3.util.retry import Retry
    from jinja2 import Template
    import json
    import orjson

    from ocular_ai_sdk import OcularClient
    from ocular_ai_sdk.exceptions import (
        OcularSDKException,
        AuthenticationError,
        ResourceNotFoundError
    )


    from secrets_manager import SecretsManager

    from WorkflowManager import WorkflowDSL, WorkflowManager
    from KnowledgebaseManager import KnowledgebaseManager
    from gitea_client import GiteaClient, WorkspaceVersionedContent
    from FilesystemManager import FilesystemManager, SupportedFilesystemType
    from Materialization import Materialization

    import ssl
    from urllib.request import Request, urlopen
    from urllib.parse import urlencode
    from urllib.error import HTTPError

    init_start_time=time.time()

    LOGGER = get_logger()
    alias_str='abcdefghijklmnopqrstuvwxyz'
    workspace = os.getenv('WORKSPACE') or 'exp360-cus-uat'
    workflow = 'service_request_metrics'
    execution_environment = os.getenv('EXECUTION_ENVIRONMENT') or 'CLUSTER'

    job_id = os.getenv("EXECUTION_ID") or str(uuid.uuid4())
    retry_job_id = os.getenv("RETRY_EXECUTION_ID") or ''

    log_info(LOGGER, f"Workspace: '{workspace}', Workflow: '{workflow}', Execution Environment: '{execution_environment}', Job Id: '{job_id}', Retry Job Id: '{retry_job_id}'")

    sm = SecretsManager(os.getenv('SECRET_MANAGER_URL'), os.getenv('SECRET_MANAGER_NAMESPACE'), os.getenv('SECRET_MANAGER_ENV'), os.getenv('SECRET_MANAGER_TOKEN'))
    secrets = sm.list_secrets(workspace)

    gitea_client=GiteaClient(os.getenv('GITEA_HOST'), os.getenv('GITEA_TOKEN'), os.getenv('GITEA_OWNER') or 'gitea_admin', os.getenv('GITEA_REPO') or 'tenant1')
    workspaceVersionedContent=WorkspaceVersionedContent(gitea_client)

    client = OcularClient(
        pat_token=secrets.get('OCULAR_AI_PAT_TOKEN')
    )

    if 'AZURE_SERVICE_PRINCIPAL' in secrets:
        _storage_options=orjson.loads(secrets['AZURE_SERVICE_PRINCIPAL'])
    else:
        _storage_options = {
            'key': secrets.get('S3_ACCESS_KEY'),
            'secret': secrets.get('S3_SECRET_KEY'),
            'region': secrets.get('S3_REGION')
        }

    filesystemManager = FilesystemManager.create(secrets.get('LAKEHOUSE_BUCKET'), storage_options=_storage_options)
    if retry_job_id:
        logs = Materialization.get_execution_history_by_job_id(filesystemManager, secrets.get('LAKEHOUSE_BUCKET'), workspace, workflow, retry_job_id, selected_components=['finalize']).to_dicts()
        if len(logs) == 1 and logs[0].get('metrics').get('execute_status') == 'SUCCESS':
            log_info(LOGGER, f"Workspace: '{workspace}', Workflow: '{workflow}', Execution Environment: '{execution_environment}', Job Id: '{job_id}' - Retry Job Id: '{retry_job_id}' was already successful. Hence exiting to forward processing to next in chain.")
            sys.exit(0)

    _conf = SparkConf()
    _params = {
     "spark.jars.ivy": "/opt/spark/.ivy2/",
     "spark.hadoop.fs.s3a.access.key": secrets.get('S3_ACCESS_KEY'),
     "spark.hadoop.fs.s3a.secret.key": secrets.get('S3_SECRET_KEY'),
     "spark.hadoop.fs.s3a.aws.region": secrets.get("S3_REGION") or "us-west-1",
     "spark.sql.catalog.dremio.warehouse" : secrets.get('LAKEHOUSE_BUCKET'),
     "spark.hadoop.fs.s3a.aws.credentials.provider": "com.amazonaws.auth.DefaultAWSCredentialsProviderChain",
     "spark.hadoop.fs.s3.aws.credentials.provider": "com.amazonaws.auth.DefaultAWSCredentialsProviderChain",
     "spark.sql.catalog.dremio" : "org.apache.iceberg.spark.SparkCatalog",
     "spark.sql.catalog.dremio.type" : "hadoop",
     "spark.hadoop.fs.s3a.impl": "org.apache.hadoop.fs.s3a.S3AFileSystem",
     "spark.hadoop.fs.s3.impl": "org.apache.hadoop.fs.s3a.S3AFileSystem",
     "spark.hadoop.fs.gs.impl": "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem",
     "spark.sql.extensions": "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions",
     "spark.driver.extraJavaOptions": "--add-opens=java.base/java.nio=ALL-UNNAMED",
     "spark.executor.extraJavaOptions": "--add-opens=java.base/java.nio=ALL-UNNAMED"
    }

    if filesystemManager.storage_type == SupportedFilesystemType.AZUREBLOB:
        _params[f"fs.azure.account.auth.type.{_storage_options['account_name']}.dfs.core.windows.net"] = "OAuth"
        _params[f"fs.azure.account.oauth.provider.type.{_storage_options['account_name']}.dfs.core.windows.net"] = "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider"
        _params[f"fs.azure.account.oauth2.client.id.{_storage_options['account_name']}.dfs.core.windows.net"] = _storage_options['client_id']
        _params[f"fs.azure.account.oauth2.client.secret.{_storage_options['account_name']}.dfs.core.windows.net"] = _storage_options['client_secret']
        _params[f"fs.azure.account.oauth2.client.endpoint.{_storage_options['account_name']}.dfs.core.windows.net"] = f"https://login.microsoftonline.com/{_storage_options['tenant_id']}/oauth2/v2.0/token"



    _conf.setAll(list(_params.items()))

    spark = SparkSession.builder.appName(workspace).config(conf=_conf).getOrCreate()
    bootstrap_udfs(spark)

    materialization = Materialization(spark, secrets.get('LAKEHOUSE_BUCKET'), workspace, workflow, job_id, retry_job_id, execution_environment, LOGGER)

    init_dependency_key="init"


    init_end_time=time.time()
    return (
        AnalysisException,
        LOGGER,
        collect_metrics,
        job_id,
        log_error,
        log_info,
        materialization,
        observe_metrics,
        os,
        secrets,
        spark,
        sum,
        time,
    )


@app.cell
def ActionsAuditData(LOGGER, log_error, observe_metrics, secrets, spark, time):


    ActionsAuditData_start_time=time.time()

    ActionsAuditData_fail_on_error=""
    try:
        _ActionsAuditData_options = {
                'jdbc':{
                    'dbtable': """actionsaudit""",
                    'url':secrets.get(''),
                    'driver':''
                },
                'kafka' : {
                    'kafka.bootstrap.servers' : secrets.get('OCULAR_KAFKA_BOOTSTRAP_SERVERS'),
                    'subscribe' : '',
                    'startingOffsets' : 'earliest'
                },
                'cobol' : {
                    'copybook' : '',
                    'encoding' : '',
                    'is_text': False,
                    'schema_retention_policy' : 'collapse_root'
                }
        }

    

        _ActionsAuditData_reader = spark.read.format('iceberg')

    

    
        ActionsAuditData_df = _ActionsAuditData_reader.load('dremio.actionsaudit')
    
    
        ActionsAuditData_df, ActionsAuditData_observer = observe_metrics("ActionsAuditData_df", ActionsAuditData_df)

    

        ActionsAuditData_df.createOrReplaceTempView('ActionsAuditData_df')


        ActionsAuditData_dependency_key="ActionsAuditData"
    
        ActionsAuditData_execute_status="SUCCESS"
    except Exception as e:
        ActionsAuditData_error = e
        log_error(LOGGER, f"Component ActionsAuditData Failed", e)
        ActionsAuditData_execute_status="ERROR"
    
        raise e
    
    finally:
        ActionsAuditData_end_time=time.time()
    return ActionsAuditData_dependency_key, ActionsAuditData_df


@app.cell
def data_mapper__1(
    ActionsAuditData_dependency_key,
    ActionsAuditData_df,
    LOGGER,
    job_id,
    log_error,
    log_info,
    observe_metrics,
    spark,
    time,
):


    data_mapper__1_start_time=time.time()

    data_mapper__1_fail_on_error=""
    try:
    
    

        _data_mapper__1_select_clause=ActionsAuditData_df.columns if False else []
    
        _data_mapper__1_select_clause.append('''DATE(action_date) AS action_date''')
    
        _data_mapper__1_select_clause.append('''sub_category AS service_type''')
    
        _data_mapper__1_select_clause.append('''action_count AS action_count''')
    
        try:
            data_mapper__1_df=spark.sql(("SELECT " + ', '.join(_data_mapper__1_select_clause) + " FROM ActionsAuditData_df").replace("{job_id}",f"'{job_id}'"))
        except Exception as e:
            data_mapper__1_df = ActionsAuditData_df.limit(0)
            log_info(LOGGER, f"error while mapping the data :{e} " )
        data_mapper__1_df, data_mapper__1_observer = observe_metrics("data_mapper__1_df", data_mapper__1_df)

    

        data_mapper__1_df.createOrReplaceTempView("data_mapper__1_df")

        data_mapper__1_dependency_key="data_mapper__1"
    
        print(ActionsAuditData_dependency_key)
    
        data_mapper__1_execute_status="SUCCESS"
    except Exception as e:
        data_mapper__1_error = e
        log_error(LOGGER, f"Component data_mapper__1 Failed", e)
        data_mapper__1_execute_status="ERROR"
    
        raise e
    
    finally:
        data_mapper__1_end_time=time.time()
    return (data_mapper__1_df,)


@app.cell
def LatestServiceRequests(
    AnalysisException,
    LOGGER,
    data_mapper__1_df,
    log_error,
    log_info,
    observe_metrics,
    spark,
    time,
):


    LatestServiceRequests_start_time=time.time()

    print(data_mapper__1_df.columns)
    LatestServiceRequests_fail_on_error=""
    try:
        try:
            LatestServiceRequests_df = spark.sql("select * from data_mapper__1_df where action_date >=  COALESCE((SELECT MAX(DATE(action_date)) FROM dremio.servicemetrics), (SELECT MIN(action_date) FROM data_mapper__1_df))")
        except AnalysisException as e:
            log_info(LOGGER, f"error while filtering data : {e}")
            LatestServiceRequests_df = data_mapper__1_df.limit(0)
        except Exception as e:
            log_info(LOGGER, f"Unexpected error: {e}")
            LatestServiceRequests_df = data_mapper__1_df.limit(0)
        LatestServiceRequests_df, LatestServiceRequests_observer = observe_metrics("LatestServiceRequests_df", LatestServiceRequests_df)

    

        LatestServiceRequests_df.createOrReplaceTempView('LatestServiceRequests_df')

        LatestServiceRequests_dependency_key="LatestServiceRequests"
    
        LatestServiceRequests_execute_status="SUCCESS"
    except Exception as e:
        LatestServiceRequests_error = e
        log_error(LOGGER, f"Component LatestServiceRequests Failed", e)
        LatestServiceRequests_execute_status="ERROR"
    
        raise e
    
    finally:
        LatestServiceRequests_end_time=time.time()
    return LatestServiceRequests_dependency_key, LatestServiceRequests_df


@app.cell
def aggregate__3(
    LOGGER,
    LatestServiceRequests_dependency_key,
    LatestServiceRequests_df,
    log_error,
    observe_metrics,
    sum,
    time,
):


    aggregate__3_start_time=time.time()

    aggregate__3_fail_on_error="True"
    try:
    
    
    
    
    
    




        aggregate__3_df = LatestServiceRequests_df.groupBy(
        
                "action_date", 
        
                "service_type"
            ).agg(
       
                sum('action_count').alias("service_count")
       
         )
        aggregate__3_df, aggregate__3_observer = observe_metrics("aggregate__3_df", aggregate__3_df)

    

        aggregate__3_df.createOrReplaceTempView('aggregate__3_df')

        aggregate__3_dependency_key="aggregate__3"
    
        print(LatestServiceRequests_dependency_key)
    
        aggregate__3_execute_status="SUCCESS"
    except Exception as e:
        aggregate__3_error = e
        log_error(LOGGER, f"Component aggregate__3 Failed", e)
        aggregate__3_execute_status="ERROR"
    
        raise e
    
    finally:
        aggregate__3_end_time=time.time()
    return aggregate__3_dependency_key, aggregate__3_df


@app.cell
def ServiceRequestMetricsWriter(
    CheckpointOutput_dependency_key,
    CheckpointOutput_df,
    LOGGER,
    log_error,
    secrets,
    time,
):


    ServiceRequestMetricsWriter_start_time=time.time()

    ServiceRequestMetricsWriter_fail_on_error=""
    try:

    
        _ServiceRequestMetricsWriter_options = {
                'jdbc':{
                    'dbtable': 'servicemetrics',
                    'url':secrets.get(''),
                    'driver':'',
                    'stringtype': 'unspecified'
                },
                'kafka' : {
                    'kafka.bootstrap.servers' : secrets.get('OCULAR_KAFKA_BOOTSTRAP_SERVERS'),
                    'topic' : ''
                }


        }
    
        _CheckpointOutput_df = CheckpointOutput_df
    
        _ServiceRequestMetricsWriter_writer = _CheckpointOutput_df.write.format('iceberg').mode('append')

    
        _ServiceRequestMetricsWriter_writer.save('dremio.servicemetrics')
    

        ServiceRequestMetricsWriter_dependency_key="ServiceRequestMetricsWriter"
    
        print(CheckpointOutput_dependency_key)
    
        ServiceRequestMetricsWriter_execute_status="SUCCESS"
    except Exception as e:
        ServiceRequestMetricsWriter_error = e
        log_error(LOGGER, f"Component ServiceRequestMetricsWriter Failed", e)
        ServiceRequestMetricsWriter_execute_status="ERROR"
    
        raise e
    
    finally:
        ServiceRequestMetricsWriter_end_time=time.time()
    return


@app.cell
def CheckpointOutput(aggregate__3_dependency_key, aggregate__3_df, time):

    CheckpointOutput_start_time=time.time()

    CheckpointOutput_df = aggregate__3_df.localCheckpoint()
    # aggregate__3_df.persist()
    CheckpointOutput_df.createOrReplaceTempView("CheckpointOutput_df")

    CheckpointOutput_end_time=time.time()

    CheckpointOutput_dependency_key="CheckpointOutput"

    print(aggregate__3_dependency_key)

    return CheckpointOutput_dependency_key, CheckpointOutput_df


@app.cell
def finalize(
    LOGGER,
    collect_metrics,
    log_info,
    materialization,
    os,
    spark,
    time,
):

    finalize_start_time=time.time()

    metrics = {
        'data': collect_metrics(locals()),
    }
    materialization.materialized_execution_history({'finalize': {'execute_status': 'SUCCESS', 'fail_on_error': 'False', 'execution_order': os.environ.get('EXECUTION_ORDER')}, **metrics['data']})
    log_info(LOGGER, f"Workflow Data metrics: {metrics['data']}")

    finalize_end_time=time.time()

    if os.getenv('EXECUTION_ENVIRONMENT'):
        spark.stop()
    return


if __name__ == "__main__":
    app.run()