diff --git a/src/snowflake/cli/_plugins/snowpark/snowpark_entity.py b/src/snowflake/cli/_plugins/snowpark/snowpark_entity.py index 55d1113748..81cb1bf2dc 100644 --- a/src/snowflake/cli/_plugins/snowpark/snowpark_entity.py +++ b/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]): diff --git a/tests/snowpark/test_snowpark_entities.py b/tests/snowpark/test_snowpark_entities.py new file mode 100644 index 0000000000..858d9eabac --- /dev/null +++ b/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;" + )