snowpark-checkpoints-collectors 0.2.1__py3-none-any.whl → 0.3.1__py3-none-any.whl

snowflake/snowpark_checkpoints_collector/__init__.py

@@ -22,9 +22,14 @@ logging.getLogger(__name__).addHandler(logging.NullHandler())
 
 # ruff: noqa: E402
 
-__all__ = ["collect_dataframe_checkpoint", "CheckpointMode"]
+__all__ = [
+    "collect_dataframe_checkpoint",
+    "CheckpointMode",
+    "xcollect_dataframe_checkpoint",
+]
 
 from snowflake.snowpark_checkpoints_collector.collection_common import CheckpointMode
 from snowflake.snowpark_checkpoints_collector.summary_stats_collector import (
     collect_dataframe_checkpoint,
+    xcollect_dataframe_checkpoint,
 )

snowflake/snowpark_checkpoints_collector/__version__.py

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

snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result.py

@@ -30,6 +30,7 @@ CHECKPOINT_NAME_KEY = "checkpoint_name"
 class CollectionResult(Enum):
     FAIL = "FAIL"
     PASS = "PASS"
+    SKIP = "SKIP"
 
 
 class CollectionPointResult:

snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result_manager.py

@@ -21,6 +21,9 @@ from typing import Optional
 from snowflake.snowpark_checkpoints_collector.collection_result.model import (
     CollectionPointResult,
 )
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    get_io_file_manager,
+)
 from snowflake.snowpark_checkpoints_collector.singleton import Singleton
 from snowflake.snowpark_checkpoints_collector.utils import file_utils
 
@@ -70,5 +73,4 @@ class CollectionPointResultManager(metaclass=Singleton):
     def _save_result(self) -> None:
         result_collection_json = self.to_json()
         LOGGER.info("Saving collection results to '%s'", self.output_file_path)
-        with open(self.output_file_path, "w") as f:
-            f.write(result_collection_json)
+        get_io_file_manager().write(self.output_file_path, result_collection_json)

snowflake/snowpark_checkpoints_collector/io_utils/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+
+__all__ = ["EnvStrategy", "IOFileManager", "IODefaultStrategy"]
+
+from snowflake.snowpark_checkpoints_collector.io_utils.io_env_strategy import (
+    EnvStrategy,
+)
+from snowflake.snowpark_checkpoints_collector.io_utils.io_default_strategy import (
+    IODefaultStrategy,
+)
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    IOFileManager,
+)

snowflake/snowpark_checkpoints_collector/io_utils/io_default_strategy.py

@@ -0,0 +1,61 @@
+# 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 glob
+import os
+import shutil
+
+from pathlib import Path
+from typing import Optional
+
+from snowflake.snowpark_checkpoints_collector.io_utils import EnvStrategy
+
+
+class IODefaultStrategy(EnvStrategy):
+    def mkdir(self, path: str, exist_ok: bool = False) -> None:
+        os.makedirs(path, exist_ok=exist_ok)
+
+    def folder_exists(self, path: str) -> bool:
+        return os.path.isdir(path)
+
+    def file_exists(self, path: str) -> bool:
+        return os.path.isfile(path)
+
+    def write(self, file_path: str, file_content: str, overwrite: bool = True) -> None:
+        mode = "w" if overwrite else "x"
+        with open(file_path, mode) as file:
+            file.write(file_content)
+
+    def read(
+        self, file_path: str, mode: str = "r", encoding: Optional[str] = None
+    ) -> str:
+        with open(file_path, mode=mode, encoding=encoding) as file:
+            return file.read()
+
+    def read_bytes(self, file_path: str) -> bytes:
+        with open(file_path, mode="rb") as f:
+            return f.read()
+
+    def ls(self, path: str, recursive: bool = False) -> list[str]:
+        return glob.glob(path, recursive=recursive)
+
+    def getcwd(self) -> str:
+        return os.getcwd()
+
+    def remove_dir(self, path: str) -> None:
+        shutil.rmtree(path)
+
+    def telemetry_path_files(self, path: str) -> Path:
+        return Path(path)

snowflake/snowpark_checkpoints_collector/io_utils/io_env_strategy.py

