Mock snowpark entities

src/snowflake/cli/_plugins/snowpark/snowpark_entity.py

@@ -1,16 +1,100 @@
-from typing import Generic, TypeVar
+import os
+from pathlib import Path
+from typing import Generic, Optional, TypeVar
 
+from snowflake.cli._plugins.nativeapp.artifacts import build_bundle
+from snowflake.cli._plugins.nativeapp.entities.application_package_child_interface import (
+    ApplicationPackageChildInterface,
+)
+from snowflake.cli._plugins.nativeapp.feature_flags import FeatureFlag
 from snowflake.cli._plugins.snowpark.snowpark_entity_model import (
     FunctionEntityModel,
     ProcedureEntityModel,
 )
 from snowflake.cli.api.entities.common import EntityBase
+from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import PathMapping
 
 T = TypeVar("T")
 
 
-class SnowparkEntity(EntityBase[Generic[T]]):
-    pass
+# WARNING: The Function/Procedure entities are not implemented yet. The logic below is only for demonstrating the
+# required interfaces for composability (used by ApplicationPackageEntity behind a feature flag).
+class SnowparkEntity(EntityBase[Generic[T]], ApplicationPackageChildInterface):
+    def __init__(self, *args, **kwargs):
+        if not FeatureFlag.ENABLE_NATIVE_APP_CHILDREN.is_enabled():
+            raise NotImplementedError("Snowpark entities are not implemented yet")
+        super().__init__(*args, **kwargs)
+
+    @property
+    def project_root(self) -> Path:
+        return self._workspace_ctx.project_root
+
+    @property
+    def deploy_root(self) -> Path:
+        return self.project_root / "output" / "deploy"
+
+    def action_bundle(
+        self,
+        *args,
+        **kwargs,
+    ):
+        return self.bundle()
+
+    def bundle(self, bundle_root=None):
+        return build_bundle(
+            self.project_root,
+            bundle_root or self.deploy_root,
+            [
+                PathMapping(src=str(artifact.src), dest=artifact.dest)
+                for artifact in self._entity_model.artifacts
+            ],
+        )
+
+    def _get_identifier_for_sql(
+        self, arg_names: bool = True, schema: Optional[str] = None
+    ) -> str:
+        model = self._entity_model
+        if arg_names:
+            signature = ", ".join(
+                f"{arg.name} {arg.arg_type}" for arg in model.signature
+            )
+        else:
+            signature = ", ".join(arg.arg_type for arg in model.signature)
+        entity_id = self.entity_id
+        object_name = f"{schema}.{entity_id}" if schema else entity_id
+        return f"{object_name}({signature})"
+
+    def get_deploy_sql(
+        self,
+        artifacts_dir: Optional[Path] = None,
+        schema: Optional[str] = None,
+    ):
+        model = self._entity_model
+        imports = [f"'{x}'" for x in model.imports]
+        if artifacts_dir:
+            for root, _, files in os.walk(self.deploy_root / artifacts_dir):
+                for f in files:
+                    file_path_relative_to_deploy_root = (
+                        Path(root).relative_to(self.deploy_root) / f
+                    )
+                    imports.append(f"'/{str(file_path_relative_to_deploy_root)}'")
+
+        entity_type = model.get_type().upper()
+
+        query = [
+            f"CREATE OR REPLACE {entity_type} {self._get_identifier_for_sql(schema=schema)}",
+            f"RETURNS {model.returns}",
+            "LANGUAGE python",
+            "RUNTIME_VERSION=3.8",
+            f"IMPORTS=({', '.join(imports)})",
+            f"HANDLER='{model.handler}'",
+            "PACKAGES=('snowflake-snowpark-python');",
+        ]
+        return "\n".join(query)
+
+    def get_usage_grant_sql(self, app_role: str, schema: Optional[str] = None):
+        entity_type = self._entity_model.get_type().upper()
+        return f"GRANT USAGE ON {entity_type} {self._get_identifier_for_sql(schema=schema, arg_names=False)} TO APPLICATION ROLE {app_role};"
 
 
 class FunctionEntity(SnowparkEntity[FunctionEntityModel]):

