snowpark-checkpoints-collectors 0.2.0rc1__py3-none-any.whl → 0.3.0__py3-none-any.whl

snowflake/snowpark_checkpoints_collector/utils/extra_config.py

@@ -0,0 +1,164 @@
+# 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 os
+
+from typing import Optional
+
+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, 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
+def _get_metadata():
+    try:
+        from snowflake.snowpark_checkpoints_configuration.checkpoint_metadata import (
+            CheckpointMetadata,
+        )
+
+        _set_conf_io_strategy()
+        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
+
+
+def is_checkpoint_enabled(checkpoint_name: str) -> bool:
+    """Check if a checkpoint is enabled.
+
+    Args:
+        checkpoint_name (str): The name of the checkpoint.
+
+    Returns:
+        bool: True if the checkpoint is enabled, False otherwise.
+
+    """
+    enabled, metadata = _get_metadata()
+    if enabled:
+        config = metadata.get_checkpoint(checkpoint_name)
+        return config.enabled
+    return True
+
+
+def get_checkpoint_sample(
+    checkpoint_name: str, sample: Optional[float] = None
+) -> float:
+    """Get the checkpoint sample.
+
+        Following this order first, the sample passed as argument, second, the sample from the checkpoint configuration,
+        third, the default sample value 1.0.
+
+    Args:
+        checkpoint_name (str): The name of the checkpoint.
+        sample (float, optional): The value passed to the function.
+
+    Returns:
+        float: returns the sample for that specific checkpoint.
+
+    """
+    default_sample = 1.0
+
+    enabled, metadata = _get_metadata()
+    if enabled:
+        config = metadata.get_checkpoint(checkpoint_name)
+        default_sample = config.sample if config.sample is not None else default_sample
+
+    return sample if sample is not None else default_sample
+
+
+def get_checkpoint_mode(
+    checkpoint_name: str, mode: Optional[CheckpointMode] = None
+) -> CheckpointMode:
+    """Get the checkpoint mode.
+
+        Following this order first, the mode passed as argument, second, the mode from the checkpoint configuration,
+        third, the default mode value 1.
+
+    Args:
+        checkpoint_name (str): The name of the checkpoint.
+        mode (int, optional): The value passed to the function.
+
+    Returns:
+        int: returns the mode for that specific checkpoint.
+
+    """
+    default_mode = CheckpointMode.SCHEMA
+
+    enabled, metadata = _get_metadata()
+    if enabled:
+        config = metadata.get_checkpoint(checkpoint_name)
+        default_mode = config.mode if config.mode is not None else default_mode
+
+    return mode if mode is not None else default_mode

snowflake/snowpark_checkpoints_collector/utils/file_utils.py

@@ -0,0 +1,137 @@
+# 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 inspect
+import os
+import tempfile
+
+from typing import Optional
+
+from snowflake.snowpark_checkpoints_collector.collection_common import (
+    COLLECTION_RESULT_FILE_NAME,
+    DOT_IPYNB_EXTENSION,
+    SNOWPARK_CHECKPOINTS_OUTPUT_DIRECTORY_NAME,
+    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:
+    """Get the output file path.
+
+    Args:
+        out_path (Optional[str], optional): the output path. Defaults to None.
+
+    Returns:
+        str: returns the output file path.
+
+    """
+    output_directory_path = get_output_directory_path(out_path)
+    output_file_path = os.path.join(output_directory_path, COLLECTION_RESULT_FILE_NAME)
+    return output_file_path
+
+
+def get_relative_file_path(path: str) -> str:
+    """Get the relative file path.
+
+    Args:
+        path (str): a file path.
+
+    Returns:
+        str: returns the relative file path of the given file.
+
+    """
+    relative_file_path = os.path.relpath(path)
+    return relative_file_path
+
+
+def get_output_directory_path(output_path: Optional[str] = None) -> str:
+    """Get the output directory path.
+
+    Returns:
+        str: returns the output directory path.
+
+    """
+    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
+    )
+    get_io_file_manager().mkdir(checkpoints_output_directory_path, exist_ok=True)
+    return checkpoints_output_directory_path
+
+
+def get_collection_point_source_file_path() -> str:
+    """Get the path of the source file where collection point it is.
+
+    Returns:
+        str: returns the path of source file where collection point it is.
+
+    """
+    try:
+        collection_point_file_path = inspect.stack()[2].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()
+            if len(ipynb_file_path_collection) == 1:
+                collection_point_file_path = ipynb_file_path_collection[0]
+            else:
+                collection_point_file_path = UNKNOWN_SOURCE_FILE
+
+        return collection_point_file_path
+
+    except Exception:
+        return UNKNOWN_SOURCE_FILE
+
+
+def get_collection_point_line_of_code() -> int:
+    """Find the line of code of the source file where collection point it 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
+        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
+        return collection_point_line_of_code
+
+    except Exception:
+        return UNKNOWN_LINE_OF_CODE
+
+
+def _is_temporal_path(path: str) -> bool:
+    temporal_directory_path = tempfile.gettempdir()
+    is_temporal_path = path.startswith(temporal_directory_path)
+    return is_temporal_path
+
+
+def _get_ipynb_file_path_collection() -> list[str]:
+    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)
+        if is_ipynb_file:
+            file_path = os.path.join(current_working_directory_path, file_name)
+            ipynb_file_path_collection.append(file_path)
+
+    return ipynb_file_path_collection

snowflake/snowpark_checkpoints_collector/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)