snowflake-ml-python 1.1.2__py3-none-any.whl → 1.2.1__py3-none-any.whl

snowflake/ml/model/_client/model/model_version_impl.py

@@ -3,11 +3,12 @@ from typing import Any, Callable, Dict, List, Optional, Union
 
 import pandas as pd
 
+from snowflake import connector
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.model import model_signature
-from snowflake.ml.model._client.model import model_method_info
 from snowflake.ml.model._client.ops import metadata_ops, model_ops
+from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.snowpark import dataframe
 
 _TELEMETRY_PROJECT = "MLOps"
@@ -49,14 +50,17 @@ class ModelVersion:
 
     @property
     def model_name(self) -> str:
+        """Return the name of the model to which the model version belongs, usable as a reference in SQL."""
         return self._model_name.identifier()
 
     @property
     def version_name(self) -> str:
+        """Return the name of the version to which the model version belongs, usable as a reference in SQL."""
         return self._version_name.identifier()
 
     @property
     def fully_qualified_model_name(self) -> str:
+        """Return the fully qualified name of the model to which the model version belongs."""
         return self._model_ops._model_version_client.fully_qualified_model_name(self._model_name)
 
     @property
@@ -65,6 +69,24 @@ class ModelVersion:
         subproject=_TELEMETRY_SUBPROJECT,
     )
     def description(self) -> str:
+        """The description for the model version. This is an alias of `comment`."""
+        return self.comment
+
+    @description.setter
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def description(self, description: str) -> None:
+        self.comment = description
+
+    @property
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def comment(self) -> str:
+        """The comment to the model version."""
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
@@ -75,18 +97,18 @@ class ModelVersion:
             statement_params=statement_params,
         )
 
