dataforge-sdk 9.2.0rc3__tar.gz

dataforge_sdk-9.2.0rc3/PKG-INFO

@@ -0,0 +1,125 @@
+Metadata-Version: 2.4
+Name: dataforge-sdk
+Version: 9.2.0rc3
+Summary: SDK for creating DataForge extensions
+Author-email: Vadim Orlov <vorlov@dataforgelabs.com>
+Project-URL: Homepage, https://docs.dataforgelabs.com
+Project-URL: Issues, https://docs.dataforgelabs.com/hc/en-us/requests/new
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# dataforge-sdk
+SDK for creating DataForge extensions.
+
+## Postgres Utilities
+
+The `dataforge.pg` module provides helper functions to execute SQL operations against the DataForge Postgres metastore:
+
+```python
+from dataforge.pg import select, update, pull
+
+# Execute a SELECT query and return a Spark DataFrame
+df = select("SELECT * FROM my_table")
+
+# Execute an UPDATE/INSERT/DELETE query
+update("UPDATE my_table SET col = 'value'")
+
+# Trigger a new data pull for source_id 123
+pull(123)
+```
+
+## IngestionSession
+
+The `IngestionSession` class manages a custom data ingestion process lifecycle.
+
+```python
+from dataforge import IngestionSession
+
+# Initialize a session (production use)
+session = IngestionSession()
+
+# Initialize a session (optional source_name/project_name for testing)
+session = IngestionSession(source_name="my_source", project_name="my_project")
+
+# Ingest data 
+# pass a function returning a DataFrame (recommended to integrate logging with DataForge)
+session.ingest(lambda: spark.read.csv("s3://bucket/path/input.csv"))
+
+# pass a DataFrame (can be used for testing, not recommended for production deployment)
+df = spark.read.csv("s3://bucket/path/input.csv")
+session.ingest(df)
+
+# ingest empty dataframe to create 0-record input
+session.ingest()
+
+
+# Fail the process with error message
+session.fail("Error message")
+
+# Retrieve latest tracking fields
+tracking = session.latest_tracking_fields()
+
+# Retrieve connection parameters for the current source
+connection_parameters = session.connection_parameters()
+
+# Retrieve custom parameters for the current source
+custom_parameters = session.custom_parameters()
+
+```
+
+## ParsingSession
+
+The `ParsingSession` class manages a custom parse process lifecycle.
+
+```python
+from dataforge import ParsingSession
+
+# Initialize a session (production use)
+session = ParsingSession()
+
+# Initialize a session (optional input_id for testing)
+session = ParsingSession(input_id=123)
+
+# Retrieve custom parameters
+params = session.custom_parameters()
+
+# Get the path of file to be parsed
+path = session.file_path
+
+# Run parsing: pass a DataFrame, a function returning a DataFrame or None (0-record file)
+session.run(lambda: spark.read.json(session.file_path))
+
+# Fail the process with error message
+session.fail("Error message")
+
+```
+
+## PostOutputSession
+
+The `PostOutputSession` class manages a custom post-output process lifecycle.
+
+```python
+from dataforge import PostOutputSession
+
+# Initialize a session (production use)
+session = PostOutputSession()
+
+# Initialize a session (optional names for testing)
+session = PostOutputSession(output_name="report", output_source_name="my_source", project_name="my_project")
+
+
+# Get the path of file generated by preceding output process
+path = session.file_path()
+
+# Retrieve connection parameters for the current output
+connection_parameters = session.connection_parameters()
+
+# Retrieve custom parameters for the current source
+custom_parameters = session.custom_parameters()
+
+# Run post-output logic: pass a function encapsulating custom code
+session.run(lambda: print(f"Uploading file from {path}"))
+
+# Fail the process with error message
+session.fail("Error message")
+```

dataforge_sdk-9.2.0rc3/README.md