@@ -0,0 +1,142 @@
+# 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.
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+
+class EnvStrategy(ABC):
+
+    """An abstract base class that defines methods for file and directory operations.
+
+    Subclasses should implement these methods to provide environment-specific behavior.
+    """
+
+    @abstractmethod
+    def mkdir(self, path: str, exist_ok: bool = False) -> None:
+        """Create a directory.
+
+        Args:
+            path: The name of the directory to create.
+            exist_ok: If False, an error is raised if the directory already exists.
+
+        """
+
+    @abstractmethod
+    def folder_exists(self, path: str) -> bool:
+        """Check if a folder exists.
+
+        Args:
+            path: The path to the folder.
+
+        Returns:
+            bool: True if the folder exists, False otherwise.
+
+        """
+
+    @abstractmethod
+    def file_exists(self, path: str) -> bool:
+        """Check if a file exists.
+
+        Args:
+            path: The path to the file.
+
+        Returns:
+            bool: True if the file exists, False otherwise.
+
+        """
+
+    @abstractmethod
+    def write(self, file_path: str, file_content: str, overwrite: bool = True) -> None:
+        """Write content to a file.
+
+        Args:
+            file_path: The name of the file to write to.
+            file_content: The content to write to the file.
+            overwrite: If True, overwrite the file if it exists.
+
+        """
+
+    @abstractmethod
+    def read(
+        self, file_path: str, mode: str = "r", encoding: Optional[str] = None
+    ) -> str:
+        """Read content from a file.
+
+        Args:
+            file_path: The path to the file to read from.
+            mode: The mode in which to open the file.
+            encoding: The encoding to use for reading the file.
+
+        Returns:
+            str: The content of the file.
+
+        """
+
+    @abstractmethod
+    def read_bytes(self, file_path: str) -> bytes:
+        """Read binary content from a file.
+
+        Args:
+            file_path: The path to the file to read from.
+
+        Returns:
+            bytes: The binary content of the file.
+
+        """
+
+    @abstractmethod
+    def ls(self, path: str, recursive: bool = False) -> list[str]:
+        """List the contents of a directory.
+
+        Args:
+            path: The path to the directory.
+            recursive: If True, list the contents recursively.
+
+        Returns:
+            list[str]: A list of the contents of the directory.
+
+        """
+
+    @abstractmethod
+    def getcwd(self) -> str:
+        """Get the current working directory.
+
+        Returns:
+            str: The current working directory.
+
+        """
+
+    @abstractmethod
+    def remove_dir(self, path: str) -> None:
+        """Remove a directory and all its contents.
+
+        Args:
+            path: The path to the directory to remove.
+
+        """
+
+    @abstractmethod
+    def telemetry_path_files(self, path: str) -> Path:
+        """Get the path to the telemetry files.
+
+        Args:
+            path: The path to the telemetry directory.
+
+        Returns:
+            Path: The path object representing the telemetry files.
+
+        """

snowflake/snowpark_checkpoints_collector/io_utils/io_file_manager.py

@@ -0,0 +1,79 @@
+# 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.
+
+from pathlib import Path
+from typing import Optional
+
+from snowflake.snowpark_checkpoints_collector.io_utils import (
+    EnvStrategy,
+    IODefaultStrategy,
+)
+from snowflake.snowpark_checkpoints_collector.singleton import Singleton
+
+
+class IOFileManager(metaclass=Singleton):
+    def __init__(self, strategy: Optional[EnvStrategy] = None):
+        self.strategy = strategy or IODefaultStrategy()
+
+    def mkdir(self, path: str, exist_ok: bool = False) -> None:
+        return self.strategy.mkdir(path, exist_ok)
+
+    def folder_exists(self, path: str) -> bool:
+        return self.strategy.folder_exists(path)
+
+    def file_exists(self, path: str) -> bool:
+        return self.strategy.file_exists(path)
+
+    def write(self, file_path: str, file_content: str, overwrite: bool = True) -> None:
+        return self.strategy.write(file_path, file_content, overwrite)
+
+    def read(
+        self, file_path: str, mode: str = "r", encoding: Optional[str] = None
+    ) -> str:
+        return self.strategy.read(file_path, mode, encoding)
+
+    def read_bytes(self, file_path: str) -> bytes:
+        return self.strategy.read_bytes(file_path)
+
+    def ls(self, path: str, recursive: bool = False) -> list[str]:
+        return self.strategy.ls(path, recursive)
+
+    def getcwd(self) -> str:
+        return self.strategy.getcwd()
+
+    def remove_dir(self, path: str) -> None:
+        return self.strategy.remove_dir(path)
+
+    def telemetry_path_files(self, path: str) -> Path:
+        return self.strategy.telemetry_path_files(path)
+
+    def set_strategy(self, strategy: EnvStrategy):
+        """Set the strategy for file and directory operations.
+
+        Args:
+            strategy (EnvStrategy): The strategy to use for file and directory operations.
+
+        """
+        self.strategy = strategy
+
+
+def get_io_file_manager():
+    """Get the singleton instance of IOFileManager.
+
+    Returns:
+        IOFileManager: The singleton instance of IOFileManager.
+
+    """
+    return IOFileManager()

snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import glob
+import io
 import logging
 import os.path
 import time
@@ -25,6 +25,9 @@ from snowflake.snowpark import Session
 from snowflake.snowpark_checkpoints_collector.collection_common import (
     DOT_PARQUET_EXTENSION,
 )
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    get_io_file_manager,
+)
 
 
 STAGE_NAME = "CHECKPOINT_STAGE"
@@ -130,11 +133,13 @@ class SnowConnection:
         )
 
         def filter_files(name: str):
-            return os.path.isfile(name) and (filter_func(name) if filter_func else True)
+            return get_io_file_manager().file_exists(name) and (
+                filter_func(name) if filter_func else True
+            )
 
         target_dir = os.path.join(input_path, "**", "*")
         LOGGER.debug("Searching for files in '%s'", input_path)
