import inspect
import json
import sys
import textwrap
import types
import warnings
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Tuple,
    Union,
    cast,
)
from uuid import uuid1

from absl import logging

from snowflake import connector, snowpark
from snowflake.ml._internal import telemetry
from snowflake.ml._internal.utils import (
    formatting,
    identifier,
    query_result_checker,
    spcs_attribution_utils,
    table_manager,
    uri,
)
from snowflake.ml.model import (
    _api as model_api,
    deploy_platforms,
    model_signature,
    type_hints as model_types,
)
from snowflake.ml.registry import _initial_schema, _schema_version_manager
from snowflake.snowpark._internal import utils as snowpark_utils

if TYPE_CHECKING:
    import pandas as pd

_DEFAULT_REGISTRY_NAME: str = "_SYSTEM_MODEL_REGISTRY"
_DEFAULT_SCHEMA_NAME: str = "_SYSTEM_MODEL_REGISTRY_SCHEMA"
_MODELS_TABLE_NAME: str = "_SYSTEM_REGISTRY_MODELS"
_METADATA_TABLE_NAME: str = "_SYSTEM_REGISTRY_METADATA"
_DEPLOYMENT_TABLE_NAME: str = "_SYSTEM_REGISTRY_DEPLOYMENTS"

# Metadata operation types.
_SET_METADATA_OPERATION: str = "SET"
_ADD_METADATA_OPERATION: str = "ADD"
_DROP_METADATA_OPERATION: str = "DROP"

# Metadata types.
_METADATA_ATTRIBUTE_DESCRIPTION: str = "DESCRIPTION"
_METADATA_ATTRIBUTE_METRICS: str = "METRICS"
_METADATA_ATTRIBUTE_REGISTRATION: str = "REGISTRATION"
_METADATA_ATTRIBUTE_TAGS: str = "TAGS"
_METADATA_ATTRIBUTE_DEPLOYMENT: str = "DEPLOYMENTS"
_METADATA_ATTRIBUTE_DELETION: str = "DELETION"

# Leaving out REGISTRATION/DEPLOYMENT events as they will be handled differently from all mutable attributes.
_LIST_METADATA_ATTRIBUTE: List[str] = [
    _METADATA_ATTRIBUTE_DESCRIPTION,
    _METADATA_ATTRIBUTE_METRICS,
    _METADATA_ATTRIBUTE_TAGS,
]
_TELEMETRY_PROJECT = "MLOps"
_TELEMETRY_SUBPROJECT = "ModelRegistry"

_STAGE_PREFIX = "@"


def _create_registry_database(
    session: snowpark.Session,
    database_name: str,
    statement_params: Dict[str, Any],
) -> None:
    """Private helper to create the model registry database.

    The creation will be skipped if the target database already exists.

    Args:
        session: Session object to communicate with Snowflake.
        database_name: Desired name of the model registry database.
        statement_params: Function usage statement parameters used in sql query executions.
    """
    registry_databases = session.sql(f"SHOW DATABASES LIKE '{identifier.get_unescaped_names(database_name)}'").collect(
        statement_params=statement_params
    )
    if len(registry_databases) > 0:
        logging.warning(f"The database {database_name} already exists. Skipping creation.")
        return

    session.sql(f"CREATE DATABASE {database_name}").collect(statement_params=statement_params)


def _create_registry_schema(
    session: snowpark.Session,
    database_name: str,
    schema_name: str,
    statement_params: Dict[str, Any],
) -> None:
    """Private helper to create the model registry schema.

    The creation will be skipped if the target schema already exists.

    Args:
        session: Session object to communicate with Snowflake.
        database_name: Desired name of the model registry database.
        schema_name: Desired name of the schema used by this model registry inside the database.
        statement_params: Function usage statement parameters used in sql query executions.
    """
    # The default PUBLIC schema is created by default so it might already exist even in a new database.
    registry_schemas = session.sql(
        f"SHOW SCHEMAS LIKE '{identifier.get_unescaped_names(schema_name)}' IN DATABASE {database_name}"
    ).collect(statement_params=statement_params)

    if len(registry_schemas) > 0:
        logging.warning(
            f"The schema {table_manager.get_fully_qualified_schema_name(database_name, schema_name)} already exists. "
            + "Skipping creation."
        )
        return

    session.sql(f"CREATE SCHEMA {table_manager.get_fully_qualified_schema_name(database_name, schema_name)}").collect(
        statement_params=statement_params
    )


def _create_registry_views(
    session: snowpark.Session,
    database_name: str,
    schema_name: str,
    registry_table_name: str,
    metadata_table_name: str,
    deployment_table_name: str,
    statement_params: Dict[str, Any],
) -> None:
    """Create views on underlying ModelRegistry tables.

    Args:
        session: Session object to communicate with Snowflake.
        database_name: Desired name of the model registry database.
        schema_name: Desired name of the schema used by this model registry inside the database.
        registry_table_name: Name for the main model registry table.
        metadata_table_name: Name for the metadata table used by the model registry.
        deployment_table_name: Name for the deployment event table.
        statement_params: Function usage statement parameters used in sql query executions.
    """
    fully_qualified_schema_name = table_manager.get_fully_qualified_schema_name(database_name, schema_name)

    # From the documentation: Each DDL statement executes as a separate transaction. Races should not be an issue.
    # https://docs.snowflake.com/en/sql-reference/transactions.html#ddl

    # Create a view on active permanent deployments.
    _create_active_permanent_deployment_view(
        session,
        fully_qualified_schema_name,
        registry_table_name,
        deployment_table_name,
        statement_params,
    )

    # Create views on most recent metadata items.
    metadata_view_name_prefix = identifier.concat_names([metadata_table_name, "_LAST_"])
    metadata_view_template = formatting.unwrap(
        """CREATE OR REPLACE TEMPORARY VIEW {database}.{schema}.{attribute_view} COPY GRANTS AS
            SELECT DISTINCT MODEL_ID, {select_expression} AS {final_attribute_name} FROM {metadata_table}
            WHERE ATTRIBUTE_NAME = '{attribute_name}'"""
    )

    # Create a separate view for the most recent item in each metadata column.
    metadata_view_names = []
    metadata_select_fields = []
    for attribute_name in _LIST_METADATA_ATTRIBUTE:
        view_name = identifier.concat_names([metadata_view_name_prefix, attribute_name])
        select_expression = (
            f"(LAST_VALUE(VALUE) OVER (PARTITION BY MODEL_ID ORDER BY EVENT_TIMESTAMP))['{attribute_name}']"
        )
        sql = metadata_view_template.format(
            database=database_name,
            schema=schema_name,
            select_expression=select_expression,
            attribute_view=view_name,
            attribute_name=attribute_name,
            final_attribute_name=attribute_name,
            metadata_table=metadata_table_name,
        )
        session.sql(sql).collect(statement_params=statement_params)
        metadata_view_names.append(view_name)
        metadata_select_fields.append(f"{view_name}.{attribute_name} AS {attribute_name}")

    # Create a special view for the registration timestamp.
    attribute_name = _METADATA_ATTRIBUTE_REGISTRATION
    final_attribute_name = identifier.concat_names([attribute_name, "_TIMESTAMP"])
    view_name = identifier.concat_names([metadata_view_name_prefix, attribute_name])
    create_registration_view_sql = metadata_view_template.format(
        database=database_name,
        schema=schema_name,
        select_expression="EVENT_TIMESTAMP",
        attribute_view=view_name,
        attribute_name=attribute_name,
        final_attribute_name=final_attribute_name,
        metadata_table=metadata_table_name,
    )
    session.sql(create_registration_view_sql).collect(statement_params=statement_params)
    metadata_view_names.append(view_name)
    metadata_select_fields.append(f"{view_name}.{final_attribute_name} AS {final_attribute_name}")

    metadata_views_join = " ".join(
        [
            "LEFT JOIN {view} ON ({view}.MODEL_ID = {registry_table}.ID)".format(
                view=view, registry_table=registry_table_name
            )
            for view in metadata_view_names
        ]
    )

    # Create view to combine all attributes.
    registry_view_name = identifier.concat_names([registry_table_name, "_VIEW"])
    metadata_select_fields_formatted = ",".join(metadata_select_fields)
    session.sql(
        f"""CREATE OR REPLACE TEMPORARY VIEW {fully_qualified_schema_name}.{registry_view_name} COPY GRANTS AS
                SELECT {registry_table_name}.*, {metadata_select_fields_formatted}
                FROM {registry_table_name} {metadata_views_join}"""
    ).collect(statement_params=statement_params)


