snowflake-cli-labs 2.3.0rc1__py3-none-any.whl → 2.4.0__py3-none-any.whl

snowflake/cli/plugins/nativeapp/codegen/compiler.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, Optional
+
+from snowflake.cli.api.console import cli_console as cc
+from snowflake.cli.api.project.schemas.native_app.native_app import NativeApp
+from snowflake.cli.api.project.schemas.native_app.path_mapping import (
+    PathMapping,
+    ProcessorMapping,
+)
+from snowflake.cli.plugins.nativeapp.artifacts import resolve_without_follow
+from snowflake.cli.plugins.nativeapp.codegen.artifact_processor import (
+    ArtifactProcessor,
+    UnsupportedArtifactProcessorError,
+)
+from snowflake.cli.plugins.nativeapp.codegen.snowpark.python_processor import (
+    SnowparkAnnotationProcessor,
+)
+
+SNOWPARK_PROCESSOR = "snowpark"
+
+
+class NativeAppCompiler:
+    """
+    Compiler class to perform custom processing on all relevant Native Apps artifacts (specified in the project definition file)
+    before an application package can be created from those artifacts.
+    An artifact can have more than one processor specified for itself, and this class will execute those processors in that order.
+    The class also maintains a dictionary of processors it creates in order to reuse them across artifacts, since processor initialization
+    is independent of the artifact to process.
+    """
+
+    def __init__(
+        self,
+        project_definition: NativeApp,
+        project_root: Path,
+        deploy_root: Path,
+        generated_root: Path,
+    ):
+        self.project_definition = project_definition
+        self.project_root = project_root
+        self.deploy_root = deploy_root
+        self.generated_root = generated_root
+
+        self.artifacts = [
+            artifact
+            for artifact in project_definition.artifacts
+            if isinstance(artifact, PathMapping)
+        ]
+        # dictionary of all processors created and shared between different artifact objects.
+        self.cached_processors: Dict[str, ArtifactProcessor] = {}
+
+    def compile_artifacts(self):
+        """
+        Go through every artifact object in the project definition of a native app, and execute processors in order of specification for each of the artifact object.
+        May have side-effects on the filesystem by either directly editing source files or the deploy root.
+        """
+        should_proceed = False
+        for artifact in self.artifacts:
+            if artifact.processors:
+                should_proceed = True
+                break
+        if not should_proceed:
+            return
+
+        with cc.phase("Invoking artifact processors"):
+            for artifact in self.artifacts:
+                for processor in artifact.processors:
+                    artifact_processor = self._try_create_processor(
+                        processor_mapping=processor,
+                    )
+                    if artifact_processor is None:
+                        raise UnsupportedArtifactProcessorError(
+                            processor_name=processor.name
+                        )
+                    else:
+                        artifact_processor.process(
+                            artifact_to_process=artifact, processor_mapping=processor
+                        )
+
+    def _try_create_processor(
+        self,
+        processor_mapping: ProcessorMapping,
+        **kwargs,
+    ) -> Optional[ArtifactProcessor]:
+        """
+        Fetch processor object if one already exists in the cached_processors dictionary.
+        Else, initialize a new object to return, and add it to the cached_processors dictionary.
+        """
+        if processor_mapping.name.lower() == SNOWPARK_PROCESSOR:
+            curr_processor = self.cached_processors.get(SNOWPARK_PROCESSOR, None)
+            if curr_processor is not None:
+                return curr_processor
+            else:
+                curr_processor = SnowparkAnnotationProcessor(
+                    project_definition=self.project_definition,
+                    project_root=resolve_without_follow(self.project_root),
+                    deploy_root=resolve_without_follow(self.deploy_root),
+                    generated_root=resolve_without_follow(self.generated_root),
+                )
+                self.cached_processors[SNOWPARK_PROCESSOR] = curr_processor
+                return curr_processor
+        else:
+            return None

snowflake/cli/plugins/nativeapp/codegen/sandbox.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import shutil
 import subprocess