-    @description.setter
+    @comment.setter
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def description(self, description: str) -> None:
+    def comment(self, comment: str) -> None:
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
         return self._model_ops.set_comment(
-            comment=description,
+            comment=comment,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -96,11 +118,11 @@ class ModelVersion:
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def list_metrics(self) -> Dict[str, Any]:
+    def show_metrics(self) -> Dict[str, Any]:
         """Show all metrics logged with the model version.
 
         Returns:
-            A dictionary showing the metrics
+            A dictionary showing the metrics.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
@@ -118,15 +140,15 @@ class ModelVersion:
         """Get the value of a specific metric.
 
         Args:
-            metric_name: The name of the metric
+            metric_name: The name of the metric.
 
         Raises:
-            KeyError: Raised when the requested metric name does not exist.
+            KeyError: When the requested metric name does not exist.
 
         Returns:
             The value of the metric.
         """
-        metrics = self.list_metrics()
+        metrics = self.show_metrics()
         if metric_name not in metrics:
             raise KeyError(f"Cannot find metric with name {metric_name}.")
         return metrics[metric_name]
@@ -136,17 +158,17 @@ class ModelVersion:
         subproject=_TELEMETRY_SUBPROJECT,
     )
     def set_metric(self, metric_name: str, value: Any) -> None:
-        """Set the value of a specific metric name
+        """Set the value of a specific metric.
 
         Args:
-            metric_name: The name of the metric
+            metric_name: The name of the metric.
             value: The value of the metric.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        metrics = self.list_metrics()
+        metrics = self.show_metrics()
         metrics[metric_name] = value
         self._model_ops._metadata_ops.save(
             metadata_ops.ModelVersionMetadataSchema(metrics=metrics),
@@ -166,13 +188,13 @@ class ModelVersion:
             metric_name: The name of the metric to be deleted.
 
         Raises:
-            KeyError: Raised when the requested metric name does not exist.
+            KeyError: When the requested metric name does not exist.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        metrics = self.list_metrics()
+        metrics = self.show_metrics()
         if metric_name not in metrics:
             raise KeyError(f"Cannot find metric with name {metric_name}.")
         del metrics[metric_name]
@@ -183,24 +205,12 @@ class ModelVersion:
             statement_params=statement_params,
         )
 
-    @telemetry.send_api_usage_telemetry(
-        project=_TELEMETRY_PROJECT,
-        subproject=_TELEMETRY_SUBPROJECT,
-    )
-    def list_methods(self) -> List[model_method_info.ModelMethodInfo]:
-        """List all method information in a model version that is callable.
-
-        Returns:
-            A list of ModelMethodInfo object containing the following information:
-            - name: The name of the method to be called (both in SQL and in Python SDK).
-            - target_method: The original method name in the logged Python object.
-            - Signature: Python signature of the original method.
-        """
+    # Only used when the model does not contains user_data with client SDK information.
+    def _legacy_show_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        # TODO(SNOW-986673, SNOW-986675): Avoid parsing manifest and meta file and put Python signature into user_data.
         manifest = self._model_ops.get_model_version_manifest(
             model_name=self._model_name,
             version_name=self._version_name,
@@ -211,7 +221,7 @@ class ModelVersion:
             version_name=self._version_name,
             statement_params=statement_params,
         )
-        return_methods_info: List[model_method_info.ModelMethodInfo] = []
+        return_functions_info: List[model_manifest_schema.ModelFunctionInfo] = []
         for method in manifest["methods"]:
             # Method's name is resolved so we need to use case_sensitive as True to get the user-facing identifier.
             method_name = sql_identifier.SqlIdentifier(method["name"], case_sensitive=True).identifier()
@@ -221,14 +231,48 @@ class ModelVersion:
             ), f"Get unexpected handler name {method['handler']}"
             target_method = method["handler"].split(".")[1]
             signature_dict = model_meta["signatures"][target_method]
-            method_info = model_method_info.ModelMethodInfo(
+            fi = model_manifest_schema.ModelFunctionInfo(
                 name=method_name,
                 target_method=target_method,
                 signature=model_signature.ModelSignature.from_dict(signature_dict),
             )
-            return_methods_info.append(method_info)
+            return_functions_info.append(fi)
+        return return_functions_info
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def show_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
+        """Show all functions information in a model version that is callable.
 
-        return return_methods_info
+        Returns:
+            A list of ModelFunctionInfo objects containing the following information:
+
+            - name: The name of the function to be called (both in SQL and in Python SDK).
+            - target_method: The original method name in the logged Python object.
+            - signature: Python signature of the original method.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        try:
+            client_data = self._model_ops.get_client_data_in_user_data(
+                model_name=self._model_name,
+                version_name=self._version_name,
+                statement_params=statement_params,
+            )
+            return [
+                model_manifest_schema.ModelFunctionInfo(
+                    name=fi["name"],
+                    target_method=fi["target_method"],
+                    signature=model_signature.ModelSignature.from_dict(fi["signature"]),
+                )
+                for fi in client_data["functions"]
+            ]
+        except (NotImplementedError, ValueError, connector.DataError):
+            return self._legacy_show_functions()
 
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
@@ -238,52 +282,52 @@ class ModelVersion:
         self,
         X: Union[pd.DataFrame, dataframe.DataFrame],
         *,
-        method_name: Optional[str] = None,
+        function_name: Optional[str] = None,
     ) -> Union[pd.DataFrame, dataframe.DataFrame]:
-        """Invoke a method in a model version object
+        """Invoke a method in a model version object.
 
         Args:
-            X: The input data. Could be pandas DataFrame or Snowpark DataFrame
-            method_name: The method name to run. It is the name you will use to call a method in SQL. Defaults to None.
-                It can only be None if there is only 1 method.
+            X: The input data, which could be a pandas DataFrame or Snowpark DataFrame.
+            function_name: The function name to run. It is the name used to call a function in SQL.
+                Defaults to None. It can only be None if there is only 1 method.
 
         Raises:
-            ValueError: No method with the corresponding name is available.
-            ValueError: There are more than 1 target methods available in the model but no method name specified.
+            ValueError: When no method with the corresponding name is available.
+            ValueError: When there are more than 1 target methods available in the model but no function name specified.
 
         Returns:
-            The prediction data.
+            The prediction data. It would be the same type dataframe as your input.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
 