def _create_active_permanent_deployment_view(
    session: snowpark.Session,
    fully_qualified_schema_name: str,
    registry_table_name: str,
    deployment_table_name: str,
    statement_params: Dict[str, Any],
) -> None:
    """Create a view which lists all available permanent deployments.

    Args:
        session: Session object to communicate with Snowflake.
        fully_qualified_schema_name: Schema name to the target table.
        registry_table_name: Name for the main model registry table.
        deployment_table_name: Name of the deployment table.
        statement_params: Function usage statement parameters used in sql query executions.
    """

    # Create a view on active permanent deployments
    # Active deployments are those whose last operation is not DROP.
    active_deployments_view_name = identifier.concat_names([deployment_table_name, "_VIEW"])
    active_deployments_view_expr = f"""
        CREATE OR REPLACE TEMPORARY VIEW {fully_qualified_schema_name}.{active_deployments_view_name}
        COPY GRANTS AS
        SELECT
            DEPLOYMENT_NAME,
            MODEL_ID,
            {registry_table_name}.NAME as MODEL_NAME,
            {registry_table_name}.VERSION as MODEL_VERSION,
            {deployment_table_name}.CREATION_TIME as CREATION_TIME,
            TARGET_METHOD,
            TARGET_PLATFORM,
            SIGNATURE,
            OPTIONS,
            STAGE_PATH,
            ROLE
        FROM {deployment_table_name}
        LEFT JOIN {registry_table_name}
            ON {deployment_table_name}.MODEL_ID = {registry_table_name}.ID
    """
    session.sql(active_deployments_view_expr).collect(statement_params=statement_params)