@@ -0,0 +1,115 @@
+# dataforge-sdk
+SDK for creating DataForge extensions.
+
+## Postgres Utilities
+
+The `dataforge.pg` module provides helper functions to execute SQL operations against the DataForge Postgres metastore:
+
+```python
+from dataforge.pg import select, update, pull
+
+# Execute a SELECT query and return a Spark DataFrame
+df = select("SELECT * FROM my_table")
+
+# Execute an UPDATE/INSERT/DELETE query
+update("UPDATE my_table SET col = 'value'")
+
+# Trigger a new data pull for source_id 123
+pull(123)
+```
+
+## IngestionSession
+
+The `IngestionSession` class manages a custom data ingestion process lifecycle.
+
+```python
+from dataforge import IngestionSession
+
+# Initialize a session (production use)
+session = IngestionSession()
+
+# Initialize a session (optional source_name/project_name for testing)
+session = IngestionSession(source_name="my_source", project_name="my_project")
+
+# Ingest data 
+# pass a function returning a DataFrame (recommended to integrate logging with DataForge)
+session.ingest(lambda: spark.read.csv("s3://bucket/path/input.csv"))
+
+# pass a DataFrame (can be used for testing, not recommended for production deployment)
+df = spark.read.csv("s3://bucket/path/input.csv")
+session.ingest(df)
+
+# ingest empty dataframe to create 0-record input
+session.ingest()
+
+
+# Fail the process with error message
+session.fail("Error message")
+
+# Retrieve latest tracking fields
+tracking = session.latest_tracking_fields()
+
+# Retrieve connection parameters for the current source
+connection_parameters = session.connection_parameters()
+
+# Retrieve custom parameters for the current source
+custom_parameters = session.custom_parameters()
+
+```
+
+## ParsingSession
+
+The `ParsingSession` class manages a custom parse process lifecycle.
+
+```python
+from dataforge import ParsingSession
+
+# Initialize a session (production use)
+session = ParsingSession()
+
+# Initialize a session (optional input_id for testing)
+session = ParsingSession(input_id=123)
+
+# Retrieve custom parameters
+params = session.custom_parameters()
+
+# Get the path of file to be parsed
+path = session.file_path
+
+# Run parsing: pass a DataFrame, a function returning a DataFrame or None (0-record file)
+session.run(lambda: spark.read.json(session.file_path))
+
+# Fail the process with error message
+session.fail("Error message")
+
+```
+
+## PostOutputSession
+
+The `PostOutputSession` class manages a custom post-output process lifecycle.
+
+```python
+from dataforge import PostOutputSession
+
+# Initialize a session (production use)
+session = PostOutputSession()
+
+# Initialize a session (optional names for testing)
+session = PostOutputSession(output_name="report", output_source_name="my_source", project_name="my_project")
+
+
+# Get the path of file generated by preceding output process
+path = session.file_path()
+
+# Retrieve connection parameters for the current output
+connection_parameters = session.connection_parameters()
+
+# Retrieve custom parameters for the current source
+custom_parameters = session.custom_parameters()
+
+# Run post-output logic: pass a function encapsulating custom code
+session.run(lambda: print(f"Uploading file from {path}"))
+
+# Fail the process with error message
+session.fail("Error message")
+```

dataforge_sdk-9.2.0rc3/dataforge/__init__.py

@@ -0,0 +1,6 @@
+from .post_output_session import PostOutputSession
+from .ingestion_session import IngestionSession
+from .parsing_session import ParsingSession
+
+__version__ = "9.2.0-rc.3"
+__all__ = ['PostOutputSession','IngestionSession', 'ParsingSession']

dataforge_sdk-9.2.0rc3/dataforge/_session.py

