Add test123/main.py.notebook, test123/main.py, test123/main.workflow

test123/main.py

@@ -0,0 +1,147 @@
+
+__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 'test_charan'
+workflow = 'test123'
+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(''),
+ "spark.hadoop.fs.s3a.secret.key": secrets.get(''),
+ "spark.hadoop.fs.s3a.aws.region": secrets.get("S3_REGION") or "",
+ "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"
+}
+
+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()
+
+# %%
+
+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()