class ModelRegistry:
    """Model Management API."""

    def __init__(
        self,
        *,
        session: snowpark.Session,
        database_name: str = _DEFAULT_REGISTRY_NAME,
        schema_name: str = _DEFAULT_SCHEMA_NAME,
        create_if_not_exists: bool = False,
    ) -> None:
        """
        Opens an already-created registry.

        Args:
            session: Session object to communicate with Snowflake.
            database_name: Desired name of the model registry database.
            schema_name: Desired name of the schema used by this model registry inside the database.
            create_if_not_exists: create model registry if it's not exists already.
        """

        warnings.warn(
            """
The `snowflake.ml.registry.model_registry.ModelRegistry` has been deprecated starting from version 1.2.0.
It will stay in the Private Preview phase. For future implementations, kindly utilize `snowflake.ml.registry.Registry`,
except when specifically required. The old model registry will be removed once all its primary functionalities are
fully integrated into the new registry.
        """,
            DeprecationWarning,
            stacklevel=2,
        )
        if create_if_not_exists:
            create_model_registry(session=session, database_name=database_name, schema_name=schema_name)

        self._name = identifier.get_inferred_name(database_name)
        self._schema = identifier.get_inferred_name(schema_name)
        self._registry_table = identifier.get_inferred_name(_MODELS_TABLE_NAME)
        self._registry_table_view = identifier.concat_names([self._registry_table, "_VIEW"])
        self._metadata_table = identifier.get_inferred_name(_METADATA_TABLE_NAME)
        self._deployment_table = identifier.get_inferred_name(_DEPLOYMENT_TABLE_NAME)
        self._permanent_deployment_view = identifier.concat_names([self._deployment_table, "_VIEW"])
        self._permanent_deployment_stage = identifier.concat_names([self._deployment_table, "_STAGE"])
        self._session = session
        self._svm = _schema_version_manager.SchemaVersionManager(self._session, self._name, self._schema)

        # A in-memory deployment info cache to store information of temporary deployments
        # TODO(zhe): Use a temporary table to replace the in-memory cache.
        self._temporary_deployments: Dict[str, model_types.Deployment] = {}

        _initial_schema.check_access(self._session, self._name, self._schema)

        statement_params = self._get_statement_params(inspect.currentframe())
        self._svm.validate_schema_version(statement_params)

        _create_registry_views(
            session,
            self._name,
            self._schema,
            self._registry_table,
            self._metadata_table,
            self._deployment_table,
            statement_params,
        )

    # Private methods

    def _get_statement_params(self, frame: Optional[types.FrameType]) -> Dict[str, Any]:
        return telemetry.get_function_usage_statement_params(
            project=_TELEMETRY_PROJECT,
            subproject=_TELEMETRY_SUBPROJECT,
            function_name=telemetry.get_statement_params_full_func_name(frame, "ModelRegistry"),
        )

    def _get_new_unique_identifier(self) -> str:
        """Create new unique identifier.

        Returns:
            String identifier."""
        return uuid1().hex

    def _fully_qualified_registry_table_name(self) -> str:
        """Get the fully qualified name to the current registry table."""
        return table_manager.get_fully_qualified_table_name(self._name, self._schema, self._registry_table)

    def _fully_qualified_registry_view_name(self) -> str:
        """Get the fully qualified name to the current registry view."""
        return table_manager.get_fully_qualified_table_name(self._name, self._schema, self._registry_table_view)

    def _fully_qualified_metadata_table_name(self) -> str:
        """Get the fully qualified name to the current metadata table."""
        return table_manager.get_fully_qualified_table_name(self._name, self._schema, self._metadata_table)

    def _fully_qualified_deployment_table_name(self) -> str:
        """Get the fully qualified name to the current deployment table."""
        return table_manager.get_fully_qualified_table_name(self._name, self._schema, self._deployment_table)

    def _fully_qualified_permanent_deployment_view_name(self) -> str:
        """Get the fully qualified name to the permanent deployment view."""
        return table_manager.get_fully_qualified_table_name(self._name, self._schema, self._permanent_deployment_view)

    def _fully_qualified_schema_name(self) -> str:
        """Get the fully qualified name to the current registry schema."""
        return table_manager.get_fully_qualified_schema_name(self._name, self._schema)

    def _fully_qualified_deployment_name(self, deployment_name: str) -> str:
        """Get the fully qualified name to the given deployment."""
        return table_manager.get_fully_qualified_table_name(self._name, self._schema, deployment_name)

    def _insert_registry_entry(
        self, *, id: str, name: str, version: str, properties: Dict[str, Any]
    ) -> List[snowpark.Row]:
        """Insert a new row into the model registry table.

        Args:
            id: Model id to register.
            name: Model Name string.
            version: Model Version string.
            properties: Dictionary of properties corresponding to table columns.

        Returns:
            snowpark.Dataframe with the result of the operation.

        Raises:
            DataError: Mismatch between different id fields.
        """
        if not id:
            raise connector.DataError("Model ID is required but none given.")
        mandatory_args = {"ID": id, "NAME": name, "VERSION": version}
        for k, v in mandatory_args.items():
            if k not in properties:
                properties[k] = v
            else:
                if v and v != properties[k]:
                    raise connector.DataError(
                        formatting.unwrap(
                            f"""Parameter '{k.lower()}' is given and parameter 'properties' has the field '{k}' set but
                            the values do not match: {k.lower()}=="{v}" properties['{k}']=="{properties[k]}"."""
                        )
                    )
        # Could do a multi-table insert here with some pros and cons:
        # [PRO] Atomic insert across multiple tables.
        # [CON] Code logic becomes messy depending on which fields are set.
        # [CON] Harder to reuse existing methods like set_model_name.
        # Context: https://docs.snowflake.com/en/sql-reference/sql/insert-multi-table.html
        return table_manager.insert_table_entry(
            self._session,
            table=self._fully_qualified_registry_table_name(),
            columns=properties,
        )

    def _insert_metadata_entry(self, *, id: str, attribute: str, value: Any, operation: str) -> List[snowpark.Row]:
        """Insert a new row into the model metadata table.

        Args:
            id: Model id to register.
            attribute: name of the metadata attribute
            value: new value of the metadata attribute
            operation: the operation type of the metadata entry.

        Returns:
            snowpark.DataFrame with the result of the operation.

        Raises:
            DataError: Missing ID field.
        """
        if not id:
            raise connector.DataError("Model ID is required but none given.")

        columns: Dict[str, Any] = {}
        columns["EVENT_TIMESTAMP"] = formatting.SqlStr("CURRENT_TIMESTAMP()")
        columns["EVENT_ID"] = self._get_new_unique_identifier()
        columns["MODEL_ID"] = id
        columns["ROLE"] = self._session.get_current_role()
        columns["OPERATION"] = operation
        columns["ATTRIBUTE_NAME"] = attribute
        columns["VALUE"] = value

        return table_manager.insert_table_entry(
            self._session,
            table=self._fully_qualified_metadata_table_name(),
            columns=columns,
        )

    def _insert_deployment_entry(
        self,
        *,
        id: str,
        name: str,
        platform: str,
        stage_path: str,
        signature: Dict[str, Any],
        target_method: str,
        options: Optional[
            Union[
                model_types.WarehouseDeployOptions,
                model_types.SnowparkContainerServiceDeployOptions,
            ]
        ] = None,
    ) -> List[snowpark.Row]:
        """Insert a new row into the model deployment table.

        Each row in the deployment table is a deployment event.

        Args:
            id: Model id of the deployed model.
            name: Name of the deployment.
            platform: The deployment target destination.
            stage_path: The stage location where the deployment UDF is stored.
            signature: The model signature.
            target_method: The method name which is used for the deployment.
            options: The deployment options.

        Returns:
            A list of snowpark rows which is the insertion result.

        Raises:
            DataError: Missing ID field.
        """
        if not id:
            raise connector.DataError("Model ID is required but none given.")

        columns: Dict[str, Any] = {}
        columns["CREATION_TIME"] = formatting.SqlStr("CURRENT_TIMESTAMP()")
        columns["MODEL_ID"] = id
        columns["DEPLOYMENT_NAME"] = name
        columns["TARGET_PLATFORM"] = platform
        columns["STAGE_PATH"] = stage_path
        columns["ROLE"] = self._session.get_current_role()
        columns["SIGNATURE"] = signature
        columns["TARGET_METHOD"] = target_method
        columns["OPTIONS"] = options

        return table_manager.insert_table_entry(
            self._session,
            table=self._fully_qualified_deployment_table_name(),
            columns=columns,
        )

    def _prepare_deployment_stage(self) -> str:
        """Create a stage in the model registry for storing all permanent deployments.

        Returns:
            Path to the stage that was created.
        """
        schema = self._fully_qualified_schema_name()
        fully_qualified_deployment_stage_name = f"{schema}.{self._permanent_deployment_stage}"
        statement_params = self._get_statement_params(inspect.currentframe())
        self._session.sql(
            f"CREATE STAGE IF NOT EXISTS {fully_qualified_deployment_stage_name} "
            f"ENCRYPTION = (TYPE= 'SNOWFLAKE_SSE')"
        ).collect(statement_params=statement_params)
        return f"@{fully_qualified_deployment_stage_name}"

    def _prepare_model_stage(self, model_id: str) -> str:
        """Create a stage in the model registry for storing the model with the given id.

        Creating a permanent stage here since we do not have a way to switch a stage from temporary to permanent.
        This can result in orphaned stages in case the process fails. It might be better to try to create a
        temporary stage, attempt to perform all operations and convert the temp stage into permanent once the
        operation is complete.

        Args:
            model_id: Internal model ID string.

        Returns:
            Name of the stage that was created.

        Raises:
            DatabaseError: Indicates that something went wrong when creating the stage.
        """
        schema = self._fully_qualified_schema_name()

        # Uppercase the model_stage_name to avoid having to quote the the stage name.
        stage_name = model_id.upper()

        model_stage_name = f"SNOWML_MODEL_{stage_name}"
        fully_qualified_model_stage_name = f"{schema}.{model_stage_name}"
        statement_params = self._get_statement_params(inspect.currentframe())

        create_stage_result = self._session.sql(
            f"CREATE OR REPLACE STAGE {fully_qualified_model_stage_name} ENCRYPTION = (TYPE= 'SNOWFLAKE_SSE')"
        ).collect(statement_params=statement_params)
        if not create_stage_result:
            raise connector.DatabaseError("Unable to create stage for model. Operation returned not result.")
        if len(create_stage_result) != 1:
            raise connector.DatabaseError(
                "Unable to create stage for model. Creating the model stage returned unexpected result: {}.".format(
                    str(create_stage_result)
                )
            )

        return fully_qualified_model_stage_name

    def _get_fully_qualified_stage_name_from_uri(self, model_uri: str) -> Optional[str]:
        """Get fully qualified stage path pointed by the URI.

        Args:
            model_uri: URI for which stage file is needed.

        Returns:
            The fully qualified Snowflake stage location encoded by the given URI. Returns None if the URI is not
                pointing to a Snowflake stage.
        """
        raw_stage_path = uri.get_snowflake_stage_path_from_uri(model_uri)
        if not raw_stage_path:
            return None
        (db, schema, stage, _) = identifier.parse_schema_level_object_identifier(raw_stage_path)
        return identifier.get_schema_level_object_identifier(db, schema, stage)

    def _list_selected_models(
        self,
        *,
        id: Optional[str] = None,
        model_name: Optional[str] = None,
        model_version: Optional[str] = None,
    ) -> snowpark.DataFrame:
        """Retrieve the Snowpark dataframe of models matching the specified ID or (name and version).

        Args:
            id: Model ID string. Required if either name or version is None.
            model_name: Model Name string. Required if id is None.
            model_version: Model Version string. Required if id is None.

        Returns:
            A Snowpark dataframe representing the models that match the given constraints.
        """
        models = self.list_models()

        if id:
            filtered_models = models.filter(snowpark.Column("ID") == id)
        else:
            self._model_identifier_is_nonempty_or_raise(model_name, model_version)

            # The following two asserts is to satisfy mypy.
            assert model_name
            assert model_version

            filtered_models = models.filter(snowpark.Column("NAME") == model_name).filter(
                snowpark.Column("VERSION") == model_version
            )

        return cast(snowpark.DataFrame, filtered_models)

    def _validate_exact_one_result(
        self, selected_model: snowpark.DataFrame, model_identifier: str
    ) -> List[snowpark.Row]:
        """Validate the filtered model has exactly one result.

        Args:
            selected_model: A snowpark dataframe representing the models that are filtered out.
            model_identifier: A string which is used to filter the model.

        Returns:
            A snowpark row which contains the metadata of the filtered model

        Raises:
            KeyError: The target model doesn't exist.
            DataError: The target model is not unique.
        """
        statement_params = self._get_statement_params(inspect.currentframe())
        model_info = None
        try:
            model_info = (
                query_result_checker.ResultValidator(result=selected_model.collect(statement_params=statement_params))
                .has_dimensions(expected_rows=1)
                .validate()
            )
        except connector.DataError:
            if model_info is None or len(model_info) == 0:
                raise KeyError(f"The model {model_identifier} does not exist in the current registry.")
            else:
                raise connector.DataError(
                    formatting.unwrap(
                        f"""There are {len(model_info)} models {model_identifier}. This might indicate a problem with
                            the integrity of the model registry data."""
                    )
                )
        return model_info

    def _get_metadata_attribute(
        self,
        attribute: str,
        id: Optional[str] = None,
        model_name: Optional[str] = None,
        model_version: Optional[str] = None,
    ) -> Any:
        """Get the value of the given metadata attribute for target model with given (model name + model version) or id.

        Args:
            attribute: Name of the attribute to get.
            id: Model ID string. Required if either name or version is None.
            model_name: Model Name string. Required if id is None.
            model_version: Model Version string. Required if version is None.

        Returns:
            The value of the attribute that was requested. Can be None if the attribute is not set.
        """
        selected_models = self._list_selected_models(id=id, model_name=model_name, model_version=model_version)
        identifier = f"id {id}" if id else f"{model_name}/{model_version}"
        model_info = self._validate_exact_one_result(selected_models, identifier)
        return model_info[0][attribute]

    def _set_metadata_attribute(
        self,
        attribute: str,
        value: Any,
        id: Optional[str] = None,
        model_name: Optional[str] = None,
        model_version: Optional[str] = None,
        operation: str = _SET_METADATA_OPERATION,
        enable_model_presence_check: bool = True,
    ) -> None:
        """Set the value of the given metadata attribute for target model with given (model name + model version) or id.

        Args:
            attribute: Name of the attribute to set.
            value: Value to set.
            id: Model ID string. Required if either name or version is None.
            model_name: Model Name string. Required if id is None.
            model_version: Model Version string. Required if version is None.
            operation: the operation type of the metadata entry.
            enable_model_presence_check: If True, we will check if the model with the given ID is currently registered
                before setting the metadata attribute. False by default meaning that by default we will check.

        Raises:
            DataError: Failed to set the metadata attribute.
            KeyError: The target model doesn't exist
        """
        selected_models = self._list_selected_models(id=id, model_name=model_name, model_version=model_version)
        identifier = f"id {id}" if id else f"{model_name}/{model_version}"
        try:
            model_info = self._validate_exact_one_result(selected_models, identifier)
        except KeyError as e:
            # If the target model doesn't exist, raise the error only if enable_model_presence_check is True.
            if enable_model_presence_check:
                raise e

        if not id:
            id = model_info[0]["ID"]
        assert id is not None

        try:
            self._insert_metadata_entry(
                id=id,
                attribute=attribute,
                value={attribute: value},
                operation=operation,
            )
        except connector.DataError:
            raise connector.DataError(f"Setting {attribute} for mode id {id} failed.")

    def _model_identifier_is_nonempty_or_raise(self, model_name: Optional[str], model_version: Optional[str]) -> None:
        """Validate model_name and model_version are non-empty strings.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.

        Raises:
            ValueError: Raised when either model_name and model_version is empty.
        """
        if not model_name or not model_version:
            raise ValueError("model_name and model_version have to be non-empty strings.")

    def _get_model_id(self, model_name: str, model_version: str) -> str:
        """Get ID of the model with the given (model name + model version).

        Args:
            model_name: Model Name string.
            model_version: Model Version string.

        Returns:
            Id of the model.

        Raises:
            DataError: The requested model could not be found.
        """
        result = self._get_metadata_attribute("ID", model_name=model_name, model_version=model_version)
        if not result:
            raise connector.DataError(f"Model {model_name}/{model_version} doesn't exist.")
        return str(result)

    def _get_model_path(
        self,
        id: Optional[str] = None,
        model_name: Optional[str] = None,
        model_version: Optional[str] = None,
    ) -> str:
        """Get the stage path for the model with the given (model name + model version) or `id` from the registry.

        Args:
            id: Id of the model to deploy. Required if either model name or model version is None.
            model_name: Model Name string. Required if id is None.
            model_version: Model Version string. Required if id is None.

        Returns:
            str: Stage path for the model.

        Raises:
            DataError: When the model cannot be found or not be restored.
        """
        statement_params = self._get_statement_params(inspect.currentframe())
        selected_models = self._list_selected_models(id=id, model_name=model_name, model_version=model_version)
        identifier = f"id {id}" if id else f"{model_name}/{model_version}"
        model_info = self._validate_exact_one_result(selected_models, identifier)
        if not id:
            id = model_info[0]["ID"]
        model_uri = model_info[0]["URI"]

        if not uri.is_snowflake_stage_uri(model_uri):
            raise connector.DataError(
                f"Artifacts with URI scheme {uri.get_uri_scheme(model_uri)} are currently not supported."
            )

        model_stage_path = self._get_fully_qualified_stage_name_from_uri(model_uri=model_uri)

        # Currently we assume only the model is on the stage.
        model_file_list = self._session.sql(f"LIST @{model_stage_path}").collect(statement_params=statement_params)
        if len(model_file_list) == 0:
            raise connector.DataError(f"No files in model artifact for id {id} located at {model_uri}.")
        return f"{_STAGE_PREFIX}{model_stage_path}"

    def _log_model_path(
        self,
        model_name: str,
        model_version: str,
    ) -> Tuple[str, str]:
        """Generate a path in the Model Registry to store a model.

        Args:
            model_name: The given name for the model.
            model_version: Version string to be set for the model.

        Returns:
            String of the auto-generate unique model identifier and path to store it.
        """
        model_id = self._get_new_unique_identifier()

        # Copy model from local disk to remote stage.
        # TODO(zhe): Check if we could use the same stage for multiple models.
        fully_qualified_model_stage_name = self._prepare_model_stage(model_id=model_id)

        return model_id, fully_qualified_model_stage_name

    def _register_model_with_id(
        self,
        model_name: str,
        model_version: str,
        model_id: str,
        *,
        type: str,
        uri: str,
        input_spec: Optional[Dict[str, str]] = None,
        output_spec: Optional[Dict[str, str]] = None,
        description: Optional[str] = None,
        tags: Optional[Dict[str, str]] = None,
    ) -> None:
        """Helper function to register model metadata.

        Args:
            model_name: Name to be set for the model. The model name can NOT be changed after registration. The
                combination of name and version is expected to be unique inside the registry.
            model_version: Version string to be set for the model. The model version string can NOT be changed after
                model registration. The combination of name and version is expected to be unique inside the registry.
            model_id: The internal id for the model.
            type: Type of the model. Only a subset of types are supported natively.
            uri: Resource identifier pointing to the model artifact. There are no restrictions on the URI format,
                however only a limited set of URI schemes is supported natively.
            input_spec: The expected input schema of the model. Dictionary where the keys are
                expected column names and the values are the value types.
            output_spec: The expected output schema of the model. Dictionary where the keys
                are expected column names and the values are the value types.
            description: A description for the model. The description can be changed later.
            tags: Key-value pairs of tags to be set for this model. Tags can be modified
                after model registration.

        Raises:
            DataError: The given model already exists.
            DatabaseError: Unable to register the model properties into table.
        """
        new_model: Dict[Any, Any] = {}
        new_model["ID"] = model_id
        new_model["NAME"] = model_name
        new_model["VERSION"] = model_version
        new_model["TYPE"] = type
        new_model["URI"] = uri
        new_model["INPUT_SPEC"] = input_spec
        new_model["OUTPUT_SPEC"] = output_spec
        new_model["CREATION_TIME"] = formatting.SqlStr("CURRENT_TIMESTAMP()")
        new_model["CREATION_ROLE"] = self._session.get_current_role()
        new_model["CREATION_ENVIRONMENT_SPEC"] = {"python": ".".join(map(str, sys.version_info[:3]))}

        existing_model_nums = self._list_selected_models(model_name=model_name, model_version=model_version).count()
        if existing_model_nums:
            raise connector.DataError(
                f"Model {model_name}/{model_version} already exists. Unable to register the model."
            )

        if self._insert_registry_entry(id=model_id, name=model_name, version=model_version, properties=new_model):
            self._set_metadata_attribute(
                model_name=model_name,
                model_version=model_version,
                attribute=_METADATA_ATTRIBUTE_REGISTRATION,
                value=new_model,
            )
            if description:
                self.set_model_description(
                    model_name=model_name,
                    model_version=model_version,
                    description=description,
                )
            if tags:
                self._set_metadata_attribute(
                    _METADATA_ATTRIBUTE_TAGS,
                    value=tags,
                    model_name=model_name,
                    model_version=model_version,
                )
        else:
            raise connector.DatabaseError("Failed to insert the model properties to the registry table.")

    def _get_deployment(self, *, model_name: str, model_version: str, deployment_name: str) -> snowpark.Row:
        statement_params = self._get_statement_params(inspect.currentframe())
        deployment_lst = (
            self._session.sql(f"SELECT * FROM {self._fully_qualified_permanent_deployment_view_name()}")
            .filter(snowpark.Column("DEPLOYMENT_NAME") == deployment_name)
            .filter(snowpark.Column("MODEL_NAME") == model_name)
            .filter(snowpark.Column("MODEL_VERSION") == model_version)
        ).collect(statement_params=statement_params)
        if len(deployment_lst) == 0:
            raise KeyError(
                f"Unable to find deployment named {deployment_name} in the model {model_name}/{model_version}."
            )
        assert len(deployment_lst) == 1, "_get_deployment should return exactly 1 deployment"
        return cast(snowpark.Row, deployment_lst[0])

    # Registry operations

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def list_models(self) -> snowpark.DataFrame:
        """Lists models contained in the registry.

        Returns:
            snowpark.DataFrame with the list of models. Access is read-only through the snowpark.DataFrame.
            The resulting snowpark.dataframe will have an "id" column that uniquely identifies each model and can be
            used to reference the model when performing operations.
        """
        # Explicitly not calling collect.
        return self._session.sql(
            "SELECT * FROM {database}.{schema}.{view}".format(
                database=self._name, schema=self._schema, view=self._registry_table_view
            )
        )

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def set_tag(
        self,
        model_name: str,
        model_version: str,
        tag_name: str,
        tag_value: Optional[str] = None,
    ) -> None:
        """Set model tag to the model with value.

        If the model tag already exists, the tag value will be overwritten.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            tag_name: Desired tag name string.
            tag_value: (optional) New tag value string. If no value is given the value of the tag will be set to None.
        """
        # This method uses a read-modify-write pattern for setting tags.
        # TODO(amauser): Investigate the use of transactions to avoid race conditions.
        model_tags = self.get_tags(model_name=model_name, model_version=model_version)
        model_tags[tag_name] = tag_value
        self._set_metadata_attribute(
            _METADATA_ATTRIBUTE_TAGS,
            model_tags,
            model_name=model_name,
            model_version=model_version,
        )

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def remove_tag(self, model_name: str, model_version: str, tag_name: str) -> None:
        """Remove target model tag.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            tag_name: Desired tag name string.

        Raises:
            DataError: If the model does not have the requested tag.
        """
        # This method uses a read-modify-write pattern for updating tags.

        model_tags = self.get_tags(model_name=model_name, model_version=model_version)
        try:
            del model_tags[tag_name]
        except KeyError:
            raise connector.DataError(
                f"Model {model_name}/{model_version} has no tag named {tag_name}. Full list of tags: {model_tags}"
            )

        self._set_metadata_attribute(
            _METADATA_ATTRIBUTE_TAGS,
            model_tags,
            model_name=model_name,
            model_version=model_version,
        )

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def has_tag(
        self,
        model_name: str,
        model_version: str,
        tag_name: str,
        tag_value: Optional[str] = None,
    ) -> bool:
        """Check if a model has a tag with the given name and value.

        If no value is given, any value for the tag will return true.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            tag_name: Desired tag name string.
            tag_value: (optional) Tag value to check. If not value is given, only the presence of the tag will be
                checked.

        Returns:
            True if the tag or tag and value combination is present for the model with the given id, False otherwise.
        """
        tags = self.get_tags(model_name=model_name, model_version=model_version)
        has_tag = tag_name in tags
        if tag_value is None:
            return has_tag
        return has_tag and tags[tag_name] == str(tag_value)

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def get_tag_value(self, model_name: str, model_version: str, tag_name: str) -> Any:
        """Return the value of the tag for the model.

        The returned value can be None. If the tag does not exist, KeyError will be raised.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            tag_name: Desired tag name string.

        Returns:
            Value string of the tag or None, if no value is set for the tag.
        """
        return self.get_tags(model_name=model_name, model_version=model_version)[tag_name]

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def get_tags(self, model_name: Optional[str] = None, model_version: Optional[str] = None) -> Dict[str, Any]:
        """Get all tags and values stored for the target model.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.

        Returns:
            String-to-string dictionary containing all tags and values. The resulting dictionary can be empty.
        """
        # Snowpark snowpark.dataframe returns dictionary objects as strings. We need to convert it back to a dictionary
        # here.
        result = self._get_metadata_attribute(
            _METADATA_ATTRIBUTE_TAGS, model_name=model_name, model_version=model_version
        )

        if result:
            ret: Dict[str, Optional[str]] = json.loads(result)
            return ret
        else:
            return dict()

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def get_model_description(self, model_name: str, model_version: str) -> Optional[str]:
        """Get the description of the model.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.

        Returns:
            Description of the model or None.
        """
        result = self._get_metadata_attribute(
            _METADATA_ATTRIBUTE_DESCRIPTION,
            model_name=model_name,
            model_version=model_version,
        )
        return None if result is None else json.loads(result)

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def set_model_description(
        self,
        model_name: str,
        model_version: str,
        description: str,
    ) -> None:
        """Set the description of the model.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            description: Desired new model description.
        """
        self._set_metadata_attribute(
            _METADATA_ATTRIBUTE_DESCRIPTION,
            description,
            model_name=model_name,
            model_version=model_version,
        )

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def get_history(self) -> snowpark.DataFrame:
        """Return a dataframe with the history of operations performed on the model registry.

        The returned dataframe is order by time and can be filtered further.

        Returns:
            snowpark.DataFrame with the history of the model.
        """
        res = (
            self._session.table(self._fully_qualified_metadata_table_name())
            .order_by("EVENT_TIMESTAMP")
            .select_expr(
                "EVENT_TIMESTAMP",
                "EVENT_ID",
                "MODEL_ID",
                "ROLE",
                "OPERATION",
                "ATTRIBUTE_NAME",
                "VALUE[ATTRIBUTE_NAME]",
            )
        )
        return cast(snowpark.DataFrame, res)

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def get_model_history(
        self,
        model_name: str,
        model_version: str,
    ) -> snowpark.DataFrame:
        """Return a dataframe with the history of operations performed on the desired model.

        The returned dataframe is order by time and can be filtered further.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.

        Returns:
            snowpark.DataFrame with the history of the model.
        """
        id = self._get_model_id(model_name=model_name, model_version=model_version)
        return cast(
            snowpark.DataFrame,
            self.get_history().filter(snowpark.Column("MODEL_ID") == id),
        )

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def set_metric(
        self,
        model_name: str,
        model_version: str,
        metric_name: str,
        metric_value: object,
    ) -> None:
        """Set scalar model metric to value.

        If a metric with that name already exists for the model, the metric value will be overwritten.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            metric_name: Desired metric name.
            metric_value: New metric value.
        """
        # This method uses a read-modify-write pattern for setting tags.
        # TODO(amauser): Investigate the use of transactions to avoid race conditions.
        model_metrics = self.get_metrics(model_name=model_name, model_version=model_version)
        model_metrics[metric_name] = metric_value
        self._set_metadata_attribute(
            _METADATA_ATTRIBUTE_METRICS,
            model_metrics,
            model_name=model_name,
            model_version=model_version,
        )

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def remove_metric(
        self,
        model_name: str,
        model_version: str,
        metric_name: str,
    ) -> None:
        """Remove a specific metric entry from the model.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            metric_name: Desired metric name.

        Raises:
            DataError: If the model does not have the requested metric.
        """
        # This method uses a read-modify-write pattern for updating tags.

        model_metrics = self.get_metrics(model_name=model_name, model_version=model_version)
        try:
            del model_metrics[metric_name]
        except KeyError:
            raise connector.DataError(
                f"Model {model_name}/{model_version} has no metric named {metric_name}. "
                f"Full list of metrics: {model_metrics}"
            )

        self._set_metadata_attribute(
            _METADATA_ATTRIBUTE_METRICS,
            model_metrics,
            model_name=model_name,
            model_version=model_version,
        )

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def has_metric(self, model_name: str, model_version: str, metric_name: str) -> bool:
        """Check if a model has a metric with the given name.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            metric_name: Desired metric name.

        Returns:
            True if the metric is present for the model with the given id, False otherwise.
        """
        metrics = self.get_metrics(model_name=model_name, model_version=model_version)
        return metric_name in metrics

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def get_metric_value(self, model_name: str, model_version: str, metric_name: str) -> object:
        """Return the value of the given metric for the model.

        The returned value can be None. If the metric does not exist, KeyError will be raised.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            metric_name: Desired metric name.

        Returns:
            Value of the metric. Can be None if the metric was set to None.
        """
        return self.get_metrics(model_name=model_name, model_version=model_version)[metric_name]

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def get_metrics(self, model_name: str, model_version: str) -> Dict[str, object]:
        """Get all metrics and values stored for the given model.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.

        Returns:
            String-to-float dictionary containing all metrics and values. The resulting dictionary can be empty.
        """
        # Snowpark snowpark.dataframe returns dictionary objects as strings. We need to convert it back to a dictionary
        # here.
        result = self._get_metadata_attribute(
            _METADATA_ATTRIBUTE_METRICS,
            model_name=model_name,
            model_version=model_version,
        )

        if result:
            ret: Dict[str, object] = json.loads(result)
            return ret
        else:
            return dict()

    # Combined Registry and Repository operations.
    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def log_model(
        self,
        model_name: str,
        model_version: str,
        *,
        model: Any,
        description: Optional[str] = None,
        tags: Optional[Dict[str, str]] = None,
        conda_dependencies: Optional[List[str]] = None,
        pip_requirements: Optional[List[str]] = None,
        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
        sample_input_data: Optional[Any] = None,
        code_paths: Optional[List[str]] = None,
        options: Optional[model_types.BaseModelSaveOption] = None,
    ) -> Optional["ModelReference"]:
        """Uploads and register a model to the Model Registry.

        Args:
            model_name: The given name for the model. The combination (name + version) must be unique for each model.
            model_version: Version string to be set for the model. The combination (name + version) must be unique for
                each model.
            model: Local model object in a supported format.
            description: A description for the model. The description can be changed later.
            tags: string-to-string dictionary of tag names and values to be set for the model.
            conda_dependencies: List of Conda package specs. Use "[channel::]package [operator version]" syntax to
                specify a dependency. It is a recommended way to specify your dependencies using conda. When channel is
                not specified, defaults channel will be used. When deploying to Snowflake Warehouse, defaults channel
                would be replaced with the Snowflake Anaconda channel.
            pip_requirements: List of PIP package specs. Model will not be able to deploy to the warehouse if there is
                pip requirements.
            signatures: Signatures of the model, which is a mapping from target method name to signatures of input and
                output, which could be inferred by calling `infer_signature` method with sample input data.
            sample_input_data: Sample of the input data for the model.
            code_paths: Directory of code to import when loading and deploying the model.
            options: Additional options when saving the model.

        Raises:
            DataError: Raised when:
                1) the given model already exists;
            ValueError: Raised when:  # noqa: DAR402
                1) Signatures and sample_input_data are both not provided and model is not a
                    snowflake estimator.
            Exception: Raised when there is any error raised when saving the model.

        Returns:
            Model Reference . None if failed.
        """
        # Ideally, the whole operation should be a single transaction. Currently, transactions do not support stage
        # operations.

        statement_params = self._get_statement_params(inspect.currentframe())
        self._svm.validate_schema_version(statement_params)

        self._model_identifier_is_nonempty_or_raise(model_name, model_version)

        existing_model_nums = self._list_selected_models(model_name=model_name, model_version=model_version).count()
        if existing_model_nums:
            raise connector.DataError(f"Model {model_name}/{model_version} already exists. Unable to log the model.")
        model_id, fully_qualified_model_stage_name = self._log_model_path(
            model_name=model_name,
            model_version=model_version,
        )
        stage_path = f"{_STAGE_PREFIX}{fully_qualified_model_stage_name}"
        model = cast(model_types.SupportedModelType, model)
        try:
            model_composer = model_api.save_model(  # type: ignore[call-overload, misc]
                name=model_name,
                session=self._session,
                stage_path=stage_path,
                model=model,
                signatures=signatures,
                metadata=tags,
                conda_dependencies=conda_dependencies,
                pip_requirements=pip_requirements,
                sample_input_data=sample_input_data,
                code_paths=code_paths,
                options=options,
            )
        except Exception:
            # When model saving fails, clean up the model stage.
            query_result_checker.SqlResultValidator(
                self._session, f"DROP STAGE {fully_qualified_model_stage_name}"
            ).has_dimensions(expected_rows=1, expected_cols=1).validate()
            raise

        self._register_model_with_id(
            model_name=model_name,
            model_version=model_version,
            model_id=model_id,
            type=model_composer.packager.meta.model_type,
            uri=uri.get_uri_from_snowflake_stage_path(stage_path),
            description=description,
            tags=tags,
        )

        return ModelReference(registry=self, model_name=model_name, model_version=model_version)

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def load_model(self, model_name: str, model_version: str) -> Any:
        """Loads the model with the given (model_name + model_version) from the registry into memory.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.

        Returns:
            Restored model object.
        """
        warnings.warn(
            (
                "Please use with caution: "
                "Using `load_model` method requires you to have the EXACT same Python environments "
                "as the one when you logged the model. Any differences will potentially lead to errors.\n"
                "Also, if your model contains custom code imported using `code_paths` argument when logging, "
                "they will be added to your `sys.path`. It might lead to unexpected module importing issues. "
                "If you run into such kind of problems, you need to restart your Python or Notebook kernel."
            ),
            category=UserWarning,
            stacklevel=2,
        )
        remote_model_path = self._get_model_path(model_name=model_name, model_version=model_version)
        restored_model = None

        restored_model = model_api.load_model(session=self._session, stage_path=remote_model_path)

        return restored_model.packager.model

    # Repository Operations

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def deploy(
        self,
        model_name: str,
        model_version: str,
        *,
        deployment_name: str,
        target_method: Optional[str] = None,
        permanent: bool = False,
        platform: deploy_platforms.TargetPlatform = deploy_platforms.TargetPlatform.WAREHOUSE,
        options: Optional[
            Union[
                model_types.WarehouseDeployOptions,
                model_types.SnowparkContainerServiceDeployOptions,
            ]
        ] = None,
    ) -> model_types.Deployment:
        """Deploy the model with the given deployment name.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            deployment_name: name of the generated UDF.
            target_method: The method name to use in deployment. Can be omitted if only 1 method in the model.
            permanent: Whether the deployment is permanent or not. Permanent deployment will generate a permanent UDF.
                (Only applicable for Warehouse deployment)
            platform: Target platform to deploy the model to. Currently supported platforms are defined as enum in
                `snowflake.ml.model.deploy_platforms.TargetPlatform`
            options: Optional options for model deployment. Defaults to None.

        Returns:
            Deployment info.

        Raises:
            RuntimeError: Raised when parameters are not properly enabled when deploying to Warehouse with temporary UDF
            RuntimeError: Raised when deploying to SPCS with db/schema that starts with underscore.
        """
        statement_params = self._get_statement_params(inspect.currentframe())
        self._svm.validate_schema_version(statement_params)

        if options is None:
            options = {}

        deployment_stage_path = ""

        if platform == deploy_platforms.TargetPlatform.SNOWPARK_CONTAINER_SERVICES:
            if self._name.startswith("_") or self._schema.startswith("_"):
                error_message = """\
                Model deployment to Snowpark Container Service does not support a database/schema name that starts with
                an underscore. Please ensure you pass in a valid db/schema name when initializing the registry with:

                model_registry.create_model_registry(
                    session=session,
                    database_name=db,
                    schema_name=schema
                )

                registry = model_registry.ModelRegistry(
                    session=session,
                    database_name=db,
                    schema_name=schema
                )
                """
                raise RuntimeError(textwrap.dedent(error_message))
            permanent = True
            options = cast(model_types.SnowparkContainerServiceDeployOptions, options)
            deployment_stage_path = f"{self._prepare_deployment_stage()}/{deployment_name}/"
        elif platform == deploy_platforms.TargetPlatform.WAREHOUSE:
            options = cast(model_types.WarehouseDeployOptions, options)
            if permanent:
                # Every deployment-generated UDF should reside in its own unique directory. As long as each deployment
                # is allocated a distinct directory, multiple deployments can coexist within the same stage.
                # Given that each permanent deployment possesses a unique deployment_name, sharing the same stage does
                # not present any issues
                deployment_stage_path = (
                    options.get("permanent_udf_stage_location")
                    or f"{self._prepare_deployment_stage()}/{deployment_name}/"
                )
                options["permanent_udf_stage_location"] = deployment_stage_path

        remote_model_path = self._get_model_path(model_name=model_name, model_version=model_version)
        model_id = self._get_model_id(model_name, model_version)

        # https://snowflakecomputing.atlassian.net/browse/SNOW-858376
        # During temporary deployment on the Warehouse, Snowpark creates an unencrypted temporary stage for UDF-related
        # artifacts. However, UDF generation fails when importing from a mix of encrypted and unencrypted stages.
        # The following workaround copies model between stages (PrPr as of July 7th, 2023) to transfer the SSE
        # encrypted model zip from model stage to the temporary unencrypted stage.
        if not permanent and platform == deploy_platforms.TargetPlatform.WAREHOUSE:
            schema = self._fully_qualified_schema_name()
            unencrypted_stage = (
                f"@{schema}.{snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.STAGE)}"
            )
            self._session.sql(f"CREATE TEMPORARY STAGE {unencrypted_stage[1:]}").collect()
            try:
                self._session.sql(f"COPY FILES INTO {unencrypted_stage} from {remote_model_path}").collect()
            except Exception:
                raise RuntimeError(
                    "Temporary deployment to the warehouse is currently not supported. Please use "
                    "permanent deployment by setting the 'permanent' parameter to True"
                )
            remote_model_path = unencrypted_stage

        # Step 1: Deploy to get the UDF
        deployment_info = model_api.deploy(
            session=self._session,
            name=self._fully_qualified_deployment_name(deployment_name),
            platform=platform,
            target_method=target_method,
            stage_path=remote_model_path,
            deployment_stage_path=deployment_stage_path,
            model_id=model_id,
            options=options,
        )

        # Step 2: Record the deployment

        # Assert to convince mypy.
        assert deployment_info
        if permanent:
            self._insert_deployment_entry(
                id=model_id,
                name=deployment_name,
                platform=deployment_info["platform"].value,
                stage_path=deployment_stage_path,
                signature=deployment_info["signature"].to_dict(),
                target_method=deployment_info["target_method"],
                options=options,
            )

        self._set_metadata_attribute(
            _METADATA_ATTRIBUTE_DEPLOYMENT,
            {"name": deployment_name, "permanent": permanent},
            id=model_id,
            operation=_ADD_METADATA_OPERATION,
        )

        # Store temporary deployment information in the in-memory cache. This allows for future referencing and
        # tracking of its availability status.
        if not permanent:
            self._temporary_deployments[deployment_name] = deployment_info

        return deployment_info

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="1.0.1")
    def list_deployments(self, model_name: str, model_version: str) -> snowpark.DataFrame:
        """List all permanent deployments that originated from the given model.

        Temporary deployment info are currently not supported for listing.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.

        Returns:
            A snowpark dataframe that contains all deployments that associated with the given model.
        """
        deployments_df = (
            self._session.sql(f"SELECT * FROM {self._fully_qualified_permanent_deployment_view_name()}")
            .filter(snowpark.Column("MODEL_NAME") == model_name)
            .filter(snowpark.Column("MODEL_VERSION") == model_version)
        )
        res = deployments_df.select(
            deployments_df["MODEL_NAME"],
            deployments_df["MODEL_VERSION"],
            deployments_df["DEPLOYMENT_NAME"],
            deployments_df["CREATION_TIME"],
            deployments_df["TARGET_METHOD"],
            deployments_df["TARGET_PLATFORM"],
            deployments_df["SIGNATURE"],
            deployments_df["OPTIONS"],
            deployments_df["STAGE_PATH"],
            deployments_df["ROLE"],
        )
        return cast(snowpark.DataFrame, res)

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="1.0.1")
    def get_deployment(self, model_name: str, model_version: str, *, deployment_name: str) -> snowpark.DataFrame:
        """Get the permanent deployment with target name of the given model.

        Temporary deployment info are currently not supported.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            deployment_name: Deployment name string.

        Returns:
            A snowpark dataframe that contains the information of the target deployment.

        Raises:
            KeyError: Raised if the target deployment is not found.
        """
        deployment = self.list_deployments(model_name, model_version).filter(
            snowpark.Column("DEPLOYMENT_NAME") == deployment_name
        )
        if deployment.count() == 0:
            raise KeyError(
                f"Unable to find deployment named {deployment_name} in the model {model_name}/{model_version}."
            )
        return cast(snowpark.DataFrame, deployment)

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="1.0.1")
    def delete_deployment(self, model_name: str, model_version: str, *, deployment_name: str) -> None:
        """Delete the target permanent deployment of the given model.

        Deleting temporary deployment are currently not supported.
        Temporary deployment will get cleaned automatically when the current session closed.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            deployment_name: Name of the deployment that is getting deleted.

        """
        deployment = self._get_deployment(
            model_name=model_name,
            model_version=model_version,
            deployment_name=deployment_name,
        )

        # TODO(SNOW-759526): The following sequence should be a transaction.
        # Step 1: Drop the UDF
        self._session.sql(
            f"DROP FUNCTION IF EXISTS {self._fully_qualified_deployment_name(deployment_name)}(OBJECT)"
        ).collect()

        # Step 2: Remove the staged artifact
        self._session.sql(f"REMOVE {deployment['STAGE_PATH']}").collect()

        # Step 3: Delete the deployment from the deployment table
        query_result_checker.SqlResultValidator(
            self._session,
            f"""DELETE FROM {self._fully_qualified_deployment_table_name()}
                WHERE MODEL_ID='{deployment['MODEL_ID']}' AND DEPLOYMENT_NAME='{deployment_name}'
            """,
        ).deletion_success(expected_num_rows=1).validate()

        # Step 4: Record the delete event
        self._set_metadata_attribute(
            _METADATA_ATTRIBUTE_DEPLOYMENT,
            {"name": deployment_name},
            id=deployment["MODEL_ID"],
            operation=_DROP_METADATA_OPERATION,
        )

        # Optional Step 5: Delete Snowpark container service.
        if deployment["TARGET_PLATFORM"] == deploy_platforms.TargetPlatform.SNOWPARK_CONTAINER_SERVICES.value:
            service_name = identifier.get_schema_level_object_identifier(
                self._name, self._schema, f"service_{deployment['MODEL_ID']}"
            )
            spcs_attribution_utils.record_service_end(self._session, service_name)
            query_result_checker.SqlResultValidator(
                self._session,
                f"DROP SERVICE IF EXISTS {service_name}",
            ).validate()

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def delete_model(
        self,
        model_name: str,
        model_version: str,
        delete_artifact: bool = True,
    ) -> None:
        """Delete model with the given ID from the registry.

        The history of the model will still be preserved.

        Args:
            model_name: Model Name string.
            model_version: Model Version string.
            delete_artifact: If True, the underlying model artifact will also be deleted, not just the entry in
                the registry table.
        """

        # Check that a model with the given ID exists and there is only one of them.
        # TODO(amauser): The following sequence should be a transaction. Transactions currently cannot contain DDL
        # statements.
        model_info = None
        selected_models = self._list_selected_models(model_name=model_name, model_version=model_version)
        identifier = f"{model_name}/{model_version}"
        model_info = self._validate_exact_one_result(selected_models, identifier)
        id = model_info[0]["ID"]
        model_uri = model_info[0]["URI"]

        # Step 1/3: Delete the registry entry.
        query_result_checker.SqlResultValidator(
            self._session,
            f"DELETE FROM {self._fully_qualified_registry_table_name()} WHERE ID='{id}'",
        ).deletion_success(expected_num_rows=1).validate()

        # Step 2/3: Delete the artifact (if desired).
        if delete_artifact:
            if uri.is_snowflake_stage_uri(model_uri):
                stage_path = self._get_fully_qualified_stage_name_from_uri(model_uri)
                query_result_checker.SqlResultValidator(self._session, f"DROP STAGE {stage_path}").has_dimensions(
                    expected_rows=1, expected_cols=1
                ).validate()

        # Step 3/3: Record the deletion event.
        self._set_metadata_attribute(
            id=id,
            attribute=_METADATA_ATTRIBUTE_DELETION,
            value={"delete_artifact": True, "URI": model_uri},
            enable_model_presence_check=False,
        )