snowflake/cli/plugins/nativeapp/codegen/snowpark/callback_source.py.jinja

@@ -0,0 +1,181 @@
+import contextlib
+import functools
+import inspect
+import sys
+from typing import Callable
+
+try:
+    import snowflake.snowpark
+except ModuleNotFoundError as exc:
+    print(
+        "An exception occurred while importing snowflake-snowpark-python package: ",
+        exc,
+        file=sys.stderr,
+    )
+    sys.exit(1)
+
+__snowflake_internal_found_correct_version = hasattr(
+    snowflake.snowpark.context, "_is_execution_environment_sandboxed_for_client"
+) and hasattr(snowflake.snowpark.context, "_should_continue_registration")
+
+if not __snowflake_internal_found_correct_version:
+    print(
+        "Did not find the minimum required version for snowflake-snowpark-python package. Please upgrade to v1.15.0 or higher.",
+        file=sys.stderr,
+    )
+    sys.exit(1)
+
+__snowflake_global_collected_extension_fn_json = []
+
+def __snowflake_internal_create_extension_fn_registration_callback():
+    def __snowflake_internal_try_extract_lineno(extension_fn):
+        try:
+            return inspect.getsourcelines(extension_fn)[1]
+        except Exception:
+            return None
+
+    def __snowflake_internal_extract_extension_fn_name(extension_fn):
+        try:
+            import snowflake.snowpark._internal.utils as snowpark_utils
+
+            if hasattr(snowpark_utils, 'TEMP_OBJECT_NAME_PREFIX'):
+                if extension_fn.object_name.startswith(snowpark_utils.TEMP_OBJECT_NAME_PREFIX):
+                    # The object name is a generated one, don't use it
+                    return None
+
+        except Exception:
+            # ignore any exception and fall back to using the object name reported from Snowpark
+            pass
+
+        return extension_fn.object_name
+
+    def __snowflake_internal_create_package_list(extension_fn):
+        if not extension_fn.all_packages.strip():
+            return []
+        return [pkg_spec.strip() for pkg_spec in extension_fn.all_packages.split(",")]
+
+    def __snowflake_internal_make_extension_fn_signature(extension_fn):
+        # Try to fetch the original argument names from the extension function
+        try:
+            args_spec = inspect.getfullargspec(extension_fn.func)
+            original_arg_names = args_spec[0]
+            start_index = len(original_arg_names) - len(extension_fn.input_sql_types)
+            signature = []
+            defaults_start_index = len(original_arg_names) - len(args_spec.defaults or [])
+            for i in range(len(extension_fn.input_sql_types)):
+                arg = {
+                    'name': original_arg_names[start_index + i],
+                    'type': extension_fn.input_sql_types[i]
+                }
+                if i >= defaults_start_index:
+                    arg['default'] = args_spec.defaults[defaults_start_index + i]
+                signature.append(arg)
+
+            return signature
+        except Exception as e:
+            msg = str(e)
+            pass # ignore, we'll use the fallback strategy
+
+        # Failed to extract the original arguments through reflection, fall back to alternative approach
+        return [
+            {"name": input_arg.name, "type": input_type}
+            for (input_arg, input_type) in zip(extension_fn.input_args, extension_fn.input_sql_types)
+        ]
+
+    def __snowflake_internal_to_extension_fn_type(object_type):
+        if object_type.name == "AGGREGATE_FUNCTION":
+            return "aggregate function"
+        if object_type.name == "TABLE_FUNCTION":
+            return "table function"
+        return object_type.name.lower()
+
+    def __snowflake_internal_imports_union_to_str_type(raw_imports):
+        final_imports = []
+        if raw_imports:
+            for raw_import in raw_imports:
+                if isinstance(raw_import, str):
+                    final_imports.append(raw_import)
+                else:
+                    final_imports.append(raw_import[0])
+        return final_imports
+
+
+    def __snowflake_internal_extension_fn_to_json(extension_fn):
+        if not isinstance(extension_fn.func, Callable):
+            # Unsupported case: extension function is a tuple
+            return
+
+        if extension_fn.anonymous:
+            # unsupported, native application extension functions need to be explicitly named
+            return
+
+        # Collect basic properties of the extension function
+        extension_fn_json = {
+            "type": __snowflake_internal_to_extension_fn_type(extension_fn.object_type),
+            "lineno": __snowflake_internal_try_extract_lineno(extension_fn.func),
+            "name": __snowflake_internal_extract_extension_fn_name(extension_fn),
+            "handler": extension_fn.func.__name__,
+            "imports": __snowflake_internal_imports_union_to_str_type(extension_fn.raw_imports),
+            "packages": __snowflake_internal_create_package_list(extension_fn),
+            "runtime": extension_fn.runtime_version,
+            "returns": extension_fn.return_sql.upper().replace("RETURNS ", "").strip(),
+            "signature": __snowflake_internal_make_extension_fn_signature(extension_fn),
+            "external_access_integrations": extension_fn.external_access_integrations or [],
+            "secrets": extension_fn.secrets or {},
+        }
+
+        if extension_fn.object_type.name == "PROCEDURE" and extension_fn.execute_as is not None:
+            extension_fn_json['execute_as_caller'] = (extension_fn.execute_as == 'caller')
+
+        if extension_fn.native_app_params is not None:
+            schema = extension_fn.native_app_params.get("schema")
+            if schema is not None:
+                extension_fn_json["schema"] = schema
+            app_roles = extension_fn.native_app_params.get("application_roles")
+            if app_roles is not None:
+                extension_fn_json["application_roles"] = app_roles
+
+        return extension_fn_json
+
+    def __snowflake_internal_collect_extension_fn(
+        collected_extension_fn_json_list, extension_function_properties
+    ):
+        extension_fn_json = __snowflake_internal_extension_fn_to_json(extension_function_properties)
+        collected_extension_fn_json_list.append(extension_fn_json)
+        return False
+
+    return functools.partial(
+        __snowflake_internal_collect_extension_fn,
+        __snowflake_global_collected_extension_fn_json,
+    )
+
+snowflake.snowpark.context._is_execution_environment_sandboxed_for_client = (  # noqa: SLF001
+    True
+)
+snowflake.snowpark.context._should_continue_registration = (  # noqa: SLF001
+    __snowflake_internal_create_extension_fn_registration_callback()
+)
+snowflake.snowpark.session._is_execution_environment_sandboxed_for_client = (  # noqa: SLF001
+    True
+)
+
+for global_key in list(globals().keys()):
+    if global_key.startswith("__snowflake_internal"):
+        del globals()[global_key]
+
+del globals()["global_key"] # make sure to clean up the loop variable as well
+
+try:
+    import importlib
+    with contextlib.redirect_stdout(None):
+        with contextlib.redirect_stderr(None):
+            __snowflake_internal_spec = importlib.util.spec_from_file_location("<string>", "{{py_file}}")
+            __snowflake_internal_module = importlib.util.module_from_spec(__snowflake_internal_spec)
+            __snowflake_internal_spec.loader.exec_module(__snowflake_internal_module)
+except Exception as exc:  # Catch any error
+    print("An exception occurred while executing file: ", exc, file=sys.stderr)
+    sys.exit(1)
+
+
+import json
+print(json.dumps(__snowflake_global_collected_extension_fn_json))

