From 384a37106507f0bc0bb27bc0d1e370228b1faafa Mon Sep 17 00:00:00 2001 From: gitea_admin Date: Thu, 23 Jul 2026 13:58:18 +0000 Subject: [PATCH] Add test123/main.py.notebook, test123/main.py, test123/main.workflow --- test123/main.py | 147 ++++++++++++++++++++++++++++++++++ test123/main.py.notebook | 167 +++++++++++++++++++++++++++++++++++++++ test123/main.workflow | 1 + 3 files changed, 315 insertions(+) create mode 100644 test123/main.py create mode 100644 test123/main.py.notebook create mode 100644 test123/main.workflow diff --git a/test123/main.py b/test123/main.py new file mode 100644 index 0000000..088c316 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/test123/main.py.notebook b/test123/main.py.notebook new file mode 100644 index 0000000..16fb483 --- /dev/null +++ b/test123/main.py.notebook @@ -0,0 +1,167 @@ +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 '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() + return LOGGER, collect_metrics, log_info, materialization, os, spark, time + + +@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() diff --git a/test123/main.workflow b/test123/main.workflow new file mode 100644 index 0000000..858748e --- /dev/null +++ b/test123/main.workflow @@ -0,0 +1 @@ +{"version":"v1alpha","kind":"Notebook","metadata":{"name":"test123","description":"test123","runtime":"spark"},"spec":{"ui":{},"blocks":[]}} \ No newline at end of file