Fixes to snow ws migrate (#1515)

* Fix for templated names * Moved conversion to api module * Migrated the conversion logic * Fixes Fixes Fixes Fix
snowflakedb · Aug 30, 2024 · 6d2999a · 6d2999a
1 parent 840db3f
commit 6d2999a
Show file tree

Hide file tree

Showing 29 changed files with 526 additions and 156 deletions.
diff --git a/src/snowflake/cli/_plugins/snowpark/commands.py b/src/snowflake/cli/_plugins/snowpark/commands.py
@@ -40,7 +40,6 @@
     SnowparkObject,
     SnowparkObjectManager,
     StageToArtefactMapping,
-    is_name_a_templated_one,
 )
 from snowflake.cli._plugins.snowpark.package.anaconda_packages import (
     AnacondaPackages,
@@ -76,7 +75,6 @@
     DEFAULT_SIZE_LIMIT_MB,
 )
 from snowflake.cli.api.exceptions import (
-    NoProjectDefinitionError,
     SecretsWithoutExternalAccessIntegrationError,
 )
 from snowflake.cli.api.identifiers import FQN
@@ -86,6 +84,9 @@
     MessageResult,
     SingleQueryResult,
 )
+from snowflake.cli.api.project.definition_conversion import (
+    convert_project_definition_to_v2,
+)
 from snowflake.cli.api.project.schemas.entities.snowpark_entity import (
     FunctionEntityModel,
     ProcedureEntityModel,
@@ -94,10 +95,6 @@
     ProjectDefinition,
     ProjectDefinitionV2,
 )
-from snowflake.cli.api.project.schemas.snowpark.callable import (
-    FunctionSchema,
-    ProcedureSchema,
-)
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.connector import DictCursor, ProgrammingError
 from snowflake.connector.cursor import SnowflakeCursor
@@ -446,66 +443,8 @@ def describe(
     )
 
 
-def migrate_v1_snowpark_to_v2(pd: ProjectDefinition):
-    if not pd.snowpark:
-        raise NoProjectDefinitionError(
-            project_type="snowpark", project_root=get_cli_context().project_root
-        )
-
-    artifact_mapping = {"src": pd.snowpark.src}
-    if pd.snowpark.project_name:
-        artifact_mapping["dest"] = pd.snowpark.project_name
-
-    snowpark_shared_mixin = "snowpark_shared"
-    data: dict = {
-        "definition_version": "2",
-        "mixins": {
-            snowpark_shared_mixin: {
-                "stage": pd.snowpark.stage_name,
-                "artifacts": [artifact_mapping],
-            }
-        },
-        "entities": {},
-    }
-
-    for index, entity in enumerate([*pd.snowpark.procedures, *pd.snowpark.functions]):
-        identifier = {"name": entity.name}
-        if entity.database is not None:
-            identifier["database"] = entity.database
-        if entity.schema_name is not None:
-            identifier["schema"] = entity.schema_name
-
-        if is_name_a_templated_one(entity.name):
-            entity_name = f"snowpark_entity_{index}"
-        else:
-            entity_name = entity.name
-
-        v2_entity = {
-            "type": "function" if isinstance(entity, FunctionSchema) else "procedure",
-            "stage": pd.snowpark.stage_name,
-            "handler": entity.handler,
-            "returns": entity.returns,
-            "signature": entity.signature,
-            "runtime": entity.runtime,
-            "external_access_integrations": entity.external_access_integrations,
-            "secrets": entity.secrets,
-            "imports": entity.imports,
-            "identifier": identifier,
-            "meta": {"use_mixins": [snowpark_shared_mixin]},
-        }
-        if isinstance(entity, ProcedureSchema):
-            v2_entity["execute_as_caller"] = entity.execute_as_caller
-
-        data["entities"][entity_name] = v2_entity
-
-    if hasattr(pd, "env") and pd.env:
-        data["env"] = {k: v for k, v in pd.env.items()}
-
-    return ProjectDefinitionV2(**data)
-
-
 def _get_v2_project_definition(cli_context) -> ProjectDefinitionV2:
     pd = cli_context.project_definition
     if not pd.meets_version_requirement("2"):
-        pd = migrate_v1_snowpark_to_v2(pd)
+        pd = convert_project_definition_to_v2(pd)
     return pd