snowflake/cli/plugins/nativeapp/codegen/snowpark/extension_function_utils.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import ast
+from typing import (
+    Any,
+    List,
+    Optional,
+    Sequence,
+)
+
+from click.exceptions import ClickException
+from snowflake.cli.api.project.schemas.snowpark.argument import Argument
+from snowflake.cli.api.project.util import (
+    is_valid_identifier,
+    is_valid_string_literal,
+    to_identifier,
+    to_string_literal,
+)
+from snowflake.cli.plugins.nativeapp.codegen.snowpark.models import (
+    ExtensionFunctionTypeEnum,
+    NativeAppExtensionFunction,
+)
+
+
+class MalformedExtensionFunctionError(ClickException):
+    """Required extension function attribute is missing."""
+
+    def __init__(self, message: str):
+        super().__init__(message=message)
+
+
+def get_sql_object_type(extension_fn: NativeAppExtensionFunction) -> Optional[str]:
+    if extension_fn.function_type == ExtensionFunctionTypeEnum.PROCEDURE:
+        return "PROCEDURE"
+    elif extension_fn.function_type in (
+        ExtensionFunctionTypeEnum.FUNCTION,
+        ExtensionFunctionTypeEnum.TABLE_FUNCTION,
+    ):
+        return "FUNCTION"
+    elif extension_fn.function_type == extension_fn.function_type.AGGREGATE_FUNCTION:
+        return "AGGREGATE FUNCTION"
+    else:
+        return None
+
+
+def get_sql_argument_signature(arg: Argument) -> str:
+    formatted = f"{arg.name} {arg.arg_type}"
+    if arg.default is not None:
+        formatted = f"{formatted} DEFAULT {arg.default}"
+    return formatted
+
+
+def get_function_type_signature_for_grant(
+    extension_fn: NativeAppExtensionFunction,
+) -> str:
+    """
+    Returns the type signature for the specified function, e.g. "int, varchar", suitable for inclusion in a GRANT statement.
+    """
+    return ", ".join([arg.arg_type for arg in extension_fn.signature])
+
+
+def get_qualified_object_name(extension_fn: NativeAppExtensionFunction) -> str:
+    qualified_name = to_identifier(extension_fn.name)
+    if extension_fn.schema_name:
+        if is_valid_identifier(extension_fn.schema_name):
+            qualified_name = f"{extension_fn.schema_name}.{qualified_name}"
+        else:
+            full_schema = ".".join(
+                [
+                    to_identifier(schema_part)
+                    for schema_part in extension_fn.schema_name.split(".")
+                ]
+            )
+            qualified_name = f"{full_schema}.{qualified_name}"
+
+    return qualified_name
+
+
+def ensure_string_literal(value: str) -> str:
+    """
+    Returns the string literal representation of the given value, or the value itself if
+    it was already a valid string literal.
+    """
+    if is_valid_string_literal(value):
+        return value
+    return to_string_literal(value)
+
+
+def ensure_all_string_literals(values: Sequence[str]) -> List[str]:
+    """
+    Ensures that all provided values are valid string literals.
+
+    Returns:
+        A list with all values transformed to be valid string literals (as necessary).
+    """
+    return [ensure_string_literal(value) for value in values]
+
+
+class _FunctionDefAccumulator(ast.NodeVisitor):
+    """
+    A NodeVisitor that collects AST nodes corresponding to function declarations, filtered by a list
+    of wanted functions. This is used to identify all Snowpark extension functions in a module's
+    source code.
+    """
+
+    def __init__(self, functions: Sequence[NativeAppExtensionFunction]):
+        self._wanted_functions_by_name = {
+            fn.handler.split(".")[-1]: fn for fn in functions
+        }
+        self.definitions: List[Any] = []
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):  # noqa: N802
+        if self._want(node):
+            self.definitions.append(node)
+        self.generic_visit(node)
+
+    def _want(self, node: Any) -> bool:
+        if not node.decorator_list:
+            # No decorators for this definition, ignore it
+            return False
+
+        return node.name in self._wanted_functions_by_name
+
+
+def _get_decorator_id(node: ast.AST) -> Optional[str]:
+    """
+    Returns the fully qualified identifier for a decorator, e.g. "foo" or "foo.bar".
+    """
+    if isinstance(node, ast.Name):
+        return node.id
+    elif isinstance(node, ast.Attribute):
+        return f"{_get_decorator_id(node.value)}.{node.attr}"
+    elif isinstance(node, ast.Call):
+        return _get_decorator_id(node.func)
+    else:
+        return None
+
+
+def _collect_ast_function_definitions(
+    tree: ast.AST, extension_functions: Sequence[NativeAppExtensionFunction]
+) -> Sequence[ast.FunctionDef]:
+    accumulator = _FunctionDefAccumulator(extension_functions)
+    accumulator.visit(tree)
+    return accumulator.definitions
+
+
+def deannotate_module_source(
+    module_source: str,
+    extension_functions: Sequence[NativeAppExtensionFunction],
+    annotations_to_preserve: Sequence[str] = (),
+) -> str:
+    """
+    Removes annotations from a set of specified extension functions.
+
+    Arguments:
+        module_source (str): The source code of the module to deannotate.
+        extension_functions (Sequence[NativeAppExtensionFunction]): The list of extension functions
+         to deannotate. Other functions encountered will be ignored.
+        annotations_to_preserve (Sequence[str], optional): The list of annotations to preserve. The
+         names should appear as they are found in the source code, e.g. "foo" for @foo or
+         "annotations.bar" for @annotations.bar.
+
+    Returns:
+        A de-annotated version of the module source if any match was found. In order to preserve
+        line numbers, annotations are simply commented out instead of completely removed.
+    """
+
+    tree = ast.parse(module_source)
+
+    definitions = _collect_ast_function_definitions(tree, extension_functions)
+    if not definitions:
+        return module_source
+
+    module_lines = module_source.splitlines()
+    for definition in definitions:
+        # Comment out all decorators. As per the python grammar, decorators must be terminated by a
+        # new line, so the line ranges can't overlap.
+        for decorator in definition.decorator_list:
+            decorator_id = _get_decorator_id(decorator)
+            if decorator_id is None:
+                continue
+            if annotations_to_preserve and decorator_id in annotations_to_preserve:
+                continue
+
+            # AST indices are 1-based
+            start_lineno = decorator.lineno - 1
+            if decorator.end_lineno is not None:
+                end_lineno = decorator.end_lineno - 1
+            else:
+                end_lineno = start_lineno
+
+            for lineno in range(start_lineno, end_lineno + 1):
+                module_lines[lineno] = "#: " + module_lines[lineno]
+
+    # we're writing files in text mode, so we should use '\n' regardless of the platform
+    return "\n".join(module_lines)

