snowpark-checkpoints-validators 0.1.4__py3-none-any.whl → 0.2.0__py3-none-any.whl

snowflake/snowpark_checkpoints/__init__.py

@@ -13,16 +13,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+
+
+# Add a NullHandler to prevent logging messages from being output to
+# sys.stderr if no logging configuration is provided.
+logging.getLogger(__name__).addHandler(logging.NullHandler())
+
+# ruff: noqa: E402
+
 from snowflake.snowpark_checkpoints.checkpoint import (
     check_dataframe_schema,
-    check_output_schema,
     check_input_schema,
+    check_output_schema,
     validate_dataframe_checkpoint,
 )
 from snowflake.snowpark_checkpoints.job_context import SnowparkJobContext
 from snowflake.snowpark_checkpoints.spark_migration import check_with_spark
 from snowflake.snowpark_checkpoints.utils.constants import CheckpointMode
 
+
 __all__ = [
     "check_with_spark",
     "SnowparkJobContext",

snowflake/snowpark_checkpoints/__version__.py

@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.1.4"
+__version__ = "0.2.0"

snowflake/snowpark_checkpoints/checkpoint.py

@@ -14,6 +14,9 @@
 # limitations under the License.
 
 # Wrapper around pandera which logs to snowflake
+
+import logging
+
 from typing import Any, Optional, Union, cast
 
 from pandas import DataFrame as PandasDataFrame
@@ -27,13 +30,13 @@ from snowflake.snowpark_checkpoints.snowpark_sampler import (
     SamplingAdapter,
     SamplingStrategy,
 )
-from snowflake.snowpark_checkpoints.utils.checkpoint_logger import CheckpointLogger
 from snowflake.snowpark_checkpoints.utils.constants import (
     FAIL_STATUS,
     PASS_STATUS,
     CheckpointMode,
 )
 from snowflake.snowpark_checkpoints.utils.extra_config import is_checkpoint_enabled
+from snowflake.snowpark_checkpoints.utils.logging_utils import log
 from snowflake.snowpark_checkpoints.utils.pandera_check_manager import (
     PanderaCheckManager,
 )
@@ -47,6 +50,10 @@ from snowflake.snowpark_checkpoints.utils.utils_checks import (
 )
 
 
+LOGGER = logging.getLogger(__name__)
+
+
+@log
 def validate_dataframe_checkpoint(
     df: SnowparkDataFrame,
     checkpoint_name: str,
@@ -84,31 +91,45 @@ def validate_dataframe_checkpoint(
     """
     checkpoint_name = _replace_special_characters(checkpoint_name)
 
-    if is_checkpoint_enabled(checkpoint_name):
-
-        if mode == CheckpointMode.SCHEMA:
-            return _check_dataframe_schema_file(
-                df,
-                checkpoint_name,
-                job_context,
-                custom_checks,
-                skip_checks,
-                sample_frac,
-                sample_number,
-                sampling_strategy,
-                output_path,
-            )
-        elif mode == CheckpointMode.DATAFRAME:
-            if job_context is None:
-                raise ValueError(
-                    "Connectionless mode is not supported for Parquet validation"
-                )
-            _check_compare_data(df, job_context, checkpoint_name, output_path)
-        else:
+    if not is_checkpoint_enabled(checkpoint_name):
+        LOGGER.warning(
+            "Checkpoint '%s' is disabled. Skipping DataFrame checkpoint validation.",
+            checkpoint_name,
+        )
+        return None
+
+    LOGGER.info(
+        "Starting DataFrame checkpoint validation for checkpoint '%s'", checkpoint_name
+    )
+
+    if mode == CheckpointMode.SCHEMA:
+        result = _check_dataframe_schema_file(
+            df,
+            checkpoint_name,
+            job_context,
+            custom_checks,
+            skip_checks,
+            sample_frac,
+            sample_number,
+            sampling_strategy,
+            output_path,
+        )
+        return result
+
+    if mode == CheckpointMode.DATAFRAME:
+        if job_context is None:
             raise ValueError(
-                """Invalid validation mode.
-                Please use for schema validation use a 1 or for a full data validation use a 2 for schema validation."""
+                "No job context provided. Please provide one when using DataFrame mode validation."
             )
+        _check_compare_data(df, job_context, checkpoint_name, output_path)
+        return None
+
+    raise ValueError(
+        (
+            "Invalid validation mode. "
+            "Please use 1 for schema validation or 2 for full data validation."
+        ),
+    )
 
 
 def _check_dataframe_schema_file(
@@ -156,7 +177,7 @@ def _check_dataframe_schema_file(
 
     schema = _generate_schema(checkpoint_name, output_path)
 
-    return check_dataframe_schema(
+    return _check_dataframe_schema(
         df,
         schema,
         checkpoint_name,
@@ -170,6 +191,7 @@ def _check_dataframe_schema_file(
     )
 
 
+@log
 def check_dataframe_schema(
     df: SnowparkDataFrame,
     pandera_schema: DataFrameSchema,
@@ -212,6 +234,9 @@ def check_dataframe_schema(
 
     """
     checkpoint_name = _replace_special_characters(checkpoint_name)
+    LOGGER.info(
+        "Starting DataFrame schema validation for checkpoint '%s'", checkpoint_name
+    )
 
     if df is None:
         raise ValueError("DataFrame is required")
@@ -219,19 +244,25 @@ def check_dataframe_schema(
     if pandera_schema is None:
         raise ValueError("Schema is required")
 
-    if is_checkpoint_enabled(checkpoint_name):
-        return _check_dataframe_schema(
-            df,
-            pandera_schema,
+    if not is_checkpoint_enabled(checkpoint_name):
+        LOGGER.warning(
+            "Checkpoint '%s' is disabled. Skipping DataFrame schema validation.",
             checkpoint_name,
-            job_context,
-            custom_checks,
-            skip_checks,
-            sample_frac,
-            sample_number,
-            sampling_strategy,
-            output_path,
         )
+        return None
+
+    return _check_dataframe_schema(
+        df,
+        pandera_schema,
+        checkpoint_name,
+        job_context,
+        custom_checks,
+        skip_checks,
+        sample_frac,
+        sample_number,
+        sampling_strategy,
+        output_path,
+    )
 
 
 @report_telemetry(
@@ -261,10 +292,22 @@ def _check_dataframe_schema(
     )
     is_valid, validation_result = _validate(pandera_schema_upper, sample_df)
     if is_valid:
+        LOGGER.info(
+            "DataFrame schema validation passed for checkpoint '%s'",
+            checkpoint_name,
+        )
         if job_context is not None:
             job_context._mark_pass(checkpoint_name)
+        else:
+            LOGGER.warning(
+                "No job context provided. Skipping result recording into Snowflake.",
+            )
         _update_validation_result(checkpoint_name, PASS_STATUS, output_path)
     else:
+        LOGGER.error(
+            "DataFrame schema validation failed for checkpoint '%s'",
+            checkpoint_name,
+        )
         _update_validation_result(checkpoint_name, FAIL_STATUS, output_path)
         raise SchemaValidationError(
             "Snowpark DataFrame schema validation error",
@@ -277,6 +320,7 @@ def _check_dataframe_schema(
 
 
 @report_telemetry(params_list=["pandera_schema"])
+@log
 def check_output_schema(
     pandera_schema: DataFrameSchema,
     checkpoint_name: str,
@@ -313,11 +357,8 @@ def check_output_schema(
             function: The decorated function.
 
         """
-        _checkpoint_name = checkpoint_name
-        if checkpoint_name is None:
-            _checkpoint_name = snowpark_fn.__name__
-        _checkpoint_name = _replace_special_characters(_checkpoint_name)
 
+        @log(log_args=False)
         def wrapper(*args, **kwargs):
             """Wrapp a function to validate the schema of the output of a Snowpark function.
 
@@ -329,7 +370,25 @@ def check_output_schema(
                 Any: The result of the Snowpark function.
 
             """
+            _checkpoint_name = checkpoint_name
+            if checkpoint_name is None:
+                LOGGER.warning(
+                    (
+                        "No checkpoint name provided for output schema validation. "
+                        "Using '%s' as the checkpoint name.",
+                    ),
+                    snowpark_fn.__name__,
+                )
+                _checkpoint_name = snowpark_fn.__name__
+            _checkpoint_name = _replace_special_characters(_checkpoint_name)
+            LOGGER.info(
+                "Starting output schema validation for Snowpark function '%s' and checkpoint '%s'",
+                snowpark_fn.__name__,
+                _checkpoint_name,
+            )
+
             # Run the sampled data in snowpark
+            LOGGER.info("Running the Snowpark function '%s'", snowpark_fn.__name__)
             snowpark_results = snowpark_fn(*args, **kwargs)
             sampler = SamplingAdapter(
                 job_context, sample_frac, sample_number, sampling_strategy
@@ -340,17 +399,25 @@ def check_output_schema(
             is_valid, validation_result = _validate(
                 pandera_schema, pandas_sample_args[0]
             )
-            logger = CheckpointLogger().get_logger()
-            logger.info(
-                f"Checkpoint {_checkpoint_name} validation result:\n{validation_result}"
-            )
-
             if is_valid:
+                LOGGER.info(
+                    "Output schema validation passed for Snowpark function '%s' and checkpoint '%s'",
+                    snowpark_fn.__name__,
+                    _checkpoint_name,
+                )
                 if job_context is not None:
                     job_context._mark_pass(_checkpoint_name)
-
+                else:
+                    LOGGER.warning(
+                        "No job context provided. Skipping result recording into Snowflake.",
+                    )
                 _update_validation_result(_checkpoint_name, PASS_STATUS, output_path)
             else:
+                LOGGER.error(
+                    "Output schema validation failed for Snowpark function '%s' and checkpoint '%s'",
+                    snowpark_fn.__name__,
+                    _checkpoint_name,
+                )
                 _update_validation_result(_checkpoint_name, FAIL_STATUS, output_path)
                 raise SchemaValidationError(
                     "Snowpark output schema validation error",
@@ -358,7 +425,6 @@ def check_output_schema(
                     _checkpoint_name,
                     validation_result,
                 )
-
             return snowpark_results
 
         return wrapper
@@ -367,6 +433,7 @@ def check_output_schema(
 
 
 @report_telemetry(params_list=["pandera_schema"])
+@log
 def check_input_schema(
     pandera_schema: DataFrameSchema,
     checkpoint_name: str,
@@ -407,11 +474,8 @@ def check_input_schema(
             Callable: A wrapper function that performs schema validation before executing the original function.
 
         """
-        _checkpoint_name = checkpoint_name
-        if checkpoint_name is None:
-            _checkpoint_name = snowpark_fn.__name__
-        _checkpoint_name = _replace_special_characters(_checkpoint_name)
 
+        @log(log_args=False)
         def wrapper(*args, **kwargs):
             """Wrapp a function to validate the schema of the input of a Snowpark function.
 
@@ -422,6 +486,23 @@ def check_input_schema(
                 Any: The result of the original function after input validation.
 
             """
+            _checkpoint_name = checkpoint_name
+            if checkpoint_name is None:
+                LOGGER.warning(
+                    (
+                        "No checkpoint name provided for input schema validation. "
+                        "Using '%s' as the checkpoint name."
+                    ),
+                    snowpark_fn.__name__,
+                )
+                _checkpoint_name = snowpark_fn.__name__
+            _checkpoint_name = _replace_special_characters(_checkpoint_name)
+            LOGGER.info(
+                "Starting input schema validation for Snowpark function '%s' and checkpoint '%s'",
+                snowpark_fn.__name__,
+                _checkpoint_name,
+            )
+
             # Run the sampled data in snowpark
             sampler = SamplingAdapter(
                 job_context, sample_frac, sample_number, sampling_strategy
@@ -429,43 +510,53 @@ def check_input_schema(
             sampler.process_args(args)
             pandas_sample_args = sampler.get_sampled_pandas_args()
 
+            LOGGER.info(
+                "Validating %s input argument(s) against a Pandera schema",
+                len(pandas_sample_args),
+            )
             # Raises SchemaError on validation issues
-            for arg in pandas_sample_args:
-                if isinstance(arg, PandasDataFrame):
-
-                    is_valid, validation_result = _validate(
-                        pandera_schema,
-                        arg,
+            for index, arg in enumerate(pandas_sample_args, start=1):
+                if not isinstance(arg, PandasDataFrame):
+                    LOGGER.info(
+                        "Arg %s: Skipping schema validation for non-DataFrame argument",
+                        index,
                     )
+                    continue
 
-                    logger = CheckpointLogger().get_logger()
-                    logger.info(
-                        f"Checkpoint {checkpoint_name} validation result:\n{validation_result}"
+                is_valid, validation_result = _validate(
+                    pandera_schema,
+                    arg,
+                )
+                if is_valid:
+                    LOGGER.info(
+                        "Arg %s: Input schema validation passed",
+                        index,
                     )
-
-                    if is_valid:
-                        if job_context is not None:
-                            job_context._mark_pass(
-                                _checkpoint_name,
-                            )
-
-                        _update_validation_result(
-                            _checkpoint_name,
-                            PASS_STATUS,
-                            output_path,
-                        )
-                    else:
-                        _update_validation_result(
-                            _checkpoint_name,
-                            FAIL_STATUS,
-                            output_path,
-                        )
-                        raise SchemaValidationError(
-                            "Snowpark input schema validation error",
-                            job_context,
+                    if job_context is not None:
+                        job_context._mark_pass(
                             _checkpoint_name,
-                            validation_result,
                         )
+                    _update_validation_result(
+                        _checkpoint_name,
+                        PASS_STATUS,
+                        output_path,
+                    )
+                else:
+                    LOGGER.error(
+                        "Arg %s: Input schema validation failed",
+                        index,
+                    )
+                    _update_validation_result(
+                        _checkpoint_name,
+                        FAIL_STATUS,
+                        output_path,
+                    )
+                    raise SchemaValidationError(
+                        "Snowpark input schema validation error",
+                        job_context,
+                        _checkpoint_name,
+                        validation_result,
+                    )
             return snowpark_fn(*args, **kwargs)
 
         return wrapper

snowflake/snowpark_checkpoints/job_context.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+
 from datetime import datetime
 from typing import Optional
 
@@ -24,6 +26,10 @@ from snowflake.snowpark import Session
 from snowflake.snowpark_checkpoints.utils.constants import SCHEMA_EXECUTION_MODE
 
 
+LOGGER = logging.getLogger(__name__)
+RESULTS_TABLE = "SNOWPARK_CHECKPOINTS_REPORT"
+
+
 class SnowparkJobContext:
 
     """Class used to record migration results in Snowflake.
@@ -45,41 +51,78 @@ class SnowparkJobContext:
     ):
         self.log_results = log_results
         self.job_name = job_name
-        self.spark_session = spark_session or SparkSession.builder.getOrCreate()
+        self.spark_session = spark_session or self._create_pyspark_session()
         self.snowpark_session = snowpark_session
 
     def _mark_fail(
         self, message, checkpoint_name, data, execution_mode=SCHEMA_EXECUTION_MODE
     ):
-        if self.log_results:
-            session = self.snowpark_session
-            df = pd.DataFrame(
-                {
-                    "DATE": [datetime.now()],
-                    "JOB": [self.job_name],
-                    "STATUS": ["fail"],
-                    "CHECKPOINT": [checkpoint_name],
-                    "MESSAGE": [message],
-                    "DATA": [f"{data}"],
-                    "EXECUTION_MODE": [execution_mode],
-                }
+        if not self.log_results:
+            LOGGER.warning(
+                (
+                    "Recording of migration results into Snowflake is disabled. "
+                    "Failure result for checkpoint '%s' will not be recorded."
+                ),
+                checkpoint_name,
             )
-            report_df = session.createDataFrame(df)
-            report_df.write.mode("append").save_as_table("SNOWPARK_CHECKPOINTS_REPORT")
+            return
+
+        LOGGER.debug(
+            "Marking failure for checkpoint '%s' in '%s' mode with message '%s'",
+            checkpoint_name,
+            execution_mode,
+            message,
+        )
+
+        session = self.snowpark_session
+        df = pd.DataFrame(
+            {
+                "DATE": [datetime.now()],
+                "JOB": [self.job_name],
+                "STATUS": ["fail"],
+                "CHECKPOINT": [checkpoint_name],
+                "MESSAGE": [message],
+                "DATA": [f"{data}"],
+                "EXECUTION_MODE": [execution_mode],
+            }
+        )
+        report_df = session.createDataFrame(df)
+        LOGGER.info("Writing failure result to table: '%s'", RESULTS_TABLE)
+        report_df.write.mode("append").save_as_table(RESULTS_TABLE)
 
     def _mark_pass(self, checkpoint_name, execution_mode=SCHEMA_EXECUTION_MODE):
-        if self.log_results:
-            session = self.snowpark_session
-            df = pd.DataFrame(
-                {
-                    "DATE": [datetime.now()],
-                    "JOB": [self.job_name],
-                    "STATUS": ["pass"],
-                    "CHECKPOINT": [checkpoint_name],
-                    "MESSAGE": [""],
-                    "DATA": [""],
-                    "EXECUTION_MODE": [execution_mode],
-                }
+        if not self.log_results:
+            LOGGER.warning(
+                (
+                    "Recording of migration results into Snowflake is disabled. "
+                    "Pass result for checkpoint '%s' will not be recorded."
+                ),
+                checkpoint_name,
             )
-            report_df = session.createDataFrame(df)
-            report_df.write.mode("append").save_as_table("SNOWPARK_CHECKPOINTS_REPORT")
+            return
+
+        LOGGER.debug(
+            "Marking pass for checkpoint '%s' in '%s' mode",
+            checkpoint_name,
+            execution_mode,
+        )
+
+        session = self.snowpark_session
+        df = pd.DataFrame(
+            {
+                "DATE": [datetime.now()],
+                "JOB": [self.job_name],
+                "STATUS": ["pass"],
+                "CHECKPOINT": [checkpoint_name],
+                "MESSAGE": [""],
+                "DATA": [""],
+                "EXECUTION_MODE": [execution_mode],
+            }
+        )
+        report_df = session.createDataFrame(df)
+        LOGGER.info("Writing pass result to table: '%s'", RESULTS_TABLE)
+        report_df.write.mode("append").save_as_table(RESULTS_TABLE)
+
+    def _create_pyspark_session(self) -> SparkSession:
+        LOGGER.info("Creating a PySpark session")
+        return SparkSession.builder.getOrCreate()

snowflake/snowpark_checkpoints/snowpark_sampler.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+
 from typing import Optional
 
 import pandas
@@ -21,6 +23,9 @@ from snowflake.snowpark import DataFrame as SnowparkDataFrame
 from snowflake.snowpark_checkpoints.job_context import SnowparkJobContext
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 class SamplingStrategy:
     RANDOM_SAMPLE = 1
     LIMIT = 2
@@ -52,23 +57,43 @@ class SamplingAdapter:
     def process_args(self, input_args):
         # create the intermediate pandas
         # data frame for the test data
+        LOGGER.info("Processing %s input argument(s) for sampling", len(input_args))
         for arg in input_args:
             if isinstance(arg, SnowparkDataFrame):
-                if arg.count() == 0:
+                df_count = arg.count()
+                if df_count == 0:
                     raise SamplingError(
                         "Input DataFrame is empty. Cannot sample from an empty DataFrame."
                     )
 
+                LOGGER.info("Sampling a Snowpark DataFrame with %s rows", df_count)
                 if self.sampling_strategy == SamplingStrategy.RANDOM_SAMPLE:
                     if self.sample_frac:
+                        LOGGER.info(
+                            "Applying random sampling with fraction %s",
+                            self.sample_frac,
+                        )
                         df_sample = arg.sample(frac=self.sample_frac).to_pandas()
                     else:
+                        LOGGER.info(
+                            "Applying random sampling with size %s", self.sample_number
+                        )
                         df_sample = arg.sample(n=self.sample_number).to_pandas()
                 else:
+                    LOGGER.info(
+                        "Applying limit sampling with size %s", self.sample_number
+                    )
                     df_sample = arg.limit(self.sample_number).to_pandas()
 
+                LOGGER.info(
+                    "Successfully sampled the DataFrame. Resulting DataFrame shape: %s",
+                    df_sample.shape,
+                )
                 self.pandas_sample_args.append(df_sample)
             else:
+                LOGGER.debug(
+                    "Argument is not a Snowpark DataFrame. No sampling is applied."
+                )
                 self.pandas_sample_args.append(arg)
 
     def get_sampled_pandas_args(self):

snowflake/snowpark_checkpoints/spark_migration.py

@@ -12,6 +12,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+import logging
+
 from typing import Callable, Optional, TypeVar
 
 import pandas as pd
@@ -27,6 +30,7 @@ from snowflake.snowpark_checkpoints.snowpark_sampler import (
     SamplingStrategy,
 )
 from snowflake.snowpark_checkpoints.utils.constants import FAIL_STATUS, PASS_STATUS
+from snowflake.snowpark_checkpoints.utils.logging_utils import log
 from snowflake.snowpark_checkpoints.utils.telemetry import STATUS_KEY, report_telemetry
 from snowflake.snowpark_checkpoints.utils.utils_checks import (
     _replace_special_characters,
@@ -35,8 +39,10 @@ from snowflake.snowpark_checkpoints.utils.utils_checks import (
 
 
 fn = TypeVar("F", bound=Callable)
+LOGGER = logging.getLogger(__name__)
 
 
+@log
 def check_with_spark(
     job_context: Optional[SnowparkJobContext],
     spark_function: fn,
@@ -67,12 +73,22 @@ def check_with_spark(
     """
 
     def check_with_spark_decorator(snowpark_fn):
-        _checkpoint_name = checkpoint_name
-        if checkpoint_name is None:
-            _checkpoint_name = snowpark_fn.__name__
-        _checkpoint_name = _replace_special_characters(_checkpoint_name)
-
+        @log(log_args=False)
         def wrapper(*args, **kwargs):
+            LOGGER.info(
+                "Starting output validation between Snowpark function '%s' and Spark function '%s'",
+                snowpark_fn.__name__,
+                spark_function.__name__,
+            )
+            _checkpoint_name = checkpoint_name
+            if checkpoint_name is None:
+                LOGGER.warning(
+                    "No checkpoint name provided. Using '%s' as the checkpoint name",
+                    snowpark_fn.__name__,
+                )
+                _checkpoint_name = snowpark_fn.__name__
+            _checkpoint_name = _replace_special_characters(_checkpoint_name)
+
             sampler = SamplingAdapter(
                 job_context,
                 sample_number=sample_number,
@@ -81,9 +97,14 @@ def check_with_spark(
             sampler.process_args(args)
             snowpark_sample_args = sampler.get_sampled_snowpark_args()
             pyspark_sample_args = sampler.get_sampled_spark_args()
+
             # Run the sampled data in snowpark
+            LOGGER.info("Running the Snowpark function with sampled args")
             snowpark_test_results = snowpark_fn(*snowpark_sample_args, **kwargs)
+            LOGGER.info("Running the Spark function with sampled args")
             spark_test_results = spark_function(*pyspark_sample_args, **kwargs)
+
+            LOGGER.info("Comparing the results of the Snowpark and Spark functions")
             result, exception = _assert_return(
                 snowpark_test_results,
                 spark_test_results,
@@ -92,7 +113,18 @@ def check_with_spark(
                 output_path,
             )
             if not result:
+                LOGGER.error(
+                    "Validation failed. The results of the Snowpark function '%s' and Spark function '%s' do not match",
+                    snowpark_fn.__name__,
+                    spark_function.__name__,
+                )
                 raise exception from None
+            LOGGER.info(
+                "Validation passed. The results of the Snowpark function '%s' and Spark function '%s' match",
+                snowpark_fn.__name__,
+                spark_function.__name__,
+            )
+
             # Run the original function in snowpark
             return snowpark_fn(*args, **kwargs)
 
@@ -126,6 +158,7 @@ def _assert_return(
     if isinstance(snowpark_results, SnowparkDataFrame) and isinstance(
         spark_results, SparkDataFrame
     ):
+        LOGGER.debug("Comparing two DataFrame results for equality")
         cmp = compare_spark_snowpark_dfs(spark_results, snowpark_results)
 
         if not cmp.empty:
@@ -137,7 +170,7 @@ def _assert_return(
         _update_validation_result(checkpoint_name, PASS_STATUS, output_path)
         return True, None
     else:
-
+        LOGGER.debug("Comparing two scalar results for equality")
         if snowpark_results != spark_results:
             exception_result = SparkMigrationError(
                 "Return value difference:\n",

snowflake/snowpark_checkpoints/utils/extra_config.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 
 from typing import Optional
@@ -22,6 +23,9 @@ from snowflake.snowpark_checkpoints.utils.constants import (
 )
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 # noinspection DuplicatedCode
 def _get_checkpoint_contract_file_path() -> str:
     return os.environ.get(SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR, os.getcwd())
@@ -35,10 +39,14 @@ def _get_metadata():
         )
 
         path = _get_checkpoint_contract_file_path()
+        LOGGER.debug("Loading checkpoint metadata from '%s'", path)
         metadata = CheckpointMetadata(path)
         return True, metadata
 
     except ImportError:
+        LOGGER.debug(
+            "snowpark-checkpoints-configuration is not installed. Cannot get a checkpoint metadata instance."
+        )
         return False, None
 
 
@@ -56,8 +64,7 @@ def is_checkpoint_enabled(checkpoint_name: Optional[str] = None) -> bool:
     if enabled and checkpoint_name is not None:
         config = metadata.get_checkpoint(checkpoint_name)
         return config.enabled
-    else:
-        return True
+    return True
 
 
 def get_checkpoint_file(checkpoint_name: str) -> Optional[str]:
@@ -78,7 +85,5 @@ def get_checkpoint_file(checkpoint_name: str) -> Optional[str]:
     enabled, metadata = _get_metadata()
     if enabled:
         config = metadata.get_checkpoint(checkpoint_name)
-
         return config.file
-    else:
-        return None
+    return None

snowflake/snowpark_checkpoints/utils/logging_utils.py

@@ -0,0 +1,67 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from functools import wraps
+from typing import Callable, Optional, TypeVar
+
+from typing_extensions import ParamSpec
+
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def log(
+    _func: Optional[Callable[P, R]] = None,
+    *,
+    logger: Optional[logging.Logger] = None,
+    log_args: bool = True,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Log the function call and any exceptions that occur.
+
+    Args:
+        _func: The function to log.
+        logger: The logger to use for logging. If not provided, a logger will be created using the
+                function's module name.
+        log_args: Whether to log the arguments passed to the function.
+
+    Returns:
+        A decorator that logs the function call and any exceptions that occur.
+
+    """
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            _logger = logging.getLogger(func.__module__) if logger is None else logger
+            if log_args:
+                args_repr = [repr(a) for a in args]
+                kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
+                formatted_args = ", ".join([*args_repr, *kwargs_repr])
+                _logger.debug("%s called with args %s", func.__name__, formatted_args)
+            try:
+                return func(*args, **kwargs)
+            except Exception:
+                _logger.exception("An error occurred in %s", func.__name__)
+                raise
+
+        return wrapper
+
+    # Handle the case where the decorator is used without parentheses
+    if _func is None:
+        return decorator
+    return decorator(_func)

snowflake/snowpark_checkpoints/utils/pandera_check_manager.py

@@ -1,9 +1,25 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
 from datetime import datetime
 from typing import Optional
 
 from pandera import Check, DataFrameSchema
 
-from snowflake.snowpark_checkpoints.utils.checkpoint_logger import CheckpointLogger
 from snowflake.snowpark_checkpoints.utils.constants import (
     COLUMNS_KEY,
     DECIMAL_PRECISION_KEY,
@@ -28,6 +44,9 @@ from snowflake.snowpark_checkpoints.utils.supported_types import (
 )
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 class PanderaCheckManager:
     def __init__(self, checkpoint_name: str, schema: DataFrameSchema):
         self.checkpoint_name = checkpoint_name
@@ -258,25 +277,28 @@ class PanderaCheckManager:
             ValueError: If the column name or type is not defined in the schema.
 
         """
-        logger = CheckpointLogger().get_logger()
-        for additional_check in custom_data.get(COLUMNS_KEY):
+        LOGGER.info("Adding checks for the checkpoint '%s'", self.checkpoint_name)
 
-            type = additional_check.get(TYPE_KEY, None)
+        for additional_check in custom_data.get(COLUMNS_KEY):
             name = additional_check.get(NAME_KEY, None)
-            is_nullable = additional_check.get(NULLABLE_KEY, False)
-
             if name is None:
                 raise ValueError(
                     f"Column name not defined in the schema {self.checkpoint_name}"
                 )
 
+            type = additional_check.get(TYPE_KEY, None)
             if type is None:
                 raise ValueError(f"Type not defined for column {name}")
 
             if self.schema.columns.get(name) is None:
-                logger.warning(f"Column {name} not found in schema")
+                LOGGER.warning(
+                    "Column '%s' was not found in the Pandera schema. Skipping checks for this column.",
+                    name,
+                )
                 continue
 
+            LOGGER.debug("Adding checks for column '%s' of type '%s'", name, type)
+
             if type in NumericTypes:
                 self._add_numeric_checks(name, additional_check)
 
@@ -289,7 +311,9 @@ class PanderaCheckManager:
             elif type == "datetime":
                 self._add_date_time_checks(name, additional_check)
 
+            is_nullable = additional_check.get(NULLABLE_KEY, False)
             if is_nullable:
+                LOGGER.debug("Column '%s' is nullable. Adding null checks.", name)
                 self._add_null_checks(name, additional_check)
 
         return self.schema
@@ -318,8 +342,19 @@ class PanderaCheckManager:
             if col in self.schema.columns:
 
                 if SKIP_ALL in checks_to_skip:
+                    LOGGER.info(
+                        "Skipping all checks for column '%s' in checkpoint '%s'",
+                        col,
+                        self.checkpoint_name,
+                    )
                     self.schema.columns[col].checks = {}
                 else:
+                    LOGGER.info(
+                        "Skipping checks %s for column '%s' in checkpoint '%s'",
+                        checks_to_skip,
+                        col,
+                        self.checkpoint_name,
+                    )
                     self.schema.columns[col].checks = [
                         check
                         for check in self.schema.columns[col].checks
@@ -350,6 +385,12 @@ class PanderaCheckManager:
         for col, checks in custom_checks.items():
 
             if col in self.schema.columns:
+                LOGGER.info(
+                    "Adding %s custom checks to column '%s' in checkpoint '%s'",
+                    len(checks),
+                    col,
+                    self.checkpoint_name,
+                )
                 col_schema = self.schema.columns[col]
                 col_schema.checks.extend(checks)
             else:

snowflake/snowpark_checkpoints/utils/utils_checks.py

@@ -15,6 +15,7 @@
 
 import inspect
 import json
+import logging
 import os
 import re
 
@@ -58,6 +59,9 @@ from snowflake.snowpark_checkpoints.validation_result_metadata import (
 from snowflake.snowpark_checkpoints.validation_results import ValidationResult
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 def _replace_special_characters(checkpoint_name: str) -> str:
     """Replace special characters in the checkpoint name with underscores.
 
@@ -147,6 +151,9 @@ def _generate_schema(
                          constraints of the DataFrame.
 
     """
+    LOGGER.info(
+        "Generating Pandera DataFrameSchema for checkpoint: '%s'", checkpoint_name
+    )
     current_directory_path = output_path if output_path else os.getcwd()
 
     output_directory_path = os.path.join(
@@ -169,6 +176,7 @@ Please run the Snowpark checkpoint collector first."""
             f"Checkpoint {checkpoint_name} JSON file not found. Please run the Snowpark checkpoint collector first."
         )
 
+    LOGGER.info("Reading schema from file: '%s'", checkpoint_schema_file_path)
     with open(checkpoint_schema_file_path) as schema_file:
         checkpoint_schema_config = json.load(schema_file)
 
@@ -182,6 +190,10 @@ Please run the Snowpark checkpoint collector first."""
     schema = DataFrameSchema.from_json(schema_dict_str)
 
     if DATAFRAME_CUSTOM_DATA_KEY not in checkpoint_schema_config:
+        LOGGER.info(
+            "No custom data found in the JSON file for checkpoint: '%s'",
+            checkpoint_name,
+        )
         return schema
 
     custom_data = checkpoint_schema_config.get(DATAFRAME_CUSTOM_DATA_KEY)
@@ -221,7 +233,7 @@ def _check_compare_data(
         SchemaValidationError: If there is a data mismatch between the DataFrame and the checkpoint table.
 
     """
-    result, err = _compare_data(df, job_context, checkpoint_name, output_path)
+    _, err = _compare_data(df, job_context, checkpoint_name, output_path)
     if err is not None:
         raise err
 
@@ -256,9 +268,18 @@ def _compare_data(
 
     """
     new_table_name = CHECKPOINT_TABLE_NAME_FORMAT.format(checkpoint_name)
-
+    LOGGER.info(
+        "Writing Snowpark DataFrame to table: '%s' for checkpoint: '%s'",
+        new_table_name,
+        checkpoint_name,
+    )
     df.write.save_as_table(table_name=new_table_name, mode="overwrite")
 
+    LOGGER.info(
+        "Comparing DataFrame to checkpoint table: '%s' for checkpoint: '%s'",
+        new_table_name,
+        checkpoint_name,
+    )
     expect_df = job_context.snowpark_session.sql(
         EXCEPT_HASH_AGG_QUERY, [checkpoint_name, new_table_name]
     )

snowflake/snowpark_checkpoints/validation_result_metadata.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 
 from typing import Optional
@@ -28,6 +29,9 @@ from snowflake.snowpark_checkpoints.validation_results import (
 )
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 class ValidationResultsMetadata(metaclass=Singleton):
 
     """ValidationResultsMetadata is a class that manages the loading, storing, and updating of validation results.
@@ -69,14 +73,26 @@ class ValidationResultsMetadata(metaclass=Singleton):
             SNOWPARK_CHECKPOINTS_OUTPUT_DIRECTORY_NAME,
         )
 
+        LOGGER.debug(
+            "Setting validation results directory to: '%s'",
+            self.validation_results_directory,
+        )
+
         self.validation_results_file = os.path.join(
             self.validation_results_directory,
             VALIDATION_RESULTS_JSON_FILE_NAME,
         )
 
+        LOGGER.debug(
+            "Setting validation results file to: '%s'", self.validation_results_file
+        )
+
         self.validation_results = ValidationResults(results=[])
 
         if os.path.exists(self.validation_results_file):
+            LOGGER.info(
+                "Loading validation results from: '%s'", self.validation_results_file
+            )
             with open(self.validation_results_file) as file:
                 try:
                     validation_result_json = file.read()
@@ -87,6 +103,11 @@ class ValidationResultsMetadata(metaclass=Singleton):
                     raise Exception(
                         f"Error reading validation results file: {self.validation_results_file} \n {e}"
                     ) from None
+        else:
+            LOGGER.info(
+                "Validation results file not found: '%s'",
+                self.validation_results_file,
+            )
 
     def clean(self):
         """Clean the validation results list.
@@ -95,6 +116,7 @@ class ValidationResultsMetadata(metaclass=Singleton):
 
         """
         if not os.path.exists(self.validation_results_file):
+            LOGGER.info("Cleaning validation results...")
             self.validation_results.results = []
 
     def add_validation_result(self, validation_result: ValidationResult):
@@ -119,7 +141,15 @@ class ValidationResultsMetadata(metaclass=Singleton):
 
         """
         if not os.path.exists(self.validation_results_directory):
+            LOGGER.debug(
+                "Validation results directory '%s' does not exist. Creating it...",
+                self.validation_results_directory,
+            )
             os.makedirs(self.validation_results_directory)
 
         with open(self.validation_results_file, "w") as output_file:
             output_file.write(self.validation_results.model_dump_json())
+            LOGGER.info(
+                "Validation results successfully saved to: '%s'",
+                self.validation_results_file,
+            )

{snowpark_checkpoints_validators-0.1.4.dist-info → snowpark_checkpoints_validators-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-checkpoints-validators
-Version: 0.1.4
+Version: 0.2.0
 Summary: Migration tools for Snowpark
 Project-URL: Bug Tracker, https://github.com/snowflakedb/snowpark-checkpoints/issues
 Project-URL: Source code, https://github.com/snowflakedb/snowpark-checkpoints/

snowpark_checkpoints_validators-0.2.0.dist-info/RECORD

@@ -0,0 +1,22 @@
+snowflake/snowpark_checkpoints/__init__.py,sha256=p7fzH3f8foD5nhNJHZ00JT3ODTXJGGkWTd3xRKx-8aQ,1435
+snowflake/snowpark_checkpoints/__version__.py,sha256=ajnGza8ucK69-PA8wEbHmWZxDwd3bsTm74yMKiIWNHY,632
+snowflake/snowpark_checkpoints/checkpoint.py,sha256=i-iDRYbGvQHy9ipW7UxHVhJhQ9BXNSO-bsCcHyg3oLA,22056
+snowflake/snowpark_checkpoints/errors.py,sha256=9KjzRf8bjDZTTNL4LeySJAwuucDOyz0Ka7EFBKWFpyg,1821
+snowflake/snowpark_checkpoints/job_context.py,sha256=RMK0g0HrbDVrOAvai4PgsGvsAn_GIo9aFmh-tWlyieY,4183
+snowflake/snowpark_checkpoints/singleton.py,sha256=7AgIHQBXVRvPBBCkmBplzkdrrm-xVWf_N8svzA2vF8E,836
+snowflake/snowpark_checkpoints/snowpark_sampler.py,sha256=Qxv-8nRGuf-ab3GoSUt8_MNL0ppjoBIMOFIMkqmwN5I,4668
+snowflake/snowpark_checkpoints/spark_migration.py,sha256=s2HqomYx76Hqn71g9TleBeHI3t1nirgfPvkggqQQdts,10253
+snowflake/snowpark_checkpoints/validation_result_metadata.py,sha256=fm2lKxjYlzlL6qsiv2icR9k5o7YNd2OwvFhiqGYrTpo,5745
+snowflake/snowpark_checkpoints/validation_results.py,sha256=J8OcpNty6hQD8RbAy8xmA0UMbPWfXSmQnHYspWWSisk,1502
+snowflake/snowpark_checkpoints/utils/__init__.py,sha256=I4srmZ8G1q9DU6Suo1S91aVfNvETyisKH95uvLAvEJ0,609
+snowflake/snowpark_checkpoints/utils/constants.py,sha256=pgFttLDQ6fTa6obSdvivWBYClS21ap41YVDNGAS4sxY,4146
+snowflake/snowpark_checkpoints/utils/extra_config.py,sha256=LvOdIhvE450AV0wLVK5P_hANvcNzAv8pLNe7Ksr598U,2802
+snowflake/snowpark_checkpoints/utils/logging_utils.py,sha256=yyi6X5DqKeTg0HRhvsH6ymYp2P0wbnyKIzI2RzrQS7k,2278
+snowflake/snowpark_checkpoints/utils/pandera_check_manager.py,sha256=tQIozLO-2kM8WZ-gGKfRwmXBx1cDPaIZB0qIcArp8xA,16100
+snowflake/snowpark_checkpoints/utils/supported_types.py,sha256=GrMX2tHdSFnK7LlPbZx20UufD6Br6TNVRkkBwIxdPy0,1433
+snowflake/snowpark_checkpoints/utils/telemetry.py,sha256=_WOVo19BxcF6cpQDplID6BEOvgJfHTGK1JZI1-OI4uc,31370
+snowflake/snowpark_checkpoints/utils/utils_checks.py,sha256=ythgWkLstEkCae_TqtdPXJ1Jjbx9iTN8sLOl1ewKxzI,14191
+snowpark_checkpoints_validators-0.2.0.dist-info/METADATA,sha256=ixLNouygrcyBFCQK3D77nmAIKsWnPIV9gCYSP_rRi1I,11470
+snowpark_checkpoints_validators-0.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+snowpark_checkpoints_validators-0.2.0.dist-info/licenses/LICENSE,sha256=pmjhbh6uVhV5MBXOlou_UZgFP7CYVQITkCCdvfcS5lY,11340
+snowpark_checkpoints_validators-0.2.0.dist-info/RECORD,,

snowflake/snowpark_checkpoints/utils/checkpoint_logger.py

@@ -1,52 +0,0 @@
-# Copyright 2025 Snowflake Inc.
-# SPDX-License-Identifier: Apache-2.0
-
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-
-# http://www.apache.org/licenses/LICENSE-2.0
-
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import logging
-import threading
-
-
-class CheckpointLogger:
-    _instance = None
-    _lock = threading.Lock()
-
-    def __new__(cls, *args, **kwargs):
-        if not cls._instance:
-            with cls._lock:
-                if not cls._instance:
-                    cls._instance = super().__new__(cls, *args, **kwargs)
-                    cls._instance._initialize()
-        return cls._instance
-
-    def _initialize(self):
-        # Create formatter
-        formatter = logging.Formatter(
-            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-        )
-
-        self.logger = logging.getLogger("CheckpointLogger")
-        self.logger.setLevel(logging.INFO)
-
-        # Create console handler and set level to debug
-        ch = logging.StreamHandler()
-        ch.setLevel(logging.DEBUG)
-
-        # Add formatter to ch
-        ch.setFormatter(formatter)
-
-        # Add ch to logger
-        self.logger.addHandler(ch)
-
-    def get_logger(self):
-        return self.logger

snowpark_checkpoints_validators-0.1.4.dist-info/RECORD

@@ -1,22 +0,0 @@
-snowflake/snowpark_checkpoints/__init__.py,sha256=1_xzSopIHWpw1i3gQqWLN0wCfWWEefjr4cl1vl0xSdY,1211
-snowflake/snowpark_checkpoints/__version__.py,sha256=XLR7FIaqd--6dSDY841wZ2c8LM8IAx8L5fVTjJIHRkQ,632
-snowflake/snowpark_checkpoints/checkpoint.py,sha256=VceodgS3muZLH5uWhqmuMsvS7cTBAOWOWFCRKt2VsNg,19150
-snowflake/snowpark_checkpoints/errors.py,sha256=9KjzRf8bjDZTTNL4LeySJAwuucDOyz0Ka7EFBKWFpyg,1821
-snowflake/snowpark_checkpoints/job_context.py,sha256=7LdJ682lC8hCJOYUn-AVXq_Llv18R9oGdK2F-amYR_o,2990
-snowflake/snowpark_checkpoints/singleton.py,sha256=7AgIHQBXVRvPBBCkmBplzkdrrm-xVWf_N8svzA2vF8E,836
-snowflake/snowpark_checkpoints/snowpark_sampler.py,sha256=-t7cg-swMK0SaU7r8y90MLSDPXGlKprc6xdVxEs29sU,3632
-snowflake/snowpark_checkpoints/spark_migration.py,sha256=DzzgUZ-XlzIqCz-aWpBICP8mgnjk8UNoL8JsomadF-U,8832
-snowflake/snowpark_checkpoints/validation_result_metadata.py,sha256=mHCIq6-F37HK-jYBAPeVtax9eIwiCvQZxFPGWi4KvQc,4765
-snowflake/snowpark_checkpoints/validation_results.py,sha256=J8OcpNty6hQD8RbAy8xmA0UMbPWfXSmQnHYspWWSisk,1502
-snowflake/snowpark_checkpoints/utils/__init__.py,sha256=I4srmZ8G1q9DU6Suo1S91aVfNvETyisKH95uvLAvEJ0,609
-snowflake/snowpark_checkpoints/utils/checkpoint_logger.py,sha256=meGl5T3Avp4Qn0GEwkJi5GSLS4MDb7zTGbTOI-8bf1E,1592
-snowflake/snowpark_checkpoints/utils/constants.py,sha256=pgFttLDQ6fTa6obSdvivWBYClS21ap41YVDNGAS4sxY,4146
-snowflake/snowpark_checkpoints/utils/extra_config.py,sha256=pmGLYT7cu9WMKzQwcEPkgk1DMnnT1fREm45p19e79hk,2567
-snowflake/snowpark_checkpoints/utils/pandera_check_manager.py,sha256=ddTwXauuZdowIRwPMT61GWYCG4XGKOFkVyfZO49bc-8,14516
-snowflake/snowpark_checkpoints/utils/supported_types.py,sha256=GrMX2tHdSFnK7LlPbZx20UufD6Br6TNVRkkBwIxdPy0,1433
-snowflake/snowpark_checkpoints/utils/telemetry.py,sha256=_WOVo19BxcF6cpQDplID6BEOvgJfHTGK1JZI1-OI4uc,31370
-snowflake/snowpark_checkpoints/utils/utils_checks.py,sha256=o9HOBrDuTxSIgzZQHfsa9pMzzXRUsRAISI7L6OURouo,13528
-snowpark_checkpoints_validators-0.1.4.dist-info/METADATA,sha256=L00oXPSPvCUBEJuYcAPPt08BPyKvbd3ZgPLpr3DgvqQ,11470
-snowpark_checkpoints_validators-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-snowpark_checkpoints_validators-0.1.4.dist-info/licenses/LICENSE,sha256=pmjhbh6uVhV5MBXOlou_UZgFP7CYVQITkCCdvfcS5lY,11340
-snowpark_checkpoints_validators-0.1.4.dist-info/RECORD,,