dataforge-sdk 10.0.dev115__tar.gz → 10.0.0rc3__tar.gz

{dataforge_sdk-10.0.dev115 → dataforge_sdk-10.0.0rc3}/PKG-INFO

@@ -1,12 +1,14 @@
 Metadata-Version: 2.4
 Name: dataforge-sdk
-Version: 10.0.dev115
+Version: 10.0.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.10
 Description-Content-Type: text/markdown
+Provides-Extra: psycopg2
+Requires-Dist: psycopg2-binary>=2.9; extra == "psycopg2"
 
 # dataforge-sdk
 SDK for creating DataForge extensions.

{dataforge_sdk-10.0.dev115 → dataforge_sdk-10.0.0rc3}/dataforge/__init__.py

@@ -2,5 +2,5 @@ from .post_output_session import PostOutputSession
 from .ingestion_session import IngestionSession
 from .parsing_session import ParsingSession
 
-__version__ = "10.0.dev115"
+__version__ = "10.0.0-rc.3"
 __all__ = ['PostOutputSession','IngestionSession', 'ParsingSession']

dataforge_sdk-10.0.0rc3/dataforge/_session.py

@@ -0,0 +1,40 @@
+import os
+
+
+def _is_databricks_environment() -> bool:
+    """Detect whether the current runtime is Databricks."""
+    spark_obj = globals().get("spark")
+    if spark_obj is not None:
+        spark_class_name = getattr(getattr(spark_obj, "__class__", None), "__name__", None)
+        if spark_class_name == "SparkSession":
+            return True
+    if os.environ.get("DATABRICKS_RUNTIME_VERSION"):
+        return True
+    try:
+        from pyspark.sql import SparkSession  # type: ignore
+    except ImportError:
+        return False
+    try:
+        return SparkSession.getActiveSession() is not None
+    except Exception:
+        return False
+
+
+if _is_databricks_environment():
+    from dataforge.databricks._databricks_base_session import _Databricks_Base_Session
+    from dataforge.databricks._databricks_ingestion_session import _Databricks_Ingestion_Session
+    from dataforge.databricks._databricks_parsing_session import _Databricks_Parsing_Session
+
+    _Session = _Databricks_Base_Session
+    _Ingestion_Session = _Databricks_Ingestion_Session
+    _Parsing_Session = _Databricks_Parsing_Session
+    _platform = "databricks"
+else:
+    from dataforge.snowflake._snowflake_base_session import _Snowflake_Base_Session
+    from dataforge.snowflake._snowflake_ingestion_session import _Snowflake_Ingestion_Session
+    from dataforge.snowflake._snowflake_parsing_session import _Snowflake_Parsing_Session
+
+    _Session = _Snowflake_Base_Session
+    _Ingestion_Session = _Snowflake_Ingestion_Session
+    _Parsing_Session = _Snowflake_Parsing_Session
+    _platform = "snowflake"

dataforge_sdk-10.0.dev115/dataforge/_databricks_session.py → dataforge_sdk-10.0.0rc3/dataforge/databricks/_databricks_base_session.py

@@ -1,24 +1,22 @@
+from __future__ import annotations
 import json
 import re
-from typing import Callable
 
-from pyspark.dbutils import DBUtils
-from pyspark.sql import SparkSession, DataFrame
-from pyspark.sql.functions import monotonically_increasing_id, lit
-from pyspark.sql.types import LongType
 
 from dataforge._base_session import _Base_Session
 
 
-class _Databricks_Session(_Base_Session):
+class _Databricks_Base_Session(_Base_Session):
     """Base session class for Databricks platform.
     Class should not be instantiated by user directly: use process-specific Session classes instead
     Adds Spark session, DBUtilsto Base_Session
     """
-    spark: SparkSession
-    dbutils: DBUtils
+
 
     def __init__(self):
+        from pyspark.sql import SparkSession, DataFrame
+        self.spark = SparkSession.builder.getOrCreate()
+        self.dbutils = self._get_dbutils()
         pg_connection_string_read = self.dbutils.secrets.get("sparky", "pg_read")
         core_jwt_token = self.dbutils.secrets.get("sparky", "coreJWT")
         try:
@@ -27,69 +25,17 @@ class _Databricks_Session(_Base_Session):
             process_id = None
 
         super().__init__(pg_connection_string_read, core_jwt_token, process_id)