tests/snowpark/test_snowpark_entities.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from pathlib import Path
+from textwrap import dedent
+
+import pytest
+from snowflake.cli._plugins.snowpark.snowpark_entity import (
+    FunctionEntity,
+    ProcedureEntity,
+)
+from snowflake.cli._plugins.snowpark.snowpark_entity_model import (
+    FunctionEntityModel,
+    ProcedureEntityModel,
+)
+from snowflake.cli._plugins.workspace.context import WorkspaceContext
+from snowflake.cli.api.console import cli_console as cc
+from snowflake.cli.api.project.definition_manager import DefinitionManager
+
+from tests.testing_utils.mock_config import mock_config_key
+
+
+def test_cannot_instantiate_without_feature_flag():
+    with pytest.raises(NotImplementedError) as err:
+        FunctionEntity()
+    assert str(err.value) == "Snowpark entities are not implemented yet"
+
+    with pytest.raises(NotImplementedError) as err:
+        ProcedureEntity()
+    assert str(err.value) == "Snowpark entities are not implemented yet"
+
+
+def test_function_implements_nativeapp_children_interface(temp_dir):
+    with mock_config_key("enable_native_app_children", True):
+        dm = DefinitionManager()
+        ctx = WorkspaceContext(
+            console=cc,
+            project_root=dm.project_root,
+            get_default_role=lambda: "mock_role",
+            get_default_warehouse=lambda: "mock_warehouse",
+        )
+        main_file = "main.py"
+        (Path(temp_dir) / main_file).touch()
+        model = FunctionEntityModel(
+            type="function",
+            handler="my_schema.my_func",
+            returns="integer",
+            signature=[
+                {"name": "input_number", "type": "integer"},
+                {"name": "input_string", "type": "text"},
+            ],
+            stage="my_stage",
+            artifacts=[main_file],
+        )
+        model._entity_id = "my_func"  # noqa: SLF001
+        schema = "my_schema"
+        fn = FunctionEntity(model, ctx)
+
+        fn.bundle()
+        bundle_artifact = Path(temp_dir) / "output" / "deploy" / main_file
+        deploy_sql_str = fn.get_deploy_sql(schema=schema)
+        grant_sql_str = fn.get_usage_grant_sql(app_role="app_role", schema=schema)
+
+        assert bundle_artifact.exists()
+        assert (
+            deploy_sql_str
+            == dedent(
+                """
+            CREATE OR REPLACE FUNCTION my_schema.my_func(input_number integer, input_string text)
+            RETURNS integer
+            LANGUAGE python
+            RUNTIME_VERSION=3.8
+            IMPORTS=()
+            HANDLER='my_schema.my_func'
+            PACKAGES=('snowflake-snowpark-python');
+            """
+            ).strip()
+        )
+        assert (
+            grant_sql_str
+            == "GRANT USAGE ON FUNCTION my_schema.my_func(integer, text) TO APPLICATION ROLE app_role;"
+        )
+
+
+def test_procedure_implements_nativeapp_children_interface(temp_dir):
+    with mock_config_key("enable_native_app_children", True):
+        dm = DefinitionManager()
+        ctx = WorkspaceContext(
+            console=cc,
+            project_root=dm.project_root,
+            get_default_role=lambda: "mock_role",
+            get_default_warehouse=lambda: "mock_warehouse",
+        )
+        main_file = "main.py"
+        (Path(temp_dir) / main_file).touch()
+        model = ProcedureEntityModel(
+            type="procedure",
+            handler="my_schema.my_sproc",
+            returns="integer",
+            signature=[
+                {"name": "input_number", "type": "integer"},
+                {"name": "input_string", "type": "text"},
+            ],
+            stage="my_stage",
+            artifacts=[main_file],
+        )
+        model._entity_id = "my_sproc"  # noqa: SLF001
+        schema = "my_schema"
+        fn = ProcedureEntity(model, ctx)
+
+        fn.bundle()
+        bundle_artifact = Path(temp_dir) / "output" / "deploy" / main_file
+        deploy_sql_str = fn.get_deploy_sql(schema=schema)
+        grant_sql_str = fn.get_usage_grant_sql(app_role="app_role", schema=schema)
+
+        assert bundle_artifact.exists()
+        assert (
+            deploy_sql_str
+            == dedent(
+                """
+            CREATE OR REPLACE PROCEDURE my_schema.my_sproc(input_number integer, input_string text)
+            RETURNS integer
+            LANGUAGE python
+            RUNTIME_VERSION=3.8
+            IMPORTS=()
+            HANDLER='my_schema.my_sproc'
+            PACKAGES=('snowflake-snowpark-python');
+            """
+            ).strip()
+        )
+        assert (
+            grant_sql_str
+            == "GRANT USAGE ON PROCEDURE my_schema.my_sproc(integer, text) TO APPLICATION ROLE app_role;"
+        )