snowflake-ml-python 1.4.1__py3-none-any.whl → 1.5.1__py3-none-any.whl

snowflake/ml/model/_client/ops/model_ops.py

@@ -1,7 +1,7 @@
+import os
 import pathlib
 import tempfile
-from contextlib import contextmanager
-from typing import Any, Dict, Generator, List, Optional, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Union, cast
 
 import yaml
 
@@ -19,7 +19,9 @@ from snowflake.ml.model._model_composer.model_manifest import (
     model_manifest,
     model_manifest_schema,
 )
+from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_meta import model_meta
+from snowflake.ml.model._packager.model_runtime import model_runtime
 from snowflake.ml.model._signatures import snowpark_handler
 from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
@@ -72,37 +74,57 @@ class ModelOperator:
             and self._model_version_client == __value._model_version_client
         )
 
-    def prepare_model_stage_path(self, *, statement_params: Optional[Dict[str, Any]] = None) -> str:
+    def prepare_model_stage_path(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> str:
         stage_name = sql_identifier.SqlIdentifier(
             snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.STAGE)
         )
-        self._stage_client.create_tmp_stage(stage_name=stage_name, statement_params=statement_params)
-        return f"@{self._stage_client.fully_qualified_stage_name(stage_name)}/model"
+        self._stage_client.create_tmp_stage(
+            database_name=database_name,
+            schema_name=schema_name,
+            stage_name=stage_name,
+            statement_params=statement_params,
+        )
+        return f"@{self._stage_client.fully_qualified_object_name(database_name, schema_name, stage_name)}/model"
 
     def create_from_stage(
         self,
         composed_model: model_composer.ModelComposer,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         stage_path = str(composed_model.stage_path)
         if self.validate_existence(
+            database_name=database_name,
+            schema_name=schema_name,
             model_name=model_name,
             statement_params=statement_params,
         ):
             if self.validate_existence(
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 version_name=version_name,
                 statement_params=statement_params,
             ):
                 raise ValueError(
-                    f"Model {self._model_version_client.fully_qualified_model_name(model_name)} "
-                    f"version {version_name} already existed."
+                    "Model "
+                    f"{self._model_version_client.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                    f" version {version_name} already existed."
                 )
             else:
                 self._model_version_client.add_version_from_stage(
+                    database_name=database_name,
+                    schema_name=schema_name,
                     stage_path=stage_path,
                     model_name=model_name,
                     version_name=version_name,
@@ -110,6 +132,8 @@ class ModelOperator:
                 )
         else:
             self._model_version_client.create_from_stage(
+                database_name=database_name,
+                schema_name=schema_name,
                 stage_path=stage_path,
                 model_name=model_name,
                 version_name=version_name,
@@ -119,17 +143,23 @@ class ModelOperator:
     def show_models_or_versions(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> List[row.Row]:
         if model_name:
             return self._model_client.show_versions(
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 validate_result=False,
                 statement_params=statement_params,
             )
         else:
             return self._model_client.show_models(
+                database_name=database_name,
+                schema_name=schema_name,
                 validate_result=False,
                 statement_params=statement_params,
             )
@@ -137,10 +167,14 @@ class ModelOperator:
     def list_models_or_versions(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> List[sql_identifier.SqlIdentifier]:
         res = self.show_models_or_versions(
+            database_name=database_name,
+            schema_name=schema_name,
             model_name=model_name,
             statement_params=statement_params,
         )
@@ -153,12 +187,16 @@ class ModelOperator:
     def validate_existence(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> bool:
         if version_name:
             res = self._model_client.show_versions(
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 version_name=version_name,
                 validate_result=False,
@@ -166,6 +204,8 @@ class ModelOperator:
             )
         else:
             res = self._model_client.show_models(
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 validate_result=False,
                 statement_params=statement_params,
@@ -175,12 +215,16 @@ class ModelOperator:
     def get_comment(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> str:
         if version_name:
             res = self._model_client.show_versions(
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 version_name=version_name,
                 statement_params=statement_params,
@@ -188,6 +232,8 @@ class ModelOperator:
             col_name = self._model_client.MODEL_VERSION_COMMENT_COL_NAME
         else:
             res = self._model_client.show_models(
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 statement_params=statement_params,
             )
@@ -198,6 +244,8 @@ class ModelOperator:
         self,
         *,
         comment: str,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, Any]] = None,
@@ -205,6 +253,8 @@ class ModelOperator:
         if version_name:
             self._model_version_client.set_comment(
                 comment=comment,
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 version_name=version_name,
                 statement_params=statement_params,
@@ -212,6 +262,8 @@ class ModelOperator:
         else:
             self._model_client.set_comment(
                 comment=comment,
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 statement_params=statement_params,
             )
@@ -219,25 +271,42 @@ class ModelOperator:
     def set_default_version(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         if not self.validate_existence(
-            model_name=model_name, version_name=version_name, statement_params=statement_params
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         ):
             raise ValueError(f"You cannot set version {version_name} as default version as it does not exist.")
         self._model_version_client.set_default_version(
-            model_name=model_name, version_name=version_name, statement_params=statement_params
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )
 
     def get_default_version(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> sql_identifier.SqlIdentifier:
-        res = self._model_client.show_models(model_name=model_name, statement_params=statement_params)[0]
+        res = self._model_client.show_models(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            statement_params=statement_params,
+        )[0]
         return sql_identifier.SqlIdentifier(
             res[self._model_client.MODEL_DEFAULT_VERSION_NAME_COL_NAME], case_sensitive=True
         )
@@ -245,14 +314,18 @@ class ModelOperator:
     def get_tag_value(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
-        tag_database_name: sql_identifier.SqlIdentifier,
-        tag_schema_name: sql_identifier.SqlIdentifier,
+        tag_database_name: Optional[sql_identifier.SqlIdentifier],
+        tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> Optional[str]:
         r = self._tag_client.get_tag_value(
-            module_name=model_name,
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
             tag_database_name=tag_database_name,
             tag_schema_name=tag_schema_name,
             tag_name=tag_name,
@@ -266,11 +339,15 @@ class ModelOperator:
     def show_tags(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, str]:
         tags_info = self._tag_client.get_tag_list(
-            module_name=model_name,
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
             statement_params=statement_params,
         )
         res: Dict[str, str] = {
@@ -286,14 +363,18 @@ class ModelOperator:
     def set_tag(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
-        tag_database_name: sql_identifier.SqlIdentifier,
-        tag_schema_name: sql_identifier.SqlIdentifier,
+        tag_database_name: Optional[sql_identifier.SqlIdentifier],
+        tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         tag_value: str,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         self._tag_client.set_tag_on_model(
+            database_name=database_name,
+            schema_name=schema_name,
             model_name=model_name,
             tag_database_name=tag_database_name,
             tag_schema_name=tag_schema_name,
@@ -305,13 +386,17 @@ class ModelOperator:
     def unset_tag(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
-        tag_database_name: sql_identifier.SqlIdentifier,
-        tag_schema_name: sql_identifier.SqlIdentifier,
+        tag_database_name: Optional[sql_identifier.SqlIdentifier],
+        tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         self._tag_client.unset_tag_on_model(
+            database_name=database_name,
+            schema_name=schema_name,
             model_name=model_name,
             tag_database_name=tag_database_name,
             tag_schema_name=tag_schema_name,
@@ -322,12 +407,16 @@ class ModelOperator:
     def get_model_version_manifest(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> model_manifest_schema.ModelManifestDict:
         with tempfile.TemporaryDirectory() as tmpdir:
             self._model_version_client.get_file(
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 version_name=version_name,
                 file_path=pathlib.PurePosixPath(model_manifest.ModelManifest.MANIFEST_FILE_REL_PATH),
@@ -337,16 +426,6 @@ class ModelOperator:
             mm = model_manifest.ModelManifest(pathlib.Path(tmpdir))
             return mm.load()
 
-    @contextmanager
-    def _enable_model_details(
-        self,
-        *,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> Generator[None, None, None]:
-        self._model_client.config_model_details(enable=True, statement_params=statement_params)
-        yield
-        self._model_client.config_model_details(enable=False, statement_params=statement_params)
-
     @staticmethod
     def _match_model_spec_with_sql_functions(
         sql_functions_names: List[sql_identifier.SqlIdentifier], target_methods: List[str]
@@ -370,68 +449,75 @@ class ModelOperator:
     def get_functions(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> List[model_manifest_schema.ModelFunctionInfo]:
-        with self._enable_model_details(statement_params=statement_params):
-            raw_model_spec_res = self._model_client.show_versions(
-                model_name=model_name,
-                version_name=version_name,
-                check_model_details=True,
-                statement_params=statement_params,
-            )[0][self._model_client.MODEL_VERSION_MODEL_SPEC_COL_NAME]
-            model_spec_dict = yaml.safe_load(raw_model_spec_res)
-            model_spec = model_meta.ModelMetadata._validate_model_metadata(model_spec_dict)
-            show_functions_res = self._model_version_client.show_functions(
-                model_name=model_name,
-                version_name=version_name,
-                statement_params=statement_params,
+        raw_model_spec_res = self._model_client.show_versions(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            check_model_details=True,
+            statement_params={**(statement_params or {}), "SHOW_MODEL_DETAILS_IN_SHOW_VERSIONS_IN_MODEL": True},
+        )[0][self._model_client.MODEL_VERSION_MODEL_SPEC_COL_NAME]
+        model_spec_dict = yaml.safe_load(raw_model_spec_res)
+        model_spec = model_meta.ModelMetadata._validate_model_metadata(model_spec_dict)
+        show_functions_res = self._model_version_client.show_functions(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
+        )
+        function_names_and_types = []
+        for r in show_functions_res:
+            function_name = sql_identifier.SqlIdentifier(
+                r[self._model_version_client.FUNCTION_NAME_COL_NAME], case_sensitive=True
             )
-            function_names_and_types = []
-            for r in show_functions_res:
-                function_name = sql_identifier.SqlIdentifier(
-                    r[self._model_version_client.FUNCTION_NAME_COL_NAME], case_sensitive=True
-                )
 
-                function_type = model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
-                try:
-                    return_type = r[self._model_version_client.FUNCTION_RETURN_TYPE_COL_NAME]
-                except KeyError:
-                    pass
-                else:
-                    if "TABLE" in return_type:
-                        function_type = model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value
-
-                function_names_and_types.append((function_name, function_type))
-
-            signatures = model_spec["signatures"]
-            function_names = [name for name, _ in function_names_and_types]
-            function_name_mapping = ModelOperator._match_model_spec_with_sql_functions(
-                function_names, list(signatures.keys())
-            )
+            function_type = model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
+            try:
+                return_type = r[self._model_version_client.FUNCTION_RETURN_TYPE_COL_NAME]
+            except KeyError:
+                pass
+            else:
+                if "TABLE" in return_type:
+                    function_type = model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value
 
-            return [
-                model_manifest_schema.ModelFunctionInfo(
-                    name=function_name.identifier(),
-                    target_method=function_name_mapping[function_name],
-                    target_method_function_type=function_type,
-                    signature=model_signature.ModelSignature.from_dict(
-                        signatures[function_name_mapping[function_name]]
-                    ),
-                )
-                for function_name, function_type in function_names_and_types
-            ]
+            function_names_and_types.append((function_name, function_type))
+
+        signatures = model_spec["signatures"]
+        function_names = [name for name, _ in function_names_and_types]
+        function_name_mapping = ModelOperator._match_model_spec_with_sql_functions(
+            function_names, list(signatures.keys())
+        )
+
+        return [
+            model_manifest_schema.ModelFunctionInfo(
+                name=function_name.identifier(),
+                target_method=function_name_mapping[function_name],
+                target_method_function_type=function_type,
+                signature=model_signature.ModelSignature.from_dict(signatures[function_name_mapping[function_name]]),
+            )
+            for function_name, function_type in function_names_and_types
+        ]
 
     def invoke_method(
         self,
         *,
         method_name: sql_identifier.SqlIdentifier,
+        method_function_type: str,
         signature: model_signature.ModelSignature,
         X: Union[type_hints.SupportedDataType, dataframe.DataFrame],
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         strict_input_validation: bool = False,
+        partition_column: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, str]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         identifier_rule = model_signature.SnowparkIdentifierRule.INFERRED
@@ -469,15 +555,31 @@ class ModelOperator:
             if output_name in original_cols:
                 original_cols.remove(output_name)
 
-        df_res = self._model_version_client.invoke_method(
-            method_name=method_name,
-            input_df=s_df,
-            input_args=input_args,
-            returns=returns,
-            model_name=model_name,
-            version_name=version_name,
-            statement_params=statement_params,
-        )
+        if method_function_type == model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value:
+            df_res = self._model_version_client.invoke_function_method(
+                method_name=method_name,
+                input_df=s_df,
+                input_args=input_args,
+                returns=returns,
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+        elif method_function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
+            df_res = self._model_version_client.invoke_table_function_method(
+                method_name=method_name,
+                input_df=s_df,
+                input_args=input_args,
+                partition_column=partition_column,
+                returns=returns,
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
 
         if keep_order:
             df_res = df_res.sort(
@@ -486,7 +588,11 @@ class ModelOperator:
             )
 
         if not output_with_input_features:
-            df_res = df_res.drop(*original_cols)
+            cols_to_drop = original_cols
+            if partition_column is not None:
+                # don't drop partition column
+                cols_to_drop.remove(partition_column.identifier())
+            df_res = df_res.drop(*cols_to_drop)
 
         # Get final result
         if not isinstance(X, dataframe.DataFrame):
@@ -497,18 +603,97 @@ class ModelOperator:
     def delete_model_or_version(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         if version_name:
             self._model_version_client.drop_version(
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
                 version_name=version_name,
                 statement_params=statement_params,
             )
         else:
             self._model_client.drop_model(
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                statement_params=statement_params,
+            )
+
+    def rename(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        new_model_db: Optional[sql_identifier.SqlIdentifier],
+        new_model_schema: Optional[sql_identifier.SqlIdentifier],
+        new_model_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self._model_client.rename(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            new_model_db=new_model_db,
+            new_model_schema=new_model_schema,
+            new_model_name=new_model_name,
+            statement_params=statement_params,
+        )
+
+    # Map indicating in different modes, the path to list and download.
+    # The boolean value indicates if it is a directory,
+    MODEL_FILE_DOWNLOAD_PATTERN = {
+        "minimal": {
+            pathlib.PurePosixPath(model_composer.ModelComposer.MODEL_DIR_REL_PATH)
+            / model_meta.MODEL_METADATA_FILE: False,
+            pathlib.PurePosixPath(model_composer.ModelComposer.MODEL_DIR_REL_PATH) / model_env._DEFAULT_ENV_DIR: True,
+            pathlib.PurePosixPath(model_composer.ModelComposer.MODEL_DIR_REL_PATH)
+            / model_runtime.ModelRuntime.RUNTIME_DIR_REL_PATH: True,
+        },
+        "model": {pathlib.PurePosixPath(model_composer.ModelComposer.MODEL_DIR_REL_PATH): True},
+        "full": {pathlib.PurePosixPath(os.curdir): True},
+    }
+
+    def download_files(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        target_path: pathlib.Path,
+        mode: Literal["full", "model", "minimal"] = "model",
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        for remote_rel_path, is_dir in self.MODEL_FILE_DOWNLOAD_PATTERN[mode].items():
+            list_file_res = self._model_version_client.list_file(
+                database_name=database_name,
+                schema_name=schema_name,
                 model_name=model_name,
+                version_name=version_name,
+                file_path=remote_rel_path,
+                is_dir=is_dir,
                 statement_params=statement_params,
             )
+            file_list = [
+                pathlib.PurePosixPath(*pathlib.PurePosixPath(row.name).parts[2:])  # versions/<version_name>/...
+                for row in list_file_res
+            ]
+            for stage_file_path in file_list:
+                local_file_dir = target_path / stage_file_path.parent
+                local_file_dir.mkdir(parents=True, exist_ok=True)
+                self._model_version_client.get_file(
+                    database_name=database_name,
+                    schema_name=schema_name,
+                    model_name=model_name,
+                    version_name=version_name,
+                    file_path=stage_file_path,
+                    target_path=local_file_dir,
+                    statement_params=statement_params,
+                )

snowflake/ml/model/_client/sql/_base.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from snowflake.ml._internal.utils import identifier, sql_identifier
+from snowflake.snowpark import session
+
+
+class _BaseSQLClient:
+    def __init__(
+        self,
+        session: session.Session,
+        *,
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        self._session = session
+        self._database_name = database_name
+        self._schema_name = schema_name
+
+    def __eq__(self, __value: object) -> bool:
+        if not isinstance(__value, _BaseSQLClient):
+            return False
+        return self._database_name == __value._database_name and self._schema_name == __value._schema_name
+
+    def fully_qualified_object_name(
+        self,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        object_name: sql_identifier.SqlIdentifier,
+    ) -> str:
+        actual_database_name = database_name or self._database_name
+        actual_schema_name = schema_name or self._schema_name
+        return identifier.get_schema_level_object_identifier(
+            actual_database_name.identifier(), actual_schema_name.identifier(), object_name.identifier()
+        )