__generated_with = "0.13.15" # %% 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() # %% 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() # %% 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() # %% 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() # %% 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() # %% 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) # %% 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() # %% 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()