-        methods: List[model_method_info.ModelMethodInfo] = self.list_methods()
-        if method_name:
-            req_method_name = sql_identifier.SqlIdentifier(method_name).identifier()
-            find_method: Callable[[model_method_info.ModelMethodInfo], bool] = (
+        functions: List[model_manifest_schema.ModelFunctionInfo] = self.show_functions()
+        if function_name:
+            req_method_name = sql_identifier.SqlIdentifier(function_name).identifier()
+            find_method: Callable[[model_manifest_schema.ModelFunctionInfo], bool] = (
                 lambda method: method["name"] == req_method_name
             )
-            target_method_info = next(
-                filter(find_method, methods),
+            target_function_info = next(
+                filter(find_method, functions),
                 None,
             )
-            if target_method_info is None:
+            if target_function_info is None:
                 raise ValueError(
-                    f"There is no method with name {method_name} available in the model"
+                    f"There is no method with name {function_name} available in the model"
                     f" {self.fully_qualified_model_name} version {self.version_name}"
                 )
-        elif len(methods) != 1:
+        elif len(functions) != 1:
             raise ValueError(
                 f"There are more than 1 target methods available in the model {self.fully_qualified_model_name}"
                 f" version {self.version_name}. Please specify a `method_name` when calling the `run` method."
             )
         else:
-            target_method_info = methods[0]
+            target_function_info = functions[0]
         return self._model_ops.invoke_method(
-            method_name=sql_identifier.SqlIdentifier(target_method_info["name"]),
-            signature=target_method_info["signature"],
+            method_name=sql_identifier.SqlIdentifier(target_function_info["name"]),
+            signature=target_function_info["signature"],
             X=X,
             model_name=self._model_name,
             version_name=self._version_name,

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

@@ -68,9 +68,7 @@ class MetadataOperator:
         version_info_list = self._model_client.show_versions(
             model_name=model_name, version_name=version_name, statement_params=statement_params
         )
-        assert len(version_info_list) == 1
-        version_info = version_info_list[0]
-        metadata_str = version_info.metadata
+        metadata_str = version_info_list[0][self._model_client.MODEL_VERSION_METADATA_COL_NAME]
         if not metadata_str:
             return {}
         res = json.loads(metadata_str)

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

@@ -1,16 +1,18 @@
+import json
 import pathlib
 import tempfile
 from typing import Any, Dict, List, Optional, Union, cast
 
 import yaml
 
-from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml._internal.utils import identifier, sql_identifier
 from snowflake.ml.model import model_signature, type_hints
 from snowflake.ml.model._client.ops import metadata_ops
 from snowflake.ml.model._client.sql import (
     model as model_sql,
     model_version as model_version_sql,
     stage as stage_sql,
+    tag as tag_sql,
 )
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._model_composer.model_manifest import (
@@ -19,7 +21,7 @@ from snowflake.ml.model._model_composer.model_manifest import (
 )
 from snowflake.ml.model._packager.model_meta import model_meta, model_meta_schema
 from snowflake.ml.model._signatures import snowpark_handler
-from snowflake.snowpark import dataframe, session
+from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
@@ -50,6 +52,11 @@ class ModelOperator:
             database_name=database_name,
             schema_name=schema_name,
         )
+        self._tag_client = tag_sql.ModuleTagSQLClient(
+            session,
+            database_name=database_name,
+            schema_name=schema_name,
+        )
         self._metadata_ops = metadata_ops.MetadataOperator(
             session,
             database_name=database_name,
@@ -109,22 +116,39 @@ class ModelOperator:
                 statement_params=statement_params,
             )
 
-    def list_models_or_versions(
+    def show_models_or_versions(
         self,
         *,
         model_name: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[sql_identifier.SqlIdentifier]:
+    ) -> List[row.Row]:
         if model_name:
-            res = self._model_client.show_versions(
+            return self._model_client.show_versions(
                 model_name=model_name,
+                validate_result=False,
                 statement_params=statement_params,
             )
         else:
-            res = self._model_client.show_models(
+            return self._model_client.show_models(
+                validate_result=False,
                 statement_params=statement_params,
             )
-        return [sql_identifier.SqlIdentifier(row.name, case_sensitive=True) for row in res]
+
+    def list_models_or_versions(
+        self,
+        *,
+        model_name: Optional[sql_identifier.SqlIdentifier] = None,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> List[sql_identifier.SqlIdentifier]:
+        res = self.show_models_or_versions(
+            model_name=model_name,
+            statement_params=statement_params,
+        )
+        if model_name:
+            col_name = self._model_client.MODEL_VERSION_NAME_COL_NAME
+        else:
+            col_name = self._model_client.MODEL_NAME_COL_NAME
+        return [sql_identifier.SqlIdentifier(row[col_name], case_sensitive=True) for row in res]
 
     def validate_existence(
         self,
@@ -137,11 +161,13 @@ class ModelOperator:
             res = self._model_client.show_versions(
                 model_name=model_name,
                 version_name=version_name,
+                validate_result=False,
                 statement_params=statement_params,
             )
         else:
             res = self._model_client.show_models(
                 model_name=model_name,
+                validate_result=False,
                 statement_params=statement_params,
             )
         return len(res) == 1
@@ -159,13 +185,14 @@ class ModelOperator:
                 version_name=version_name,
                 statement_params=statement_params,
             )
+            col_name = self._model_client.MODEL_VERSION_COMMENT_COL_NAME
         else:
             res = self._model_client.show_models(
                 model_name=model_name,
                 statement_params=statement_params,
             )
-        assert len(res) == 1
-        return cast(str, res[0].comment)
+            col_name = self._model_client.MODEL_COMMENT_COL_NAME
+        return cast(str, res[0][col_name])
 
     def set_comment(
         self,
@@ -189,6 +216,109 @@ class ModelOperator:
                 statement_params=statement_params,
             )
 
+    def set_default_version(
+        self,
+        *,
+        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
+        ):
+            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
+        )
+
+    def get_default_version(
+        self,
+        *,
+        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]
+        return sql_identifier.SqlIdentifier(
+            res[self._model_client.MODEL_DEFAULT_VERSION_NAME_COL_NAME], case_sensitive=True
+        )
+
+    def get_tag_value(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        tag_database_name: sql_identifier.SqlIdentifier,
+        tag_schema_name: 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,
+            tag_database_name=tag_database_name,
+            tag_schema_name=tag_schema_name,
+            tag_name=tag_name,
+            statement_params=statement_params,
+        )
+        value = r.TAG_VALUE
+        if value is None:
+            return value
+        return str(value)
+
+    def show_tags(
+        self,
+        *,
+        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,
+            statement_params=statement_params,
+        )
+        res: Dict[str, str] = {
+            identifier.get_schema_level_object_identifier(
+                sql_identifier.SqlIdentifier(r.TAG_DATABASE, case_sensitive=True),
+                sql_identifier.SqlIdentifier(r.TAG_SCHEMA, case_sensitive=True),
+                sql_identifier.SqlIdentifier(r.TAG_NAME, case_sensitive=True),
+            ): str(r.TAG_VALUE)
+            for r in tags_info
+        }
+        return res
+
+    def set_tag(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        tag_database_name: sql_identifier.SqlIdentifier,
+        tag_schema_name: 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(
+            model_name=model_name,
+            tag_database_name=tag_database_name,
+            tag_schema_name=tag_schema_name,
+            tag_name=tag_name,
+            tag_value=tag_value,
+            statement_params=statement_params,
+        )
+
+    def unset_tag(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        tag_database_name: sql_identifier.SqlIdentifier,
+        tag_schema_name: sql_identifier.SqlIdentifier,
+        tag_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self._tag_client.unset_tag_on_model(
+            model_name=model_name,
+            tag_database_name=tag_database_name,
+            tag_schema_name=tag_schema_name,
+            tag_name=tag_name,
+            statement_params=statement_params,
+        )
+
     def get_model_version_manifest(
         self,
         *,
@@ -228,6 +358,22 @@ class ModelOperator:
                 raw_model_meta = yaml.safe_load(f)
             return model_meta.ModelMetadata._validate_model_metadata(raw_model_meta)
 
+    def get_client_data_in_user_data(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> model_manifest_schema.SnowparkMLDataDict:
+        raw_user_data_json_string = self._model_client.show_versions(
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
+        )[0][self._model_client.MODEL_VERSION_USER_DATA_COL_NAME]
+        raw_user_data = json.loads(raw_user_data_json_string)
+        assert isinstance(raw_user_data, dict), "user data should be a dictionary"
+        return model_manifest.ModelManifest.parse_client_data_from_user_data(raw_user_data)
+
     def invoke_method(
         self,
         *,

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

@@ -1,10 +1,23 @@
 from typing import Any, Dict, List, Optional
 
-from snowflake.ml._internal.utils import identifier, sql_identifier
+from snowflake.ml._internal.utils import (
+    identifier,
+    query_result_checker,
+    sql_identifier,
+)
 from snowflake.snowpark import row, session
 
 
 class ModelSQLClient:
+    MODEL_NAME_COL_NAME = "name"
+    MODEL_COMMENT_COL_NAME = "comment"
+    MODEL_DEFAULT_VERSION_NAME_COL_NAME = "default_version_name"
+
+    MODEL_VERSION_NAME_COL_NAME = "name"
+    MODEL_VERSION_COMMENT_COL_NAME = "comment"
+    MODEL_VERSION_METADATA_COL_NAME = "metadata"
+    MODEL_VERSION_USER_DATA_COL_NAME = "user_data"
+
     def __init__(
         self,
         session: session.Session,
@@ -30,29 +43,56 @@ class ModelSQLClient:
         self,
         *,
         model_name: Optional[sql_identifier.SqlIdentifier] = None,
+        validate_result: bool = True,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> List[row.Row]:
         fully_qualified_schema_name = ".".join([self._database_name.identifier(), self._schema_name.identifier()])
         like_sql = ""
         if model_name:
             like_sql = f" LIKE '{model_name.resolved()}'"
-        res = self._session.sql(f"SHOW MODELS{like_sql} IN SCHEMA {fully_qualified_schema_name}")
 
-        return res.collect(statement_params=statement_params)
+        res = (
+            query_result_checker.SqlResultValidator(
+                self._session,
+                f"SHOW MODELS{like_sql} IN SCHEMA {fully_qualified_schema_name}",
+                statement_params=statement_params,
+            )
+            .has_column(ModelSQLClient.MODEL_NAME_COL_NAME, allow_empty=True)
+            .has_column(ModelSQLClient.MODEL_COMMENT_COL_NAME, allow_empty=True)
+            .has_column(ModelSQLClient.MODEL_DEFAULT_VERSION_NAME_COL_NAME, allow_empty=True)
+        )
+        if validate_result and model_name:
+            res = res.has_dimensions(expected_rows=1)
+
+        return res.validate()
 
     def show_versions(
         self,
         *,
         model_name: sql_identifier.SqlIdentifier,
         version_name: Optional[sql_identifier.SqlIdentifier] = None,
+        validate_result: bool = True,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> List[row.Row]:
         like_sql = ""
         if version_name:
             like_sql = f" LIKE '{version_name.resolved()}'"
-        res = self._session.sql(f"SHOW VERSIONS{like_sql} IN MODEL {self.fully_qualified_model_name(model_name)}")
 
-        return res.collect(statement_params=statement_params)
+        res = (
+            query_result_checker.SqlResultValidator(
+                self._session,
+                f"SHOW VERSIONS{like_sql} IN MODEL {self.fully_qualified_model_name(model_name)}",
+                statement_params=statement_params,
+            )
+            .has_column(ModelSQLClient.MODEL_VERSION_NAME_COL_NAME, allow_empty=True)
+            .has_column(ModelSQLClient.MODEL_VERSION_COMMENT_COL_NAME, allow_empty=True)
+            .has_column(ModelSQLClient.MODEL_VERSION_METADATA_COL_NAME, allow_empty=True)
+            .has_column(ModelSQLClient.MODEL_VERSION_USER_DATA_COL_NAME, allow_empty=True)
+        )
+        if validate_result and version_name:
+            res = res.has_dimensions(expected_rows=1)
+
+        return res.validate()
 
     def set_comment(
         self,
@@ -61,8 +101,11 @@ class ModelSQLClient:
         model_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        comment_sql = f"COMMENT ON MODEL {self.fully_qualified_model_name(model_name)} IS $${comment}$$"
-        self._session.sql(comment_sql).collect(statement_params=statement_params)
+        query_result_checker.SqlResultValidator(
+            self._session,
+            f"COMMENT ON MODEL {self.fully_qualified_model_name(model_name)} IS $${comment}$$",
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
     def drop_model(
         self,
@@ -70,6 +113,8 @@ class ModelSQLClient:
         model_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        self._session.sql(f"DROP MODEL {self.fully_qualified_model_name(model_name)}").collect(
-            statement_params=statement_params
-        )
+        query_result_checker.SqlResultValidator(
+            self._session,
+            f"DROP MODEL {self.fully_qualified_model_name(model_name)}",
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()