@@ -0,0 +1,302 @@
+"""Session management for DataForge SDK.
+
+This module provides the Session class for managing DataForge SDK sessions,
+
+Classes:
+    CoreAPIError: Raised when a Core API call fails.
+    VersionMismatchError: Raised when the SDK version does not match server expectation.
+    Session: Main class for session management.
+"""
+import json
+import re
+import traceback
+from typing import Optional, Literal
+from pyspark.sql import SparkSession, DataFrame
+import requests
+import logging
+from pyspark.dbutils import DBUtils
+from pyspark.sql.functions import monotonically_increasing_id, lit
+from pyspark.sql.types import LongType
+
+from dataforge.process_record import ProcessRecord
+from dataforge.system_configuration import SystemConfiguration
+from dataforge.postgres_connection import PostgresConnection
+from importlib.metadata import version
+
+from .utils import _setup_logger
+
+
+class CoreAPIError(Exception):
+    """Exception raised when a Core API call fails."""
+    pass
+
+class VersionMismatchError(Exception):
+    """Exception raised when the SDK version does not match server expectation."""
+    pass
+
+class _Session:
+    """Base session class for DataForge SDK operations.
+    Class should not be instantiated by user directly: use process-specific Session classes instead
+    Manages Spark session, DBUtils, system configuration, Core API calls,
+    logging, and process lifecycle.
+
+    Attributes:
+        spark (SparkSession): Spark session for data processing.
+        dbutils (DBUtils): Databricks DBUtils instance.
+        logger (logging.Logger): Logger instance for session logs.
+        _systemConfiguration (SystemConfiguration): System configuration loaded from Postgres.
+        _pg (PostgresConnection): Postgres connection to core API database.
+        process (ProcessRecord): The current process record.
+        version (str): DataForge SDK version.
+    """
+    _systemConfiguration: SystemConfiguration
+    _pg: PostgresConnection
+    spark: SparkSession
+    dbutils: DBUtils
+    logger: logging.Logger
+    process: ProcessRecord
+    version = version("dataforge-sdk")
+    process_parameters: dict
+    _is_open: bool = False
+
+    def __init__(self):
+        self.logger = _setup_logger(self.__class__.__name__)
+        self.spark = SparkSession.builder.getOrCreate()
+        self.dbutils = self._get_dbutils()
+        pg_connection_string_read = self.dbutils.secrets.get("sparky", "pg_read")
+        pg_read = PostgresConnection(f"{pg_connection_string_read}&application_name=sdk", self.logger)
+        self._systemConfiguration = SystemConfiguration(pg_read)
+        pg_connection_string: str = self._core_api_call("core/sparky-db-connection")["connection"]
+        self._pg = PostgresConnection(f"{pg_connection_string}&application_name=sdk", self.logger)
+        self._is_open = True
+        self._check_version()
+        self.process_parameters = {
+            "version": self.version,
+            "packageName": "dataforge-sdk"
+        }
+        try:
+            self.process_parameters['process_id'] = int(self.dbutils.widgets.get("process_id"))
+        except Exception as e:
+            pass
+        self.logger.info(f"Initialized base session for {self.__class__.__name__}")
+
+
+    def _get_dbutils(self):
+        return DBUtils(self.spark)
+
+    def _core_api_call(self, route: str):
+        """Make a GET request to the Core API.
+
+        Args:
+            route (str): API route to call, appended to core URI.
+
+        Returns:
+            dict: Parsed JSON response from Core API.
+
+        Raises:
+            CoreAPIError: If response status is not 200.
+        """
+        add_core =  "core/" if route.startswith("process-complete") and self._systemConfiguration.saas_flag else ""
+        end_point = f"{self._systemConfiguration.coreUri}/{add_core}{route}"
+        core_jwt_token = self.dbutils.secrets.get("sparky", "coreJWT")
+        headers = {
+            "Authorization": f"Bearer {core_jwt_token}",
+            "Content-Type": "application/json",
+            "Accept": "application/json"
+        }
+        self.logger.debug(f"Executing core API call to {end_point}")
+        response = requests.get(end_point, headers=headers)
+        if response.status_code == 200:
+            data = response.json() if response.content else None
+            return data
+        else:
+            raise CoreAPIError(f"Core API call to {end_point} failed with status code {response.status_code}")
+
+    def _check_version(self):
+        """Checks that the SDK version is compatible with the server.
+
+        Calls the database function to verify the SDK version. Raises
+        VersionMismatchError if there is a version mismatch.
+        """
+        self.logger.info(f"Executing package dataforge-sdk version {self.version}")
+        check_res = self._pg.sql("select sparky.sdk_check_version(%s,%s)", [self.version, "dataforge-sdk"])
+        if check_res.get("error"):
+            raise VersionMismatchError(check_res.get("error"))
+
+    def log(self, message:str, severity: str = "I"):
+        """Log a message to both the local logger and the process log in the database.
+
+        Args:
+            message (str): The log message.
+            severity (str): Severity level, "I" for info, "E" for error. Defaults to "I".
+        """
+        match severity:
+            case "I":
+                self.logger.info(message)
+            case "E":
+                self.logger.error(message)
+        payload = {"log_id": self.process.logId, "operation_type": self.process.operationType,
+         "severity": severity, "message": message}
+        self._pg.sql("SELECT sparky.write_log(%s)", [json.dumps(payload)], fetch=False)
+
+    def connection_parameters(self):
+        """Retrieve connection parameters for the current process.
+
+        Returns:
+            dict: Dictionary containing private and public connection parameters.
+        """
+        connection_id = self.process.connectionId
+        self.log(f"Reading connection parameters of connection id {connection_id}")
+        params = self._core_api_call(f"core/connection/{connection_id}").get("parameters")
+        return {"private_connection_parameters": params.get("private_connection_parameters"),
+          "public_connection_parameters": params.get("public_connection_parameters")}
+
+
+    def _end_process(self, status_code: Literal['P', 'F', 'Z'] = 'P', parameters: Optional[dict] = None):
+        """Private method. End the current process
+        """
+        payload = parameters if parameters else {}
+        payload["process_id"] = self.process.processId
+        payload["status_code"] = status_code
+        self._pg.sql("select meta.prc_process_end(%s)", [json.dumps(payload)], fetch=False)
+        self._core_api_call(f"process-complete/{self.process.processId}")
+        if status_code in ('P','Z'):
+            self.logger.info("Session completed successfully")
+        else:
+            self.logger.error("Process failed")
+        self.close()
+
+    def fail(self,message: str):
+        """Fail current session and log the message.
+
+        Args:
+            message (str): The log message.
+        """
+        self.log(f"Process failed with error: {message}", "E")
+        self._end_process("F")
+
+    def _log_fail(self, e: Exception):
+        """Log failure for the given exception.
+
+        Args:
+            e (Exception): The exception that caused the failure.
+        """
+        traceback.print_exception(e)
+        self.log(f"Process failed with error: {e}", "E")
+
+    def close(self):
+        """Close the session."""
+        self._pg.close()
+        self.logger.info("Session closed")
+        self._is_open = False
+
+    def custom_parameters(self):
+        """Retrieve custom parameters from the process.
+
+        Returns:
+            dict: The custom parameters dictionary from the process.
+        """
+        return self.process.parameters.get('custom_parameters')
+
+    def _write_parsed_data(self, in_df: DataFrame, dest_file_path: str) -> tuple[int, int]:
+        """Process input DataFrame, write to Parquet, and update metadata.
+
+        Args:
+            in_df (DataFrame): Input Spark DataFrame to process and write.
+            dest_file_path (str): Destination path for saving Parquet file.
+
+        Returns:
+            tuple[int, int]: A tuple containing the total file size in bytes and the number of records written.
+
+        Raises:
+            Exception: If duplicate columns are detected or metadata update fails.
+        """
+        self.log("Data read successfully. Checking schema.")
+
+        select_list = self._pg.sql("SELECT sparky.get_select_list(%s)", (self.process.sourceId,))
+        df_sel = in_df.selectExpr(*select_list)
+        self.log(f"Applied select list {select_list}")
+
+        # Duplicate column check
+        cols = df_sel.columns
+        dup_columns = [col for col in set(cols) if cols.count(col) > 1]
+        if dup_columns:
+            raise Exception(f"Duplicate columns detected: {', '.join(dup_columns)}")
+
+        # Cast binary/void to string
+        binary_casts = [
+            f"CAST(`{f.name}` AS STRING) `{f.name}`" if f.dataType.typeName() in ("binary", "void")
+            else f"`{f.name}`"
+            for f in df_sel.schema.fields
+        ]
+        df = df_sel.selectExpr(*binary_casts)
+
+        # Schema as JSON array
+        schema = []
+        for f in df.schema.fields:
+            field_name = f.name.lower() if self.process.forceCaseInsensitive else f.name
+            name_normalized =  re.sub(r'\W+', '_', field_name)
+            column_normalized = ("_" if field_name[0].isdigit() else "") + name_normalized # add leading underscore
+
+            if f.dataType.simpleString().startswith("struct"):
+                spark_type = "StructType"
+            elif f.dataType.simpleString().startswith("array"):
+                spark_type = "ArrayType"
+            elif f.dataType.simpleString().startswith("decimal"):
+                spark_type = "DecimalType"
+            else:
+                spark_type = f.dataType.simpleString()
+
+            attr_schema = json.loads(f.dataType.json())
+            self.logger.info(f"Column `{column_normalized}` schema: {attr_schema}")
+            schema.append({
+                "name": field_name,
+                "column_normalized": column_normalized,
+                "spark_type": spark_type,
+                "schema": attr_schema
+            })
+
+        self.log("Schema read successfully. Updating source raw metadata.")
+
+        metadata_update_json = {
+            "source_id": self.process.sourceId,
+            "input_id": self.process.inputId,
+            "raw_attributes": schema,
+            "ingestion_type": "sparky"
+        }
+
+        result = self._pg.sql("SELECT meta.prc_n_normalize_raw_attribute(%s)", (json.dumps(metadata_update_json),))
+        if "error" in result:
+            raise Exception(result["error"])
+
+        normalize_attributes = result["normalized_metadata"]
+
+        self.log("Source metadata updated. Renaming and upcasting attributes")
+
+        cast_rename_expr = []
+        for att in normalize_attributes:
+            base_expr = att.get("upcastExpr") or att["raw_attribute_name"]
+            if att["raw_attribute_name"] != att["column_alias"] or self.process.forceCaseInsensitive:
+                base_expr += f" AS {att['column_alias']}"
+            cast_rename_expr.append(base_expr)
+
+        self.logger.info("Normalized SQL: " + ", ".join(cast_rename_expr))
+        df_update = df.selectExpr(*cast_rename_expr)
+
+        if self.process.parameters.get("generate_row_id", False):
+            self.logger.info("Added s_row_id to data.")
+            df_final = df_update.withColumn("s_row_id", monotonically_increasing_id())
+        else:
+            self.logger.info("generate_row_id = false, added null s_row_id.")
+            df_final = df_update.withColumn("s_row_id", lit(None).cast(LongType()))
+
+        self.log("Writing file")
+        df_final.write.format("parquet").mode("overwrite").save(dest_file_path)
+        self.log(f"Wrote file {dest_file_path}")
+
+        row_count = self.spark.read.format("parquet").load(dest_file_path).count()
+        self.log(f"{row_count} records counted")
+
+        file_size = sum(f.size for f in self.dbutils.fs.ls(dest_file_path))
+        return file_size, row_count
+