-        self.spark = SparkSession.builder.getOrCreate()
-        self.dbutils = self._get_dbutils()
         self.process_parameters["start_process_flag"] = True
         self.logger.info(f"Initialized databricks base session for {self.__class__.__name__} with parameters {self.process_parameters}")
 
 
     def _get_dbutils(self):
+        from pyspark.dbutils import DBUtils
         return DBUtils(self.spark)
 
 
-    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 _write_parsed_data(self, in_df: DataFrame, dest_file_path: str) -> tuple[int, int]:
+    def _write_parsed_data(self, in_df: pyspark.sql.DataFrame, dest_file_path: str) -> tuple[int, int]:
         """Process input DataFrame, write to Parquet, and update metadata.
 
         Args:
@@ -102,6 +48,10 @@ class _Databricks_Session(_Base_Session):
         Raises:
             Exception: If duplicate columns are detected or metadata update fails.
         """
+        from pyspark.sql.functions import monotonically_increasing_id, lit
+        from pyspark.sql.types import LongType
+
+
         self.log("Data read successfully. Checking schema.")
 
         select_list = self._pg.sql("SELECT sparky.get_select_list(%s)", (self.process.sourceId,))
@@ -126,8 +76,6 @@ class _Databricks_Session(_Base_Session):
         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"
@@ -139,10 +87,9 @@ class _Databricks_Session(_Base_Session):
                 spark_type = type(f.dataType).__name__
 
             attr_schema = json.loads(f.dataType.json())
-            self.logger.info(f"Column `{column_normalized}` schema: {attr_schema}")
+            self.logger.info(f"Column `{field_name}` schema: {attr_schema}")
             schema.append({
                 "name": field_name,
-                "column_normalized": column_normalized,
                 "spark_type": spark_type,
                 "schema": attr_schema
             })