-        files_collection = glob.glob(target_dir, recursive=True)
+        files_collection = get_io_file_manager().ls(target_dir, recursive=True)
 
         files = [file for file in files_collection if filter_files(file)]
         files_count = len(files)
@@ -152,17 +157,14 @@ class SnowConnection:
                 if not os.path.isabs(file)
                 else str(Path(file).resolve())
             )
-            # Snowflake required URI format for input in the put.
-            normalize_file_path = Path(file_full_path).as_uri()
             new_file_path = file_full_path.replace(input_path, folder_name)
             # as Posix to convert Windows dir to posix
             new_file_path = Path(new_file_path).as_posix()
             stage_file_path = STAGE_PATH_FORMAT.format(stage_name, new_file_path)
-            put_statement = PUT_FILE_IN_STAGE_STATEMENT_FORMAT.format(
-                normalize_file_path, stage_file_path
-            )
+            parquet_file = get_io_file_manager().read_bytes(file_full_path)
+            binary_parquet = io.BytesIO(parquet_file)
             LOGGER.info("Loading file '%s' to %s", file_full_path, stage_file_path)
-            self.session.sql(put_statement).collect()
+            self.session.file.put_stream(binary_parquet, stage_file_path)
 
     def create_table_from_parquet(
         self, table_name: str, stage_directory_path: str

snowflake/snowpark_checkpoints_collector/summary_stats_collector.py

@@ -12,12 +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 glob
 import json
 import logging
 import os
-import shutil
 
 from typing import Optional
 
@@ -54,6 +51,9 @@ from snowflake.snowpark_checkpoints_collector.column_collection import (
 from snowflake.snowpark_checkpoints_collector.column_pandera_checks import (
     PanderaColumnChecksManager,
 )
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    get_io_file_manager,
+)
 from snowflake.snowpark_checkpoints_collector.snow_connection_model import (
     SnowConnection,
 )
@@ -117,11 +117,10 @@ def collect_dataframe_checkpoint(
             "Checkpoint names must only contain alphanumeric characters, underscores and dollar signs."
         )
     if not is_checkpoint_enabled(normalized_checkpoint_name):
-        LOGGER.info(
-            "Checkpoint '%s' is disabled. Skipping collection.",
-            normalized_checkpoint_name,
+        raise Exception(
+            f"Checkpoint '{normalized_checkpoint_name}' is disabled. Please enable it in the checkpoints.json file.",
+            "In case you want to skip it, use the xcollect_dataframe_checkpoint method instead.",
         )
-        return
 
     LOGGER.info("Starting to collect checkpoint '%s'", normalized_checkpoint_name)
     LOGGER.debug("DataFrame size: %s rows", df.count())
@@ -184,6 +183,68 @@ def collect_dataframe_checkpoint(
         collection_point_result_manager.add_result(collection_point_result)
 
 
+@log
+def xcollect_dataframe_checkpoint(
+    df: SparkDataFrame,
+    checkpoint_name: str,
+    sample: Optional[float] = None,
+    mode: Optional[CheckpointMode] = None,
+    output_path: Optional[str] = None,
+) -> None:
+    """Skips the collection of metadata from a Dataframe checkpoint.
+
+    Args:
+        df (SparkDataFrame): The input Spark DataFrame to skip.
+        checkpoint_name (str): The name of the checkpoint.
+        sample (float, optional): Fraction of DataFrame to sample for schema inference.
+            Defaults to 1.0.
+        mode (CheckpointMode): The mode to execution the collection.
+            Defaults to CheckpointMode.Schema
+        output_path (str, optional): The output path to save the checkpoint.
+            Defaults to Current working Directory.
+
+    Raises:
+        Exception: Invalid mode value.
+        Exception: Invalid checkpoint name. Checkpoint names must only contain alphanumeric characters,
+                     underscores and dollar signs.
+
+    """
+    normalized_checkpoint_name = checkpoint_name_utils.normalize_checkpoint_name(
+        checkpoint_name
+    )
+    if normalized_checkpoint_name != checkpoint_name:
+        LOGGER.warning(
+            "Checkpoint name '%s' was normalized to '%s'",
+            checkpoint_name,
+            normalized_checkpoint_name,
+        )
+    is_valid_checkpoint_name = checkpoint_name_utils.is_valid_checkpoint_name(
+        normalized_checkpoint_name
+    )
+    if not is_valid_checkpoint_name:
+        raise Exception(
+            f"Invalid checkpoint name: {normalized_checkpoint_name}. "
+            "Checkpoint names must only contain alphanumeric characters, underscores and dollar signs."
+        )
+
+    LOGGER.warning(
+        "Checkpoint '%s' is disabled. Skipping collection.",
+        normalized_checkpoint_name,
+    )
+
+    collection_point_file_path = file_utils.get_collection_point_source_file_path()
+    collection_point_line_of_code = file_utils.get_collection_point_line_of_code()
+    collection_point_result = CollectionPointResult(
+        collection_point_file_path,
+        collection_point_line_of_code,
+        normalized_checkpoint_name,
+    )
+
+    collection_point_result.result = CollectionResult.SKIP
+    collection_point_result_manager = CollectionPointResultManager(output_path)
+    collection_point_result_manager.add_result(collection_point_result)
+
+
 @report_telemetry(params_list=["column_type_dict"])
 def _collect_dataframe_checkpoint_mode_schema(
     checkpoint_name: str,
@@ -321,8 +382,7 @@ def _generate_json_checkpoint_file(
     output_directory_path = file_utils.get_output_directory_path(output_path)
     checkpoint_file_path = os.path.join(output_directory_path, checkpoint_file_name)
     LOGGER.info("Writing DataFrame JSON schema file to '%s'", checkpoint_file_path)
-    with open(checkpoint_file_path, "w") as f:
-        f.write(dataframe_schema_contract)
+    get_io_file_manager().write(checkpoint_file_path, dataframe_schema_contract)
 
 
 @report_telemetry(params_list=["df"])
@@ -366,17 +426,17 @@ def generate_parquet_for_spark_df(spark_df: SparkDataFrame, output_path: str) ->
     ]
     converted_df = spark_df.select(new_cols)
 
-    if os.path.exists(output_path):
+    if get_io_file_manager().folder_exists(output_path):
         LOGGER.warning(
             "Output directory '%s' already exists. Deleting it...", output_path
         )
-        shutil.rmtree(output_path)
+        get_io_file_manager().remove_dir(output_path)
 
     LOGGER.info("Writing DataFrame to parquet files at '%s'", output_path)
     converted_df.write.parquet(output_path, mode="overwrite")
 
     target_dir = os.path.join(output_path, "**", f"*{DOT_PARQUET_EXTENSION}")
-    parquet_files = glob.glob(target_dir, recursive=True)
+    parquet_files = get_io_file_manager().ls(target_dir, recursive=True)
     parquet_files_count = len(parquet_files)
     if parquet_files_count == 0:
         raise Exception("No parquet files were generated.")

snowflake/snowpark_checkpoints_collector/utils/extra_config.py

@@ -22,13 +22,57 @@ from snowflake.snowpark_checkpoints_collector.collection_common import (
     SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR,
     CheckpointMode,
 )
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    get_io_file_manager,
+)
 
 
 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())
+    return os.environ.get(
+        SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR, get_io_file_manager().getcwd()
+    )
+
+
+def _set_conf_io_strategy() -> None:
+    try:
+        from snowflake.snowpark_checkpoints_collector.io_utils.io_default_strategy import (
+            IODefaultStrategy,
+        )
+        from snowflake.snowpark_checkpoints_configuration.io_utils.io_file_manager import (
+            EnvStrategy as ConfEnvStrategy,
+        )
+        from snowflake.snowpark_checkpoints_configuration.io_utils.io_file_manager import (
+            get_io_file_manager as get_conf_io_file_manager,
+        )
+
+        is_default_strategy = isinstance(
+            get_io_file_manager().strategy, IODefaultStrategy
+        )
+
+        if is_default_strategy:
+            return
+
+        class CustomConfEnvStrategy(ConfEnvStrategy):
+            def file_exists(self, path: str) -> bool:
+                return get_io_file_manager().file_exists(path)
+
+            def read(
+                self, file_path: str, mode: str = "r", encoding: Optional[str] = None
+            ) -> Optional[str]:
+                return get_io_file_manager().read(file_path, mode, encoding)
+
+            def getcwd(self) -> str:
+                return get_io_file_manager().getcwd()
+
+        get_conf_io_file_manager().set_strategy(CustomConfEnvStrategy())
+
+    except ImportError:
+        LOGGER.debug(
+            "snowpark-checkpoints-configuration is not installed. Cannot get a checkpoint metadata instance."
+        )
 
 
 # noinspection DuplicatedCode
@@ -38,6 +82,7 @@ def _get_metadata():
             CheckpointMetadata,
         )
 
+        _set_conf_io_strategy()
         path = _get_checkpoint_contract_file_path()
         LOGGER.debug("Loading checkpoint metadata from '%s'", path)
         metadata = CheckpointMetadata(path)

snowflake/snowpark_checkpoints_collector/utils/file_utils.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 import inspect
 import os
+import re
 import tempfile
 
 from typing import Optional
@@ -25,6 +26,9 @@ from snowflake.snowpark_checkpoints_collector.collection_common import (
     UNKNOWN_LINE_OF_CODE,
     UNKNOWN_SOURCE_FILE,
 )
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    get_io_file_manager,
+)
 
 
 def get_output_file_path(out_path: Optional[str] = None) -> str:
@@ -63,11 +67,13 @@ def get_output_directory_path(output_path: Optional[str] = None) -> str:
         str: returns the output directory path.
 
     """
-    current_working_directory_path = output_path if output_path else os.getcwd()
+    current_working_directory_path = (
+        output_path if output_path else get_io_file_manager().getcwd()
+    )
     checkpoints_output_directory_path = os.path.join(
         current_working_directory_path, SNOWPARK_CHECKPOINTS_OUTPUT_DIRECTORY_NAME
     )
-    os.makedirs(checkpoints_output_directory_path, exist_ok=True)
+    get_io_file_manager().mkdir(checkpoints_output_directory_path, exist_ok=True)
     return checkpoints_output_directory_path
 
 
@@ -79,7 +85,10 @@ def get_collection_point_source_file_path() -> str:
 
     """
     try:
-        collection_point_file_path = inspect.stack()[2].filename
+        stack_frame = _get_stack_frame()
+        if not stack_frame:
+            return UNKNOWN_SOURCE_FILE
+        collection_point_file_path = stack_frame.filename
         is_temporal_file_path = _is_temporal_path(collection_point_file_path)
         if is_temporal_file_path:
             ipynb_file_path_collection = _get_ipynb_file_path_collection()
@@ -95,15 +104,18 @@ def get_collection_point_source_file_path() -> str:
 
 
 def get_collection_point_line_of_code() -> int:
-    """Find the line of code of the source file where collection point it is.
+    """Find the line of code of the source file where collection point is.
 
     Returns:
         int: returns the line of code of the source file where collection point it is.
 
     """
     try:
-        collection_point_file_path = inspect.stack()[2].filename
-        collection_point_line_of_code = inspect.stack()[2].lineno
+        stack_frame = _get_stack_frame()
+        if not stack_frame:
+            return UNKNOWN_LINE_OF_CODE
+        collection_point_file_path = stack_frame.filename
+        collection_point_line_of_code = stack_frame.lineno
         is_temporal_file_path = _is_temporal_path(collection_point_file_path)
         if is_temporal_file_path:
             collection_point_line_of_code = UNKNOWN_LINE_OF_CODE
@@ -119,9 +131,24 @@ def _is_temporal_path(path: str) -> bool:
     return is_temporal_path
 
 
+def _get_stack_frame() -> inspect.FrameInfo:
+    batch = inspect.stack()[:7]
+    batch.reverse()
+    collect_frame_regex = r"(collect_dataframe_checkpoint)"
+
+    for frame in batch:
+        if (
+            frame.code_context is not None
+            and len(frame.code_context) >= 0
+            and re.search(collect_frame_regex, frame.code_context[0])
+        ):
+            return frame
+    return None
+
+
 def _get_ipynb_file_path_collection() -> list[str]:
-    current_working_directory_path = os.getcwd()
-    cwd_file_name_collection = os.listdir(current_working_directory_path)
+    current_working_directory_path = get_io_file_manager().getcwd()
+    cwd_file_name_collection = get_io_file_manager().ls(current_working_directory_path)
     ipynb_file_path_collection = []
     for file_name in cwd_file_name_collection:
         is_ipynb_file = file_name.endswith(DOT_IPYNB_EXTENSION)

snowflake/snowpark_checkpoints_collector/utils/telemetry.py

@@ -19,16 +19,34 @@ from sys import platform
 from typing import Any, Callable, Optional, TypeVar
 from uuid import getnode
 
-from snowflake.connector import (
-    SNOWFLAKE_CONNECTOR_VERSION,
-    time_util,
-)
-from snowflake.connector.constants import DIRS as SNOWFLAKE_DIRS
-from snowflake.connector.network import SnowflakeRestful
+from snowflake.connector.description import PLATFORM as CONNECTOR_PLATFORM
 from snowflake.connector.telemetry import TelemetryClient
 from snowflake.snowpark import VERSION as SNOWPARK_VERSION
 from snowflake.snowpark import dataframe as snowpark_dataframe
 from snowflake.snowpark.session import Session
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    get_io_file_manager,
+)
+
+
+try:
+    """
+    The following imports are used to log telemetry events in the Snowflake Connector.
+    """
+    from snowflake.connector import (
+        SNOWFLAKE_CONNECTOR_VERSION,
+        time_util,
+    )
+    from snowflake.connector.constants import DIRS as SNOWFLAKE_DIRS
+    from snowflake.connector.network import SnowflakeRestful
+except Exception:
+    """
+    Set default import values for the Snowflake Connector when using snowpark-checkpoints in stored procedures.
+    """
+    SNOWFLAKE_CONNECTOR_VERSION = ""
+    time_util = None
+    SNOWFLAKE_DIRS = ""
+    SnowflakeRestful = None
 
 
 try:
@@ -81,7 +99,7 @@ class TelemetryManager(TelemetryClient):
             path: path to write telemetry.
 
         """