dataforge_sdk-9.2.0rc3/dataforge/ingestion_session.py

@@ -0,0 +1,119 @@
+"""IngestionSession management for DataForge SDK.
+
+This module provides the IngestionSession class
+that manages custom data ingestion process, including initialization, data writing,
+metadata updates, and failure handling.
+
+Classes:
+    IngestionSession: Manages the custom ingestion process lifecycle.
+"""
+import json
+from datetime import datetime
+from typing import Optional, Callable
+from pyspark.sql import DataFrame
+from .process_record import ProcessRecord
+from ._session import _Session
+
+
+class IngestionSession(_Session):
+
+    """Session class to manage custom ingestion process lifecycle.
+
+    Initializes an ingestion process record and provides methods to ingest
+    data frames and handle failures, updating metadata and notifying the Core API.
+    """
+    def __init__(self, source_name: Optional[str] = None, project_name: Optional[str] = None):
+        """Initialize ingestion session and start a new ingestion process.
+
+        Args:
+            source_name (Optional[str]): Optional name of the data source being ingested.
+                Used for interactive testing. Leave blank for production use.
+            project_name (Optional[str]): Optional project context name. Used for interactive testing.
+                Leave blank for production use.
+        """
+        super().__init__()
+        # initialize process
+        self.process_parameters["start_process_flag"] = True
+        if source_name is not None:
+            self.process_parameters["source_name"] = source_name
+        if project_name is not None:
+            self.process_parameters["project_name"] = project_name
+
+        process = self._pg.sql("select sparky.sdk_new_ingestion(%s)", (json.dumps(self.process_parameters),))
+        self.process = ProcessRecord(process)
+
+
+    def latest_tracking_fields(self):
+        """Get the latest tracking fields from the ingestion process record.
+
+        Returns:
+            dict: The latest tracking data from the process record, or None if not available.
+        """
+        return self.process.process.get("tracking")
+
+    def ingest(self,df: DataFrame | Callable[[], DataFrame] | None = None):
+        """Ingest the provided DataFrame into the DataForge and update input record.
+
+        Writes the DataFrame to raw Parquet file,
+        updates the input record with status, file size, record count, and notifies
+        the Core API of process completion. On failure, updates logs and flags the input and process
+        records as failed.
+
+        Args:
+            df (Callable[[], DataFrame] | DataFrame): parameterless def that you defined, returning the Spark DataFrame to ingest (recommended),
+                or spark DataFrame
+        """
+        try:
+            if not self._is_open:
+                raise Exception("Session is closed")
+            if df is None:
+                status = "Z"
+                row_count = 0
+                file_size = 0
+            else:
+                if callable(df):
+                    result_df = df()  # call it to get the DataFrame
+                else:
+                    result_df = df
+                dest_file_path = f"{self._systemConfiguration.datalakePath}/source_{self.process.sourceId}/parsed/parsed_input_{self.process.inputId}"
+                file_size, row_count = self._write_parsed_data(result_df, dest_file_path)
+                status = "P" if row_count > 0 else "Z"
+            input_update_json = {
+                "ingestion_status_code": status,
+                "extract_datetime": datetime.now().isoformat(),
+                "file_size": file_size,
+                "process_id": self.process.processId,
+                "input_id": self.process.inputId,
+                "record_counts": {"Total": row_count}
+            }
+
+            self._pg.sql("SELECT meta.prc_iw_in_update_input_record(%s)",
+                         (json.dumps(input_update_json),), fetch=False)
+            self.logger.info("Ingestion completed successfully")
+
+        except Exception as e:
+            self._log_fail(e)
+            failure_update_json = {
+                "process_id": self.process.processId,
+                "ingestion_status_code": "F"
+            }
+            self._pg.sql("SELECT meta.prc_iw_in_update_input_record(%s)",
+                         (json.dumps(failure_update_json),), fetch=False)
+        finally:
+            self._core_api_call(f"process-complete/{self.process.processId}")
+            self.close()
+
+
+    def fail(self, message: str):
+        """Mark the ingestion as failed with a custom message.
+
+        Logs the provided error message, updates the input and process record to failure status,
+        and notifies the Core API of process completion.
+
+        Args:
+            message (str): Custom error message explaining the failure.
+        """
+        self.log(f"Custom Ingestion failed with error: {message}", "E")
+        self._pg.sql("select meta.prc_iw_in_update_input_record(%s)",
+                     [json.dumps({"process_id": self.process.processId, "ingestion_status_code" : "F"})], fetch=False)
+        self._core_api_call(f"process-complete/{self.process.processId}")