snowflake/cli/plugins/nativeapp/codegen/snowpark/models.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import List, Optional
+
+from pydantic import Field
+from snowflake.cli.api.project.schemas.snowpark.callable import _CallableBase
+from snowflake.cli.api.project.schemas.updatable_model import IdentifierField
+
+
+class ExtensionFunctionTypeEnum(str, Enum):
+    PROCEDURE = "procedure"
+    FUNCTION = "function"
+    TABLE_FUNCTION = "table function"
+    AGGREGATE_FUNCTION = "aggregate function"
+
+
+class NativeAppExtensionFunction(_CallableBase):
+    function_type: ExtensionFunctionTypeEnum = Field(
+        title="The type of extension function, one of 'procedure', 'function', 'table function' or 'aggregate function'",
+        alias="type",
+    )
+    lineno: Optional[int] = Field(
+        title="The starting line number of the extension function (1-based)",
+        default=None,
+    )
+    name: Optional[str] = Field(
+        title="The name of the extension function", default=None
+    )
+    packages: Optional[List[str]] = Field(
+        title="List of packages (with optional version constraints) to be loaded for the function",
+        default=[],
+    )
+    schema_name: Optional[str] = IdentifierField(
+        title=f"Name of the schema for the function",
+        default=None,
+        alias="schema",
+    )
+    application_roles: Optional[List[str]] = Field(
+        title="Application roles granted usage to the function",
+        default=[],
+    )
+    execute_as_caller: Optional[bool] = Field(
+        title="Determine whether the procedure is executed with the privileges of "
+        "the owner or with the privileges of the caller",
+        default=False,
+    )