snowpark-checkpoints-collectors 0.2.1__tar.gz → 0.3.0__tar.gz

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-checkpoints-collectors
-Version: 0.2.1
+Version: 0.3.0
 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'

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.0}/pyproject.toml

@@ -60,6 +60,7 @@ development = [
   "pyarrow>=18.0.0",
   "deepdiff>=8.0.0",
   "pyspark>=3.5.0",
+  "certifi==2025.1.31",
 ]
 
 [project.urls]

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.0}/src/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.0"

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.0}/src/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)

snowpark_checkpoints_collectors-0.3.0/src/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,
+)

snowpark_checkpoints_collectors-0.3.0/src/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)

snowpark_checkpoints_collectors-0.3.0/src/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.
+
+        """

snowpark_checkpoints_collectors-0.3.0/src/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()

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.0}/src/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

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.0}/src/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,
 )
@@ -321,8 +321,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 +365,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.")

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.0}/src/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)

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.0}/src/snowflake/snowpark_checkpoints_collector/utils/file_utils.py

@@ -25,6 +25,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 +66,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
 
 
@@ -120,8 +125,8 @@ def _is_temporal_path(path: str) -> bool:
 
 
 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)