dataforge_sdk-9.2.0rc3/dataforge/parsing_session.py

@@ -0,0 +1,82 @@
+
+import json
+from typing import Optional, Callable
+from pyspark.sql import DataFrame
+from .process_record import ProcessRecord
+from ._session import _Session
+
+
+class ParsingSession(_Session):
+
+    """Session class to manage custom parse process lifecycle.
+
+    Initializes parse process record and provides methods to parse
+    data frames and handle failures, updating metadata and notifying the Core API.
+    """
+    _parsing_parameters: dict
+    def __init__(self, input_id: Optional[int] = None):
+        """Initialize custom parse session and start a new custom parse process.
+
+        Args:
+            input_id (Optional[int]): Optional input_id of the batch for interactive testing.
+                Leave blank for production use.
+        """
+        super().__init__()
+        # initialize process
+        if input_id is not None:
+            self.process_parameters["input_id"] = input_id
+
+        process = self._pg.sql("select sparky.sdk_new_parse(%s)", (json.dumps(self.process_parameters),))
+        self.process = ProcessRecord(process)
+        self._parsing_parameters = self._pg.sql("select meta.prc_n_start_parse(%s)",
+                                                (json.dumps({"input_id": self.process.inputId}),))
+        # Extract the file extension
+        file_name = self._parsing_parameters["source_file_name"]
+        file_extension = file_name.split(".")[-1]
+        # Construct the path
+        self.file_path = f"{self._systemConfiguration.datalakePath}/source_{self.process.sourceId}/raw/raw_input_{self.process.inputId}.{file_extension}"
+
+    def custom_parameters(self):
+        """Retrieve custom parsing parameters.
+
+        Returns:
+            dict: The custom parsing parameters dictionary.
+        """
+        return self._parsing_parameters.get('custom_parameters')
+
+    def run(self,df: DataFrame | Callable[[], DataFrame] | None = None):
+        """Save parsed file from the provided DataFrame, and upload it into the DataForge data lake.
+
+        Writes the DataFrame to parsed Parquet file,
+        updates the input record with status, file size, record count, and notifies
+        the Core API of process completion. On failure, updates logs and flags the input and process
+        records as failed.
+
+        Args:
+            df (DataFrame): parameterless def that you defined, returning the Spark DataFrame containing parsed file data (recommended),
+                or spark DataFrame
+        """
+        try:
+            if not self._is_open:
+                raise Exception("Session is closed")
+            if callable(df):
+                result_df = df()  # call it to get the DataFrame
+            else:
+                result_df = df
+
+            if result_df is None or result_df.isEmpty():
+                file_size, row_count = (0, 0)
+            else:
+                dest_file_path = f"{self._systemConfiguration.datalakePath}/source_{self.process.sourceId}/parsed/parsed_input_{self.process.inputId}"
+                file_size, row_count = self._write_parsed_data(result_df, dest_file_path)
+            input_update_json = {
+                "file_size": file_size,
+                "input_id": self.process.inputId,
+                "record_counts": {"Total": row_count}
+            }
+            self._end_process('P' if row_count > 0 else 'Z', input_update_json)
+
+        except Exception as e:
+            self._log_fail(e)
+            self._end_process("F")
+