class ModelReference:
    """Wrapper class for ModelReference objects that proxy model metadata operations."""

    def _remove_arg_from_docstring(self, arg: str, docstring: Optional[str]) -> Optional[str]:
        """Remove the given parameter from a function docstring (Google convention)."""
        if docstring is None:
            return None
        docstring_lines = docstring.split("\n")

        args_section_start = None
        args_section_end = None
        args_section_indent = None
        arg_start = None
        arg_end = None
        arg_indent = None
        for i in range(len(docstring_lines)):
            line = docstring_lines[i]
            lstrip_line = line.lstrip()
            indent = len(line) - len(lstrip_line)

            if line.strip() == "Args:":
                # Starting the Args section of the docstring (assuming Google-style).
                args_section_start = i
                # logging.info("TEST: args_section_start=" + str(args_section_start))
                args_section_indent = indent
                continue

            # logging.info("TEST: " + lstrip_line)
            if args_section_start and lstrip_line.startswith(f"{arg}:"):
                # This is the arg we are looking for.
                arg_start = i
                # logging.info("TEST: arg_start=" + str(arg_start))
                arg_indent = indent
                continue

            if arg_start and not arg_end and indent == arg_indent:
                # We got the next arg, previous line was the last of the cut out arg docstring
                # and we do have other args. Saving arg_end for python slice/range notation.
                arg_end = i
                continue

            if arg_start and (len(lstrip_line) == 0 or indent == args_section_indent):
                # Arg section ends.
                args_section_end = i
                arg_end = arg_end if arg_end else i
                # We have learned everything we need to know, no need to continue.
                break

        if arg_start and not arg_end:
            arg_end = len(docstring_lines)

        if args_section_start and not args_section_end:
            args_section_end = len(docstring_lines)

        # Determine which lines from the "Args:" section of the docstring to skip or if we
        # should skip the entire section.
        keep_lines = set(range(len(docstring_lines)))
        if args_section_start:
            if arg_start == args_section_start + 1 and arg_end == args_section_end:
                # Removed arg was the only arg, remove the entire section.
                assert args_section_end
                keep_lines.difference_update(range(args_section_start, args_section_end))
            else:
                # Just remove the arg.
                assert arg_start
                assert arg_end
                keep_lines.difference_update(range(arg_start, arg_end))

        return "\n".join([docstring_lines[i] for i in sorted(keep_lines)])

    def __init__(
        self,
        *,
        registry: ModelRegistry,
        model_name: str,
        model_version: str,
    ) -> None:
        self._registry = registry
        self._id = registry._get_model_id(model_name=model_name, model_version=model_version)
        self._model_name = model_name
        self._model_version = model_version

        # Wrap all functions of the ModelRegistry that have an "id" parameter and bind that parameter
        # the the "_id" member of this class.
        if hasattr(self.__class__, "init_complete"):
            # Already did the generation of wrapped method.
            return

        for name, obj in self._registry.__class__.__dict__.items():
            if (
                not inspect.isfunction(obj)
                or "model_name" not in inspect.signature(obj).parameters
                or "model_version" not in inspect.signature(obj).parameters
            ):
                continue

            # Ensure that we are not silently overwriting existing functions.
            assert not hasattr(self.__class__, name)

            def build_method(m: Callable[..., Any]) -> Callable[..., Any]:
                return lambda self, *args, **kwargs: m(
                    self._registry,
                    self._model_name,
                    self._model_version,
                    *args,
                    **kwargs,
                )

            method = build_method(m=obj)
            setattr(self.__class__, name, method)

            docstring = self._remove_arg_from_docstring("model_name", obj.__doc__)
            if docstring and "model_version" in docstring:
                docstring = self._remove_arg_from_docstring("model_version", docstring)
            setattr(self.__class__.__dict__[name], "__doc__", docstring)  # noqa: B010

        setattr(self.__class__, "init_complete", True)  # noqa: B010

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    def get_name(self) -> str:
        return self._model_name

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    def get_version(self) -> str:
        return self._model_version

    @telemetry.send_api_usage_telemetry(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def predict(self, deployment_name: str, data: Any) -> "pd.DataFrame":
        """Predict using the deployed model in Snowflake.

        Args:
            deployment_name: name of the generated UDF.
            data: Data to run predict.

        Raises:
            ValueError: The deployment with given name haven't been deployed.

        Returns:
            A dataframe containing the result of prediction.
        """
        # We will search temporary deployments from the local in-memory cache.
        # If there is no hit, we try to search the remote deployment table.
        di = self._registry._temporary_deployments.get(deployment_name)

        statement_params = telemetry.get_function_usage_statement_params(
            project=_TELEMETRY_PROJECT,
            subproject=_TELEMETRY_SUBPROJECT,
            function_name=telemetry.get_statement_params_full_func_name(
                inspect.currentframe(), self.__class__.__name__
            ),
        )

        self._registry._svm.validate_schema_version(statement_params)

        if di:
            return model_api.predict(
                session=self._registry._session,
                deployment=di,
                X=data,
                statement_params=statement_params,
            )

        # Mypy enforce to refer to the registry for calling the function
        deployment_collect = self._registry.get_deployment(
            self._model_name, self._model_version, deployment_name=deployment_name
        ).collect(statement_params=statement_params)
        if not deployment_collect:
            raise ValueError(f"The deployment with name {deployment_name} haven't been deployed")
        deployment = deployment_collect[0]
        platform = deploy_platforms.TargetPlatform(deployment["TARGET_PLATFORM"])
        target_method = deployment["TARGET_METHOD"]
        signature = model_signature.ModelSignature.from_dict(json.loads(deployment["SIGNATURE"]))
        options_dict = cast(Dict[str, Any], json.loads(deployment["OPTIONS"]))
        platform_options = {
            deploy_platforms.TargetPlatform.WAREHOUSE: model_types.WarehouseDeployOptions,
            deploy_platforms.TargetPlatform.SNOWPARK_CONTAINER_SERVICES: (
                model_types.SnowparkContainerServiceDeployOptions
            ),
        }

        if platform not in platform_options:
            raise ValueError(f"Unsupported target Platform: {platform}")
        options = platform_options[platform](options_dict)
        di = model_types.Deployment(
            name=self._registry._fully_qualified_deployment_name(deployment_name),
            platform=platform,
            target_method=target_method,
            signature=signature,
            options=options,
        )
        return model_api.predict(
            session=self._registry._session,
            deployment=di,
            X=data,
            statement_params=statement_params,
        )


@telemetry.send_api_usage_telemetry(
    project=_TELEMETRY_PROJECT,
    subproject=_TELEMETRY_SUBPROJECT,
)
@snowpark._internal.utils.private_preview(version="0.2.0")
def create_model_registry(
    *,
    session: snowpark.Session,
    database_name: str = _DEFAULT_REGISTRY_NAME,
    schema_name: str = _DEFAULT_SCHEMA_NAME,
) -> bool:
    """Setup a new model registry. This should be run once per model registry by an administrator role.

    Args:
        session: Session object to communicate with Snowflake.
        database_name: Desired name of the model registry database.
        schema_name: Desired name of the schema used by this model registry inside the database.

    Returns:
        True if the creation of the model registry internal data structures was successful,
        False otherwise.
    """
    # Get the db & schema of the current session
    old_db = session.get_current_database()
    old_schema = session.get_current_schema()

    # These might be exposed as parameters in the future.
    database_name = identifier.get_inferred_name(database_name)
    schema_name = identifier.get_inferred_name(schema_name)

    statement_params = telemetry.get_function_usage_statement_params(
        project=_TELEMETRY_PROJECT,
        subproject=_TELEMETRY_SUBPROJECT,
        function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), ""),
    )
    try:
        _create_registry_database(session, database_name, statement_params)
        _create_registry_schema(session, database_name, schema_name, statement_params)

        svm = _schema_version_manager.SchemaVersionManager(session, database_name, schema_name)
        deployed_schema_version = svm.get_deployed_version(statement_params)
        if deployed_schema_version == _initial_schema._INITIAL_VERSION:
            # We do not know if registry is being created for the first time.
            # So let's start with creating initial schema, which is idempotent anyways.
            _initial_schema.create_initial_registry_tables(session, database_name, schema_name, statement_params)

        svm.try_upgrade(statement_params)

    finally:
        if not snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
            # Restore the db & schema to the original ones
            if old_db is not None and old_db != session.get_current_database():
                session.use_database(old_db)
            if old_schema is not None and old_schema != session.get_current_schema():
                session.use_schema(old_schema)
    return True