-        os.makedirs(path, exist_ok=True)
+        get_io_file_manager().mkdir(str(path), exist_ok=True)
         self.sc_folder_path = path
 
     def sc_log_error(
@@ -189,7 +207,7 @@ class TelemetryManager(TelemetryClient):
 
         """
         try:
-            os.makedirs(self.sc_folder_path, exist_ok=True)
+            get_io_file_manager().mkdir(str(self.sc_folder_path), exist_ok=True)
             for event in batch:
                 message = event.get("message")
                 if message is not None:
@@ -199,8 +217,7 @@ class TelemetryManager(TelemetryClient):
                         f'_telemetry_{message.get("type")}.json'
                     )
                     json_content = self._sc_validate_folder_space(event)
-                    with open(file_path, "w") as json_file:
-                        json_file.write(json_content)
+                    get_io_file_manager().write(str(file_path), json_content)
         except Exception:
             pass
 
@@ -227,10 +244,10 @@ class TelemetryManager(TelemetryClient):
         if not self.sc_is_enabled or self.sc_is_testing or not self._rest:
             return
         batch = []
-        for file in self.sc_folder_path.glob("*.json"):
-            with open(file) as json_file:
-                data_dict = json.load(json_file)
-                batch.append(data_dict)
+        for file in get_io_file_manager().ls(f"{self.sc_folder_path}/*.json"):
+            json_content = get_io_file_manager().read(file)
+            data_dict = json.loads(json_content)
+            batch.append(data_dict)
         if batch == []:
             return
         body = {"logs": batch}
@@ -242,14 +259,17 @@ class TelemetryManager(TelemetryClient):
             timeout=5,
         )
         if ret.get("success"):
-            for file in self.sc_folder_path.glob("*.json"):
+            for file_path in get_io_file_manager().ls(f"{self.sc_folder_path}/*.json"):
+                file = get_io_file_manager().telemetry_path_files(file_path)
                 file.unlink()
 
     def _sc_is_telemetry_testing(self) -> bool:
         is_testing = os.getenv("SNOWPARK_CHECKPOINTS_TELEMETRY_TESTING") == "true"
         if is_testing:
             local_telemetry_path = (
-                Path(os.getcwd()) / "snowpark-checkpoints-output" / "telemetry"
+                Path(get_io_file_manager().getcwd())
+                / "snowpark-checkpoints-output"
+                / "telemetry"
             )
             self.set_sc_output_path(local_telemetry_path)
             self.sc_is_enabled = True
@@ -348,7 +368,7 @@ def _get_metadata() -> dict:
     }
 
 
-def _get_version() -> str:
+def _get_version() -> Optional[str]:
     """Get the version of the package.
 
     Returns:
@@ -359,11 +379,10 @@ def _get_version() -> str:
         directory_levels_up = 1
         project_root = Path(__file__).resolve().parents[directory_levels_up]
         version_file_path = project_root / VERSION_FILE_NAME
-        with open(version_file_path) as file:
-            content = file.read()
-            version_match = re.search(VERSION_VARIABLE_PATTERN, content, re.MULTILINE)
-            if version_match:
-                return version_match.group(1)
+        content = get_io_file_manager().read(str(version_file_path))
+        version_match = re.search(VERSION_VARIABLE_PATTERN, content, re.MULTILINE)
+        if version_match:
+            return version_match.group(1)
         return None
     except Exception:
         return None
@@ -379,7 +398,10 @@ def _get_folder_size(folder_path: Path) -> int:
         int: The size of the folder in bytes.
 
     """
-    return sum(f.stat().st_size for f in folder_path.glob("*.json") if f.is_file())
+    sum_size = 0
+    for f in get_io_file_manager().ls(f"{folder_path}/*.json"):
+        sum_size += get_io_file_manager().telemetry_path_files(f).stat().st_size
+    return sum_size
 
 
 def _free_up_space(folder_path: Path, max_size: int) -> None:
@@ -390,9 +412,13 @@ def _free_up_space(folder_path: Path, max_size: int) -> None:
         max_size (int): The maximum allowed size of the folder in bytes.
 
     """
-    files = sorted(folder_path.glob("*.json"), key=lambda f: f.stat().st_mtime)
+    files = sorted(
+        get_io_file_manager().ls(f"{folder_path}/*.json"),
+        key=lambda f: f.stat().st_mtime,
+    )
     current_size = _get_folder_size(folder_path)
-    for file in files:
+    for file_path in files:
+        file = get_io_file_manager().telemetry_path_files(file_path)
         if current_size <= max_size:
             break
         current_size -= file.stat().st_size
@@ -471,12 +497,22 @@ def get_load_json(json_schema: str) -> dict:
 
     """
     try:
-        with open(json_schema, encoding="utf-8") as file:
-            return json.load(file)
+        file_content = get_io_file_manager().read(json_schema, encoding="utf-8")
+        return json.loads(file_content)
     except (OSError, json.JSONDecodeError) as e:
         raise ValueError(f"Error reading JSON schema file: {e}") from None
 
 
+def _is_in_stored_procedure() -> bool:
+    """Check if the code is running in a stored procedure.
+
+    Returns:
+        bool: True if the code is running in a stored procedure, False otherwise.
+
+    """
+    return CONNECTOR_PLATFORM == "XP"
+
+
 def extract_parameters(
     func: Callable, args: tuple, kwargs: dict, params_list: Optional[list[str]]
 ) -> dict:
@@ -824,7 +860,10 @@ def report_telemetry(
             except Exception as err:
                 func_exception = err
 
-            if os.getenv("SNOWPARK_CHECKPOINTS_TELEMETRY_ENABLED") == "false":
+            if (
+                os.getenv("SNOWPARK_CHECKPOINTS_TELEMETRY_ENABLED") == "false"
+                or _is_in_stored_procedure()
+            ):
                 return result
             telemetry_event = None
             data = None

{snowpark_checkpoints_collectors-0.2.1.dist-info → snowpark_checkpoints_collectors-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-checkpoints-collectors
-Version: 0.2.1
+Version: 0.3.1
 Summary: Snowpark column and table statistics collection
 Project-URL: Bug Tracker, https://github.com/snowflakedb/snowpark-checkpoints/issues
 Project-URL: Source code, https://github.com/snowflakedb/snowpark-checkpoints/
@@ -30,6 +30,7 @@ Requires-Dist: pandera[io]==0.20.4
 Requires-Dist: snowflake-connector-python
 Requires-Dist: snowflake-snowpark-python>=1.23.0
 Provides-Extra: development
+Requires-Dist: certifi==2025.1.31; extra == 'development'
 Requires-Dist: coverage>=7.6.7; extra == 'development'
 Requires-Dist: deepdiff>=8.0.0; extra == 'development'
 Requires-Dist: hatchling==1.25.0; extra == 'development'
@@ -100,6 +101,28 @@ def collect_dataframe_checkpoint(
 - `output_path`: The output path to save the checkpoint, defaults to current working directory.
 
 
+### Skip DataFrame Checkpoint Collection
+
+
+
+```python
+from pyspark.sql import DataFrame as SparkDataFrame
+from snowflake.snowpark_checkpoints_collector.collection_common import CheckpointMode
+from typing import Optional
+
+# Signature of the function
+def xcollect_dataframe_checkpoint(
+    df: SparkDataFrame,
+    checkpoint_name: str,
+    sample: Optional[float] = None,
+    mode: Optional[CheckpointMode] = None,
+    output_path: Optional[str] = None,
+) -> None:
+    ...
+```
+
+The signature of the method is the same of `collect_dataframe_checkpoint`.
+
 ## Usage Example
 
 ### Schema mode

{snowpark_checkpoints_collectors-0.2.1.dist-info → snowpark_checkpoints_collectors-0.3.1.dist-info}/RECORD

@@ -1,11 +1,11 @@
-snowflake/snowpark_checkpoints_collector/__init__.py,sha256=GIESlH2W6g_qdcnyRqw9yjsvEkt0aniFvGixKlF4K7A,1096
-snowflake/snowpark_checkpoints_collector/__version__.py,sha256=jEnm4p_P4FqdYsTq3hnGQnhLZ4KwL0_Ew8fDF8BRL98,632
+snowflake/snowpark_checkpoints_collector/__init__.py,sha256=g4NemuA6Mj4O2jkK0yLQ8sEV3owHiiJnBEz_OWvlW1I,1179
+snowflake/snowpark_checkpoints_collector/__version__.py,sha256=uSRs7fRupFeQ-z3PtU_6qh6ry8YBaSAnEIAvLhJKUR8,632
 snowflake/snowpark_checkpoints_collector/collection_common.py,sha256=ff5vYffrTRjoJXZQvVQBaOlegAUj_vXBbl1IZidz8Qo,4510
 snowflake/snowpark_checkpoints_collector/singleton.py,sha256=7AgIHQBXVRvPBBCkmBplzkdrrm-xVWf_N8svzA2vF8E,836
-snowflake/snowpark_checkpoints_collector/summary_stats_collector.py,sha256=Jd881L0Ju3I-aGhj0dkZZ_ZjW7x-NXKqcMD4sCRCaDw,14619
+snowflake/snowpark_checkpoints_collector/summary_stats_collector.py,sha256=kRJpVRE9Iy_uqeIPT-__Aan-YLWxQbgSjkJ3w4LpvCc,17214
 snowflake/snowpark_checkpoints_collector/collection_result/model/__init__.py,sha256=jZzx29WzrjH7C_6ZsBGoe4PxbW_oM4uIjySS1axIM34,1000
-snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result.py,sha256=8xD9zGnFJ7Rz9RUXIys7JnV3kQD4mk8QwNOTxAihSjQ,2908
-snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result_manager.py,sha256=6XbjHiehEm_RN_9y2MRlr0MaSgk3cWTczwZEYqUHCpM,2565
+snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result.py,sha256=XelL7LughZpKl1B_6bJoKOc_PqQg3UleX6zdgVXqTus,2926
+snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result_manager.py,sha256=EY6WIIXRbvkTYC4bQn7jFALHh7D2PirVoiLZ5Kq8dNs,2659
 snowflake/snowpark_checkpoints_collector/column_collection/__init__.py,sha256=hpTh1V7hqBSHxNUqISwfxdz-NLD-7oZEMLXDUuRsoOU,783
 snowflake/snowpark_checkpoints_collector/column_collection/column_collector_manager.py,sha256=Vav_vbiipHFIAdHxeQG4ZK1BAmWTi_18hBnVeIeXFRs,9670
 snowflake/snowpark_checkpoints_collector/column_collection/model/__init__.py,sha256=d0WNMeayDyUKYFLLaVAMIC5Qt-DoWoWgOjj2ygJaHWA,2919
@@ -26,14 +26,18 @@ snowflake/snowpark_checkpoints_collector/column_collection/model/timestamp_colum
 snowflake/snowpark_checkpoints_collector/column_collection/model/timestamp_ntz_column_collector.py,sha256=glUUnCLgTbGiPLpF2pSZ11KCgKSpHDRt5uhi1ZT9bxA,2578
 snowflake/snowpark_checkpoints_collector/column_pandera_checks/__init__.py,sha256=JNZPOYx8rUTONGz_d7xyfAvEC2_umHmGkJLoNSATLs4,793
 snowflake/snowpark_checkpoints_collector/column_pandera_checks/pandera_column_checks_manager.py,sha256=X1Mm37DKt-WZ5AegvoUA3itU1nBUxvhBxvjO85QqcGE,7893
+snowflake/snowpark_checkpoints_collector/io_utils/__init__.py,sha256=RhADOBizQJq4CoewWYleuZZErthjzRMHNagObTW-IsI,984
+snowflake/snowpark_checkpoints_collector/io_utils/io_default_strategy.py,sha256=RG5uL6OM5R55AmyKhrNGw78nlIvSLy9iGw0Rd7WdOl8,1967
+snowflake/snowpark_checkpoints_collector/io_utils/io_env_strategy.py,sha256=kJMbg2VOKNXXdkGCt_tMMLGEZ2aUl1_nie1qYvx5M-c,3770
+snowflake/snowpark_checkpoints_collector/io_utils/io_file_manager.py,sha256=M17EtANswD5gcgGnmT13OImO_W1uH4K3ewu2CXL9aes,2597
 snowflake/snowpark_checkpoints_collector/snow_connection_model/__init__.py,sha256=kLjZId-aGCljK7lF6yeEw-syEqeTOJDxdXfpv9YxvZA,755
-snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py,sha256=odKGTzc0xov8WOgJSR6WmVs0IT-f6O4YoaLqH6CbbFo,7263
+snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py,sha256=r3IPnmDMb8151PTgE4YojOhWnxWGPLyBWlgFvvhOfRY,7314
 snowflake/snowpark_checkpoints_collector/utils/checkpoint_name_utils.py,sha256=Xc4k3JU6A96-79VFRR8NrNAUPeO3V1DEAhngg-hLlU4,1787
-snowflake/snowpark_checkpoints_collector/utils/extra_config.py,sha256=t8WakSiHA3sgnXxz0WXE7q2MG7czWlnSYB5XR9swIhs,3643
-snowflake/snowpark_checkpoints_collector/utils/file_utils.py,sha256=deetkhQZOB0GUxQJvUHw4Ridp_rNYiCqmK9li3uwBL0,4324
+snowflake/snowpark_checkpoints_collector/utils/extra_config.py,sha256=3kVf6WVA-EuyMpTO3ycTlXMSCHtytGtT6wkV4U2Hyjw,5195
+snowflake/snowpark_checkpoints_collector/utils/file_utils.py,sha256=5ztlNCv9GdSktUvtdfydv86cCFcmSXCdD4axZXJrOQQ,5125
 snowflake/snowpark_checkpoints_collector/utils/logging_utils.py,sha256=yyi6X5DqKeTg0HRhvsH6ymYp2P0wbnyKIzI2RzrQS7k,2278
-snowflake/snowpark_checkpoints_collector/utils/telemetry.py,sha256=7S0yFE3Zq96SEGmVuVbpYc_wtXIQUpL--6KfGoxwJcA,30837
-snowpark_checkpoints_collectors-0.2.1.dist-info/METADATA,sha256=xu9uudc0KMeTcai8EF89QVaJZRN00E3kH79aRzfRkw0,6003
-snowpark_checkpoints_collectors-0.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-snowpark_checkpoints_collectors-0.2.1.dist-info/licenses/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
-snowpark_checkpoints_collectors-0.2.1.dist-info/RECORD,,
+snowflake/snowpark_checkpoints_collector/utils/telemetry.py,sha256=ueN9vM8j5YNax7jMcnEj_UrgGkoeMv_hJHVKjN7hiJE,32161
+snowpark_checkpoints_collectors-0.3.1.dist-info/METADATA,sha256=dPrRK0GO-5U3bQUntCpYyRfG3J_z-v6bV9hX2s8CxdM,6613
+snowpark_checkpoints_collectors-0.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+snowpark_checkpoints_collectors-0.3.1.dist-info/licenses/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
+snowpark_checkpoints_collectors-0.3.1.dist-info/RECORD,,