dataforge_sdk-9.2.0rc3/dataforge/pg.py

@@ -0,0 +1,91 @@
+"""Postgres utilities for data operations.
+
+This module provides functions to execute SQL queries against a Postgres database
+using Spark JDBC for reads and a direct write connection for write operations.
+"""
+from pyspark.dbutils import DBUtils
+from pyspark.sql import SparkSession, DataFrame
+from .postgres_connection import PostgresConnection
+
+spark = SparkSession.builder.getOrCreate()
+dbutils = DBUtils(spark)
+pg_connection_string_read = dbutils.secrets.get("sparky", "pg_read")
+
+
+def update(query: str):
+    """Execute an update SQL query on the DataForge metastore Postgres database.
+
+    Args:
+        query (str): SQL query string to execute.
+
+    Returns:
+        None
+
+    Raises:
+        Exception: If write connection cannot be established or SQL execution fails.
+    """
+    pg = _get_pg_write_connection()
+    pg.sql(query, fetch=False)
+
+
+def execute(query: str):
+    """Alias for update() to execute write SQL queries.
+
+    Args:
+        query (str): SQL query string to execute.
+
+    Returns:
+        None
+    """
+    update(query)
+
+
+def select(query: str) -> DataFrame:
+    """Execute a SELECT SQL query on the DataForge metastore Postgres database and return a DataFrame with results.
+
+    Args:
+        query (str): SQL SELECT query string.
+
+    Returns:
+        DataFrame: Spark DataFrame containing query results.
+
+    Raises:
+        Exception: If Spark fails to load data or connection issues.
+    """
+    return spark.read.format("jdbc") \
+        .option("url", pg_connection_string_read) \
+        .option("query", query) \
+        .load()
+
+
+def pull(source_id: int):
+    """Trigger new ingestion (pull data) on DataForge source for a given source ID.
+
+    Args:
+        source_id (int): Identifier for the source to pull.
+
+    Returns:
+        None
+
+    Raises:
+        Exception: If write connection cannot be established or SQL execution fails.
+    """
+    pg = _get_pg_write_connection()
+    pg.sql("SELECT meta.svc_pull_source(%s, %s)", (source_id,'sdk'), fetch=False)
+
+
+def _get_pg_write_connection() -> PostgresConnection:
+    """Internal method to retrieve a PostgresConnection for write operations using secured secrets.
+
+    Returns:
+        PostgresConnection: Connection object for executing write queries.
+
+    Raises:
+        Exception: If the 'pg_write' secret is not defined in the 'sparky' scope.
+    """
+    secrets = dbutils.secrets.list("sparky")
+    if any(secret.key == "pg_write" for secret in secrets):
+        conn_string = dbutils.secrets.get("sparky", "pg_write")
+        return PostgresConnection(conn_string + "&application_name=sdk-pg")
+    else:
+        raise Exception("pg_write secret is not defined in sparky scope")

dataforge_sdk-9.2.0rc3/dataforge/post_output_session.py

@@ -0,0 +1,60 @@
+
+import json
+from typing import Optional, Callable
+from .process_record import ProcessRecord
+from ._session import _Session
+
+
+class PostOutputSession(_Session):
+
+    """Session class to manage custom post-output process lifecycle.
+
+    Initializes post-output process record and provides run() method to execute custom logic post-output completion,
+     handle failures and update metadata.
+    """
+    def __init__(self, output_name: Optional[str] = None, output_source_name: Optional[str] = None, project_name: Optional[str] = None):
+        """Initialize custom post-output session and start a new post-output process.
+
+        Args:
+            input_id (Optional[int]): Optional input_id of the batch for interactive testing.
+                Leave blank for production use.
+        """
+        super().__init__()
+        # initialize process
+        if output_name is not None:
+            self.process_parameters["output_name"] = output_name
+            self.process_parameters["output_source_name"] = output_source_name
+            self.process_parameters["project_name"] = project_name
+
+        process = self._pg.sql("select sparky.sdk_new_post_output(%s)", (json.dumps(self.process_parameters),))
+        self.process = ProcessRecord(process)
+
+    def file_path(self):
+        """Returns file path for the file generated by output process"""
+        return self.process.parameters.get("file_path")
+
+    def run(self, po_commands: Callable[[],None]):
+
+        """Save parsed file from the provided DataFrame, and upload it into the DataForge data lake.
+
+        Writes the DataFrame to parsed Parquet file,
+        updates the input record with status, file size, record count, and notifies
+        the Core API of process completion. On failure, updates logs and flags the input and process
+        records as failed.
+
+        Args:
+            po_commands (po_commands: Callable[[],None]): The custom-defined, parameterless def containing post-output logic.
+            Returns None
+        """
+        try:
+            if not self._is_open:
+                raise Exception("Session is closed")
+            self.log(f"Running post output def {po_commands}")
+            po_commands()
+            self.log(f"Post output execution completed successfully")
+            self._end_process("P")
+
+        except Exception as e:
+            self._log_fail(e)
+            self._end_process("F")
+