diff --git a/src/snowflake/cli/_plugins/snowpark/common.py b/src/snowflake/cli/_plugins/snowpark/common.py
@@ -23,7 +23,13 @@
 from snowflake.cli._plugins.snowpark.models import Requirement
 from snowflake.cli._plugins.snowpark.snowpark_project_paths import Artefact
 from snowflake.cli.api.console import cli_console
-from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.constants import (
+    INIT_TEMPLATE_VARIABLE_CLOSING,
+    INIT_TEMPLATE_VARIABLE_OPENING,
+    PROJECT_TEMPLATE_VARIABLE_CLOSING,
+    PROJECT_TEMPLATE_VARIABLE_OPENING,
+    ObjectType,
+)
 from snowflake.cli.api.project.schemas.entities.snowpark_entity import (
     ProcedureEntityModel,
     SnowparkEntityModel,
@@ -253,4 +259,10 @@ def _compare_imports(
 
 
 def is_name_a_templated_one(name: str) -> bool:
-    return ("<!" in name and "!>" in name) or ("<%" in name and "%>" in name)
+    return (
+        PROJECT_TEMPLATE_VARIABLE_OPENING in name
+        and PROJECT_TEMPLATE_VARIABLE_CLOSING in name
+    ) or (
+        INIT_TEMPLATE_VARIABLE_OPENING in name
+        and INIT_TEMPLATE_VARIABLE_CLOSING in name
+    )
diff --git a/src/snowflake/cli/_plugins/streamlit/commands.py b/src/snowflake/cli/_plugins/streamlit/commands.py
@@ -25,7 +25,6 @@
     add_object_command_aliases,
     scope_option,
 )
-from snowflake.cli._plugins.snowpark.common import is_name_a_templated_one
 from snowflake.cli._plugins.streamlit.manager import StreamlitManager
 from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.decorators import (
@@ -47,13 +46,12 @@
     MessageResult,
     SingleQueryResult,
 )
+from snowflake.cli.api.project.definition_conversion import (
+    convert_project_definition_to_v2,
+)
 from snowflake.cli.api.project.schemas.entities.streamlit_entity_model import (
     StreamlitEntityModel,
 )
-from snowflake.cli.api.project.schemas.project_definition import (
-    ProjectDefinition,
-    ProjectDefinitionV2,
-)
 
 app = SnowTyperFactory(
     name="streamlit",
@@ -140,7 +138,7 @@ def streamlit_deploy(
             raise NoProjectDefinitionError(
                 project_type="streamlit", project_root=cli_context.project_root
             )
-        pd = migrate_v1_streamlit_to_v2(pd)
+        pd = convert_project_definition_to_v2(pd)
 
     streamlits: Dict[str, StreamlitEntityModel] = pd.get_entities_by_type(
         entity_type="streamlit"
@@ -172,64 +170,6 @@ def streamlit_deploy(
     return MessageResult(f"Streamlit successfully deployed and available under {url}")
 
 
-def migrate_v1_streamlit_to_v2(pd: ProjectDefinition):
-    default_env_file = "environment.yml"
-    default_pages_dir = "pages"
-
-    # Process env file
-    environment_file = pd.streamlit.env_file
-    if environment_file and not Path(environment_file).exists():
-        raise ClickException(f"Provided file {environment_file} does not exist")
-    elif environment_file is None and Path(default_env_file).exists():
-        environment_file = default_env_file
-    # Process pages dir
-    pages_dir = pd.streamlit.pages_dir
-    if pages_dir and not Path(pages_dir).exists():
-        raise ClickException(f"Provided file {pages_dir} does not exist")
-    elif pages_dir is None and Path(default_pages_dir).exists():
-        pages_dir = default_pages_dir
-
-    # Build V2 definition
-    artifacts = [
-        pd.streamlit.main_file,
-        environment_file,
-        pages_dir,
-    ]
-    artifacts = [a for a in artifacts if a is not None]
-    if pd.streamlit.additional_source_files:
-        artifacts.extend(pd.streamlit.additional_source_files)
-
-    identifier = {"name": pd.streamlit.name}
-    if pd.streamlit.schema_name:
-        identifier["schema"] = pd.streamlit.schema_name
-    if pd.streamlit.database:
-        identifier["database"] = pd.streamlit.database
-
-    if is_name_a_templated_one(pd.streamlit.name):
-        streamlit_name = "streamlit_entity_1"  # we don't care about numbering, as only 1 streamlit is allowed in V1
-    else:
-        streamlit_name = pd.streamlit.name
-
-    data = {
-        "definition_version": "2",
-        "entities": {
-            streamlit_name: {
-                "type": "streamlit",
-                "identifier": identifier,
-                "title": pd.streamlit.title,
-                "query_warehouse": pd.streamlit.query_warehouse,
-                "main_file": str(pd.streamlit.main_file),
-                "pages_dir": str(pd.streamlit.pages_dir),
-                "stage": pd.streamlit.stage,
-                "artifacts": artifacts,
-            }
-        },
-    }
-    if hasattr(pd, "env") and pd.env:
-        data["env"] = {k: v for k, v in pd.env.items()}
-    return ProjectDefinitionV2(**data)
-
-
 @app.command("get-url", requires_connection=True)
 def get_url(
     name: FQN = StreamlitNameArgument,

diff --git a/src/snowflake/cli/_plugins/workspace/commands.py b/src/snowflake/cli/_plugins/workspace/commands.py
@@ -21,18 +21,18 @@
 
 import typer
 import yaml
-from click import ClickException
 from snowflake.cli._plugins.nativeapp.artifacts import BundleMap
 from snowflake.cli._plugins.nativeapp.common_flags import ValidateOption
-from snowflake.cli._plugins.snowpark.commands import migrate_v1_snowpark_to_v2
-from snowflake.cli._plugins.streamlit.commands import migrate_v1_streamlit_to_v2
 from snowflake.cli._plugins.workspace.manager import WorkspaceManager
 from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.decorators import with_project_definition
 from snowflake.cli.api.commands.snow_typer import SnowTyper
 from snowflake.cli.api.entities.common import EntityActions
 from snowflake.cli.api.exceptions import IncompatibleParametersError
 from snowflake.cli.api.output.types import MessageResult
+from snowflake.cli.api.project.definition_conversion import (
+    convert_project_definition_to_v2,
+)
 from snowflake.cli.api.project.definition_manager import DefinitionManager
 from snowflake.cli.api.secure_path import SecurePath
 
@@ -56,24 +56,7 @@ def migrate(
     if pd.meets_version_requirement("2"):
         return MessageResult("Project definition is already at version 2.")
 
-    if "<% ctx." in str(pd):
-        if not accept_templates:
-            raise ClickException(
-                "Project definition contains templates. They may not be migrated correctly, and require manual migration."
-                "You can try again with --accept-templates  option, to attempt automatic migration."
-            )
-        log.warning(
-            "Your V1 definition contains templates. We cannot guarantee the correctness of the migration."
-        )
-
-    if pd.streamlit:
-        pd_v2 = migrate_v1_streamlit_to_v2(pd)
-    elif pd.snowpark:
-        pd_v2 = migrate_v1_snowpark_to_v2(pd)
-    else:
-        raise ValueError(
-            "Only Snowpark and Streamlit entities are supported for migration."
-        )
+    pd_v2 = convert_project_definition_to_v2(pd, accept_templates)
 
     SecurePath("snowflake.yml").rename("snowflake_V1.yml")
     with open("snowflake.yml", "w") as file:

diff --git a/src/snowflake/cli/api/constants.py b/src/snowflake/cli/api/constants.py
@@ -77,3 +77,14 @@ def __str__(self):
 DEFAULT_SIZE_LIMIT_MB = 128
 
 SF_REST_API_URL_PREFIX = "/api/v2"
+
+PROJECT_TEMPLATE_VARIABLE_OPENING = "<%"
+PROJECT_TEMPLATE_VARIABLE_CLOSING = "%>"
+
+INIT_TEMPLATE_VARIABLE_OPENING = "<!"
+INIT_TEMPLATE_VARIABLE_CLOSING = "!>"
+
+SNOWPARK_SHARED_MIXIN = "snowpark_shared"
+
+DEFAULT_ENV_FILE = "environment.yml"
+DEFAULT_PAGES_DIR = "pages"