dataforge_sdk-10.0.0rc3/dataforge/databricks/_databricks_ingestion_session.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+import json
+import re
+from datetime import datetime
+from typing import Callable
+from dataforge.databricks._databricks_base_session import _Databricks_Base_Session
+
+
+class _Databricks_Ingestion_Session(_Databricks_Base_Session):
+    """Base session class for Databricks platform.
+    Class should not be instantiated by user directly: use process-specific Session classes instead
+    Adds Spark session, DBUtilsto Base_Session
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def ingest(self,df: pyspark.sql.DataFrame | Callable[[], pyspark.sql.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()
+
+

dataforge_sdk-10.0.0rc3/dataforge/databricks/_databricks_parsing_session.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+from typing import Callable
+from ._databricks_base_session import _Databricks_Base_Session
+
+
+class _Databricks_Parsing_Session(_Databricks_Base_Session):
+    """Implements run method for Databricks.
+    """
+
+
+    def run(self,df: pyspark.sql.DataFrame | Callable[[], pyspark.sql.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-10.0.0rc3/dataforge/databricks/_databricks_pg.py

@@ -0,0 +1,94 @@
+"""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 dataforge.postgres_connection import PostgresConnection
+
+
+class DataBricksPg:
+
+
+    def __init__(self):
+        from pyspark.dbutils import DBUtils
+        from pyspark.sql import SparkSession, DataFrame
+        from .postgres_connection import PostgresConnection
+
+        self.spark = SparkSession.builder.getOrCreate()
+        self.dbutils = DBUtils(spark)
+        self.pg_connection_string_read = dbutils.secrets.get("sparky", "pg_read")
+
+    def update(self,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 = self._get_pg_write_connection()
+        pg.sql(query, fetch=False)
+
+    def execute(self,query: str):
+        """Alias for update() to execute write SQL queries.
+
+        Args:
+            query (str): SQL query string to execute.
+
+        Returns:
+            None
+        """
+        self.update(query)
+
+    def select(self,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 self.spark.read.format("jdbc") \
+            .option("url", self.pg_connection_string_read) \
+            .option("query", query) \
+            .load()
+
+    def pull(self,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 = self._get_pg_write_connection()
+        pg.sql("SELECT meta.svc_pull_source(%s, %s)", (source_id,'sdk'), fetch=False)
+
+
+    def _get_pg_write_connection(self) -> 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 = self.dbutils.secrets.list("sparky")
+        if any(secret.key == "pg_write" for secret in secrets):
+            conn_string = self.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-10.0.dev115 → dataforge_sdk-10.0.0rc3}/dataforge/ingestion_session.py

@@ -9,11 +9,11 @@ Classes:
 """
 import json
 from typing import Optional
-from ._session import _Session
+from ._session import _Ingestion_Session
 from .process_record import ProcessRecord
 
 
-class IngestionSession(_Session):
+class IngestionSession(_Ingestion_Session):
 
     """Session class to manage custom ingestion process lifecycle.
 

{dataforge_sdk-10.0.dev115 → dataforge_sdk-10.0.0rc3}/dataforge/parsing_session.py

@@ -1,12 +1,13 @@
+from __future__ import annotations
 
 import json
 from typing import Optional, Callable
-from pyspark.sql import DataFrame
 from .process_record import ProcessRecord
-from ._session import _Session
+from ._session import _Parsing_Session
 
 
-class ParsingSession(_Session):
+class ParsingSession(_Parsing_Session):
+
 
     """Session class to manage custom parse process lifecycle.
 
@@ -44,39 +45,5 @@ class ParsingSession(_Session):
         """
         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-10.0.0rc3/dataforge/pg.py

@@ -0,0 +1,10 @@
+"""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 dataforge._session import _platform
+
+if _platform=='databricks':
+    from dataforge.databricks._databricks_pg import DataBricksPg
+    pg = DataBricksPg()

{dataforge_sdk-10.0.dev115 → dataforge_sdk-10.0.0rc3}/dataforge/post_output_session.py

@@ -16,7 +16,6 @@ class PostOutputSession(_Session):
         """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__()

{dataforge_sdk-10.0.dev115 → dataforge_sdk-10.0.0rc3}/dataforge/postgres_connection.py

@@ -16,7 +16,8 @@ class PostgresConnection:
 
         except Exception as e:
             logger.error(f"Error connecting to Postgres: {e}")
-            sys.exit(1)
+            raise
+            # sys.exit(1)
 
     def sql(self, query: str, params=None, fetch=True):
         try:
@@ -29,7 +30,8 @@ class PostgresConnection:
             return res[0]
         except Exception as e:
             self.logger.error(f"Error executing query {query}({params}) on Postgres: {e}")
-            sys.exit(1)
+            # sys.exit(1)
+            raise
 
     def connect(self, connection_string: str):
         # Execute a query
@@ -40,7 +42,8 @@ class PostgresConnection:
             # Change connection
         except Exception as e:
             self.logger.error(f"Error connecting to Postgres database or insufficient permissions. Details: {e}")
-            sys.exit(1)
+            # sys.exit(1)
+            raise
 
     def close(self ):
         self.conn.close()

dataforge_sdk-10.0.0rc3/dataforge/snowflake/_snowflake_base_session.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+from typing import Callable
+from dataforge._base_session import _Base_Session
+import sys
+
+class _Snowflake_Base_Session(_Base_Session):
+    """Base session class for Snowflake platform.
+    Class should not be instantiated by user directly: use process-specific Session classes instead
+    Adds Snowpark session
+    """
+
+    def __init__(self):
+        from snowflake.snowpark.context import get_active_session
+        self.snowpark_session = get_active_session()
+        pg_connection_string_read = self.snowpark_session.sql("SELECT get_secret('DATAFORGE_PG_READ')").first(1)[0][0]
+        core_jwt_token = self.snowpark_session.sql("SELECT get_secret('DATAFORGE_CORE_JWT')").first(1)[0][0]
+        params = self.parse_key_value_args()
+        process_id = params.get('process_id')
+        self.input_id = params.get('input_id')
+
+        super().__init__(pg_connection_string_read, core_jwt_token, process_id)
+
+        self.process_parameters["start_process_flag"] = process_id is None
+
+        self.logger.info(f"Initialized Snowflake base session for {self.__class__.__name__} with parameters {self.process_parameters}")
+
+    def _write_input_table(self,df: snowflake.snowpark.dataframe.DataFrame | Callable[[],
+        snowflake.snowpark.dataframeDataFrame] | None = None) -> snowflake.snowpark.dataframe.DataFrame:
+        from snowflake.snowpark.types import StructType, StructField, StringType
+        if not self._is_open:
+            raise Exception("Session is closed")
+        if df is None:
+            # create empty df
+            result_df = df = self.snowpark_session.create_dataframe([], StructType([StructField("id", StringType())]))
+        else:
+            if callable(df):
+                result_df = df()  # call it to get the DataFrame
+            else:
+                result_df = df
+        table = f"{self._systemConfiguration.dataLakeDbName}.{self._systemConfiguration.dataLakeSchemaName}.RAW_INPUT_{self.process.inputId}"
+        self.log(f"Writing dataframe to table {table}")
+        result_df.write.save_as_table(
+            table_name=table,
+            mode="overwrite",
+            table_type="transient"
+        )
+        self.log(f"Table {table} written")
+        if self.process.startProcessFlag:
+            # process started by IngestionSession, tell Core to continue and not run Notebook
+            self._pg.sql("SELECT sparky.sdk_complete_manual_process(%s)", [self.process.processId], fetch=False)
+
+    @staticmethod
+    def parse_key_value_args():
+        """
+        Parse command line arguments formatted as key=value into a dict.
+        Example: python script.py foo=123 bar=hello
+
+        Returns: {'foo': '123', 'bar': 'hello'}
+        """
+        argv = sys.argv
+        params: dict[str,str] = {}
+        for arg in argv:
+            if "=" in arg:
+                key, value = arg.split("=", 1)  # split only on first '='
+                params[key] = value
+        return params

dataforge_sdk-10.0.0rc3/dataforge/snowflake/_snowflake_ingestion_session.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+import json
+from typing import Callable
+from dataforge.snowflake._snowflake_base_session import _Snowflake_Base_Session
+
+
+class _Snowflake_Ingestion_Session(_Snowflake_Base_Session):
+    """Base ingestion session class for Snowflake platform.
+    Class should not be instantiated by user directly: use process-specific Session classes instead
+    Adds Snowpark session
+    """
+
+    def __init__(self):
+        super().__init__()
+
+
+    def ingest(self,df: snowflake.snowpark.dataframe.DataFrame | Callable[[], snowflake.snowpark.dataframeDataFrame] | None = None):
+        """Ingest the provided DataFrame into the DataForge and update input record.
+
+        Writes the DataFrame to raw Snowflake table
+
+        Args:
+            df (Callable[[], DataFrame] | DataFrame): parameterless def that you defined, returning the Snowpark DataFrame to ingest (recommended),
+                or spark DataFrame
+        """
+        try:
+            self._write_input_table(df)
+        except Exception as e:
+            self._log_fail(e)
+            if self.process.startProcessFlag:
+                # Fail input and process to prevent core from executing it
+                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.close()

dataforge_sdk-10.0.0rc3/dataforge/snowflake/_snowflake_parsing_session.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+from typing import Callable
+from ._snowflake_base_session import _Snowflake_Base_Session
+
+
+class _Snowflake_Parsing_Session(_Snowflake_Base_Session):
+    """Implements run method for Databricks.
+    """
+
+
+    def run(self,df: snowflake.snowpark.dataframe.DataFrame | Callable[[], snowflake.snowpark.dataframeDataFrame] | 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:
+            self._write_input_table(df)
+
+        except Exception as e:
+            self._log_fail(e)
+            if self.process.startProcessFlag:
+                self._end_process("F")
+

{dataforge_sdk-10.0.dev115 → dataforge_sdk-10.0.0rc3}/dataforge_sdk.egg-info/PKG-INFO

@@ -1,12 +1,14 @@
 Metadata-Version: 2.4
 Name: dataforge-sdk
-Version: 10.0.dev115
+Version: 10.0.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.10
 Description-Content-Type: text/markdown
+Provides-Extra: psycopg2
+Requires-Dist: psycopg2-binary>=2.9; extra == "psycopg2"
 
 # dataforge-sdk
 SDK for creating DataForge extensions.

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

@@ -0,0 +1,28 @@
+README.md
+pyproject.toml
+setup.cfg
+dataforge/__init__.py
+dataforge/_base_session.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/databricks/__init__.py
+dataforge/databricks/_databricks_base_session.py
+dataforge/databricks/_databricks_ingestion_session.py
+dataforge/databricks/_databricks_parsing_session.py
+dataforge/databricks/_databricks_pg.py
+dataforge/snowflake/__init__.py
+dataforge/snowflake/_snowflake_base_session.py
+dataforge/snowflake/_snowflake_ingestion_session.py
+dataforge/snowflake/_snowflake_parsing_session.py
+dataforge_sdk.egg-info/PKG-INFO
+dataforge_sdk.egg-info/SOURCES.txt
+dataforge_sdk.egg-info/dependency_links.txt
+dataforge_sdk.egg-info/requires.txt
+dataforge_sdk.egg-info/top_level.txt

dataforge_sdk-10.0.0rc3/dataforge_sdk.egg-info/requires.txt

@@ -0,0 +1,3 @@
+
+[psycopg2]
+psycopg2-binary>=2.9

{dataforge_sdk-10.0.dev115 → dataforge_sdk-10.0.0rc3}/pyproject.toml

@@ -6,15 +6,15 @@ requires = [
 build-backend = "setuptools.build_meta"
 [project]
 name = "dataforge-sdk"
-version = "10.0.dev115"
+version = "10.0.0-rc.3"
 authors = [
     {name="Vadim Orlov", email="vorlov@dataforgelabs.com"}
 ]
 description = "SDK for creating DataForge extensions"
 readme = "README.md"
 requires-python = ">=3.10"
-dependencies = [
-    ]
+[project.optional-dependencies]
+psycopg2 = ["psycopg2-binary>=2.9"]
 [project.urls]
 Homepage = "https://docs.dataforgelabs.com"
 Issues = "https://docs.dataforgelabs.com/hc/en-us/requests/new"

dataforge_sdk-10.0.dev115/dataforge/_session.py

@@ -1,7 +0,0 @@
-# Pick base class at import time
-if "spark" in globals() and type(spark).__name__ == 'SparkSession':
-    from dataforge._databricks_session import _Databricks_Session
-    _Session = _Databricks_Session
-else:
-    from dataforge._snowflake_session import _Snowflake_Session
-    _Session = _Snowflake_Session

dataforge_sdk-10.0.dev115/dataforge/_snowflake_session.py

@@ -1,82 +0,0 @@
-import json
-from typing import Callable
-
-from dataforge._base_session import _Base_Session
-import sys
-from snowflake.snowpark.context import get_active_session
-import streamlit as st
-from snowflake.snowpark.dataframe import DataFrame
-
-class _Snowflake_Session(_Base_Session):
-    """Base session class for Snowflake platform.
-    Class should not be instantiated by user directly: use process-specific Session classes instead
-    Adds Snowpark session
-    """
-
-    def __init__(self):
-        pg_connection_string_read = st.secrets['DATAFORGE_PG_READ']
-        core_jwt_token = st.secrets['DATAFORGE_CORE_JWT']
-        params = self.parse_key_value_args()
-        process_id = params.get('process_id')
-        self.input_id = params.get('input_id')
-
-        super().__init__(pg_connection_string_read, core_jwt_token, process_id)
-        self.snowpark_session = get_active_session()
-        self.process_parameters["start_process_flag"] = process_id is None
-
-        self.logger.info(f"Initialized Snowflake base session for {self.__class__.__name__} with parameters {self.process_parameters}")
-
-
-    @staticmethod
-    def parse_key_value_args():
-        """
-        Parse command line arguments formatted as key=value into a dict.
-        Example: python script.py foo=123 bar=hello
-
-        Returns: {'foo': '123', 'bar': 'hello'}
-        """
-        argv = sys.argv
-        params: dict[str,str] = {}
-        for arg in argv:
-            if "=" in arg:
-                key, value = arg.split("=", 1)  # split only on first '='
-                params[key] = value
-            else:
-                raise ValueError(f"Invalid argument format (expected key=value): {arg}")
-        return params
-
-    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 Snowflake table
-
-        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")
-            table = f"{self._systemConfiguration.dataLakeDbName}.{self._systemConfiguration.dataLakeSchemaName}.INPUT_{self.process.inputId}"
-            self.log(f"Writing dataframe to table {table}")
-            df.write.save_as_table(
-                name=table,
-                mode="overwrite",
-                table_type="transient"
-            )
-            self.log(f"Table {table} written")
-            if self.process.startProcessFlag:
-                # process started by IngestionSession, tell Core to continue and not run Notebook
-                self._pg.sql("SELECT sparky.sdk_complete_manual_process(%s)", [self.process.processId], fetch=False)
-        except Exception as e:
-            self._log_fail(e)
-            if self.process.startProcessFlag:
-                # Fail input and process to prevent core from executing it
-                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.close()

dataforge_sdk-10.0.dev115/dataforge/pg.py

@@ -1,91 +0,0 @@
-"""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-10.0.dev115/dataforge_sdk.egg-info/SOURCES.txt

@@ -1,20 +0,0 @@
-README.md
-pyproject.toml
-setup.cfg
-dataforge/__init__.py
-dataforge/_base_session.py
-dataforge/_databricks_session.py
-dataforge/_session.py
-dataforge/_snowflake_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