dataforge_sdk-9.2.0rc3/dataforge/postgres_connection.py

@@ -0,0 +1,49 @@
+import logging
+import sys
+import psycopg2
+from .utils import _setup_logger
+
+
+class PostgresConnection:
+    logger: logging.Logger
+    def __init__(self, connection_string: str, logger: logging.Logger | None = None):
+        try:
+            if logger:
+                self.logger = logger
+            else:
+                self.logger = _setup_logger("PG")
+            self.connect(connection_string)
+
+        except Exception as e:
+            logger.error(f"Error connecting to Postgres: {e}")
+            sys.exit(1)
+
+    def sql(self, query: str, params=None, fetch=True):
+        try:
+            # Execute a query
+            cur = self.conn.cursor()
+            cur.execute(query, params)
+            # Retrieve query results
+            res = cur.fetchone() if fetch else [None]
+            cur.close()
+            return res[0]
+        except Exception as e:
+            self.logger.error(f"Error executing query {query}({params}) on Postgres: {e}")
+            sys.exit(1)
+
+    def connect(self, connection_string: str):
+        # Execute a query
+        try:
+            self.conn = psycopg2.connect(connection_string.removeprefix('jdbc:'))
+            self.conn.set_session(autocommit=True)
+            self.sql("select 1")  # execute test query
+            # Change connection
+        except Exception as e:
+            self.logger.error(f"Error connecting to Postgres database or insufficient permissions. Details: {e}")
+            sys.exit(1)
+
+    def close(self ):
+        self.conn.close()
+
+
+

dataforge_sdk-9.2.0rc3/dataforge/process_record.py

@@ -0,0 +1,21 @@
+class ProcessRecord:
+    processId: int
+    logId: int
+    operationType: str
+    sourceId: int
+    inputId: int
+    process: dict
+    parameters: dict
+    forceCaseInsensitive: bool
+    connectionId: int
+
+    def __init__(self, process):
+        self.process = process
+        self.processId = process['process_id']
+        self.logId = process['log_id']
+        self.operationType = process['operation_type']
+        self.inputId = process['input_id']
+        self.sourceId = process['source_id']
+        self.connectionId = process.get('connection_id')
+        self.parameters = process['parameters']
+        self.forceCaseInsensitive = self.parameters.get('force_case_insensitive', False)

dataforge_sdk-9.2.0rc3/dataforge/system_configuration.py

@@ -0,0 +1,22 @@
+from .postgres_connection import PostgresConnection
+
+
+class SystemConfiguration:
+    datalakePath: str
+    databricksDbName: str
+    targetParquetFileSize: int
+    # etlURI: str
+    architecture:str
+    coreUri: str
+    saas_flag: bool
+
+    def __init__(self, pg: PostgresConnection):
+        conf = pg.sql("SELECT meta.prc_get_system_configuration('sparky')")
+        self.datalakePath = conf['data-lake-path'].replace("s3://", "s3a://")
+        self.databricksDbName = conf['databricks-db-name']
+        self.targetParquetFileSize = conf['target-parquet-file-size']
+        # self.etlURL = conf['etl-url']
+        self.architecture = conf['architecture']
+        self.saas_flag = self.architecture == "saas"
+        self.coreUri = f"https://{conf['api-url']}" if self.saas_flag else f"http://{conf['etl-url']}:7131"
+

dataforge_sdk-9.2.0rc3/dataforge/utils.py

@@ -0,0 +1,25 @@
+import logging
+
+def _setup_logger(name: str = "SDK"):
+    """Set up and return a configured logger for the SDK.
+
+    Returns:
+        logging.Logger: Configured logger for SDK outputs.
+    """
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO)
+    logger.propagate = False
+
+    # Create handler
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.INFO)
+
+    # Create formatter with timestamp
+    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(name)s - %(message)s')
+    console_handler.setFormatter(formatter)
+
+    if not logger.hasHandlers():
+        # Add handler to the logger
+        logger.addHandler(console_handler)
+
+    return logger

dataforge_sdk-9.2.0rc3/dataforge_sdk.egg-info/PKG-INFO

@@ -0,0 +1,125 @@
+Metadata-Version: 2.4
+Name: dataforge-sdk
+Version: 9.2.0rc3
+Summary: SDK for creating DataForge extensions
+Author-email: Vadim Orlov <vorlov@dataforgelabs.com>
+Project-URL: Homepage, https://docs.dataforgelabs.com
+Project-URL: Issues, https://docs.dataforgelabs.com/hc/en-us/requests/new
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# dataforge-sdk
+SDK for creating DataForge extensions.
+
+## Postgres Utilities
+
+The `dataforge.pg` module provides helper functions to execute SQL operations against the DataForge Postgres metastore:
+
+```python
+from dataforge.pg import select, update, pull
+
+# Execute a SELECT query and return a Spark DataFrame
+df = select("SELECT * FROM my_table")
+
+# Execute an UPDATE/INSERT/DELETE query
+update("UPDATE my_table SET col = 'value'")
+
+# Trigger a new data pull for source_id 123
+pull(123)
+```
+
+## IngestionSession
+
+The `IngestionSession` class manages a custom data ingestion process lifecycle.
+
+```python
+from dataforge import IngestionSession
+
+# Initialize a session (production use)
+session = IngestionSession()
+
+# Initialize a session (optional source_name/project_name for testing)
+session = IngestionSession(source_name="my_source", project_name="my_project")
+
+# Ingest data 
+# pass a function returning a DataFrame (recommended to integrate logging with DataForge)
+session.ingest(lambda: spark.read.csv("s3://bucket/path/input.csv"))
+
+# pass a DataFrame (can be used for testing, not recommended for production deployment)
+df = spark.read.csv("s3://bucket/path/input.csv")
+session.ingest(df)
+
+# ingest empty dataframe to create 0-record input
+session.ingest()
+
+
+# Fail the process with error message
+session.fail("Error message")
+
+# Retrieve latest tracking fields
+tracking = session.latest_tracking_fields()
+
+# Retrieve connection parameters for the current source
+connection_parameters = session.connection_parameters()
+
+# Retrieve custom parameters for the current source
+custom_parameters = session.custom_parameters()
+
+```
+
+## ParsingSession
+
+The `ParsingSession` class manages a custom parse process lifecycle.
+
+```python
+from dataforge import ParsingSession
+
+# Initialize a session (production use)
+session = ParsingSession()
+
+# Initialize a session (optional input_id for testing)
+session = ParsingSession(input_id=123)
+
+# Retrieve custom parameters
+params = session.custom_parameters()
+
+# Get the path of file to be parsed
+path = session.file_path
+
+# Run parsing: pass a DataFrame, a function returning a DataFrame or None (0-record file)
+session.run(lambda: spark.read.json(session.file_path))
+
+# Fail the process with error message
+session.fail("Error message")
+
+```
+
+## PostOutputSession
+
+The `PostOutputSession` class manages a custom post-output process lifecycle.
+
+```python
+from dataforge import PostOutputSession
+
+# Initialize a session (production use)
+session = PostOutputSession()
+
+# Initialize a session (optional names for testing)
+session = PostOutputSession(output_name="report", output_source_name="my_source", project_name="my_project")
+
+
+# Get the path of file generated by preceding output process
+path = session.file_path()
+
+# Retrieve connection parameters for the current output
+connection_parameters = session.connection_parameters()
+
+# Retrieve custom parameters for the current source
+custom_parameters = session.custom_parameters()
+
+# Run post-output logic: pass a function encapsulating custom code
+session.run(lambda: print(f"Uploading file from {path}"))
+
+# Fail the process with error message
+session.fail("Error message")
+```

dataforge_sdk-9.2.0rc3/dataforge_sdk.egg-info/SOURCES.txt

@@ -0,0 +1,17 @@
+README.md
+pyproject.toml
+setup.cfg
+dataforge/__init__.py
+dataforge/_session.py
+dataforge/ingestion_session.py
+dataforge/parsing_session.py
+dataforge/pg.py
+dataforge/post_output_session.py
+dataforge/postgres_connection.py
+dataforge/process_record.py
+dataforge/system_configuration.py
+dataforge/utils.py
+dataforge_sdk.egg-info/PKG-INFO
+dataforge_sdk.egg-info/SOURCES.txt
+dataforge_sdk.egg-info/dependency_links.txt
+dataforge_sdk.egg-info/top_level.txt

dataforge_sdk-9.2.0rc3/dataforge_sdk.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

dataforge_sdk-9.2.0rc3/dataforge_sdk.egg-info/top_level.txt

@@ -0,0 +1 @@
+dataforge

dataforge_sdk-9.2.0rc3/pyproject.toml

@@ -0,0 +1,20 @@
+[build-system]
+requires = [
+    "setuptools>=61.0",
+    "wheel"
+]
+build-backend = "setuptools.build_meta"
+[project]
+name = "dataforge-sdk"
+version = "9.2.0-rc.3"
+authors = [
+    {name="Vadim Orlov", email="vorlov@dataforgelabs.com"}
+]
+description = "SDK for creating DataForge extensions"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    ]
+[project.urls]
+Homepage = "https://docs.dataforgelabs.com"
+Issues = "https://docs.dataforgelabs.com/hc/en-us/requests/new"

dataforge_sdk-9.2.0rc3/setup.cfg

@@ -0,0 +1,22 @@
+[metadata]
+name = sdk-python
+version = attr: dataforge.__version__
+description = SDK for creating DataForge extensions
+long_description = file: README.md
+long_description_content_type = text/markdown
+classifiers = 
+	Programming Language :: Python :: 3
+	Programming Language :: Python :: 3.10
+	Operating System :: OS Independent
+
+[options]
+packages = find:
+install_requires = 
+	pyspark>=3.0
+	psycopg2-binary>=2.9
+python_requires = >=3.10
+
+[egg_info]
+tag_build = 
+tag_date = 0
+