snowflake-cli-labs 2.4.1__py3-none-any.whl → 2.5.0rc1__py3-none-any.whl

snowflake/cli/plugins/nativeapp/artifacts.py

@@ -1,11 +1,24 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from __future__ import annotations
 
 import itertools
 import os
-from dataclasses import dataclass
 from pathlib import Path
 from textwrap import dedent
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, Union
 
 from click.exceptions import ClickException
 from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB
@@ -43,7 +56,7 @@ class SourceNotFoundError(ClickException):
 
 class TooManyFilesError(ClickException):
     """
-    Multiple files were mapped to one output file.
+    Multiple file or directories were mapped to one output destination.
     """
 
     dest_path: Path
@@ -85,17 +98,130 @@ class NotInDeployRootError(ClickException):
         self.src_path = src_path
 
 
-@dataclass
-class ArtifactMapping:
+ArtifactPredicate = Callable[[Path, Path], bool]
+
+
+class _ArtifactPathMap:
     """
-    Used to keep track of equivalent paths / globs so we can copy artifacts from the project folder to the deploy root.
+    A specialized version of an ordered multimap used to keep track of artifact
+    source-destination mappings. The mapping is bidirectional, so it can be queried
+    by source or destination paths. All paths manipulated by this class must be in
+    relative, canonical form (relative to the project or deploy roots, as appropriate).
     """
 
-    src: str
-    dest: str
+    def __init__(self, project_root: Path):
+        self._project_root = project_root
 
+        # All (src,dest) pairs in inserting order, for iterating
+        self.__src_dest_pairs: List[Tuple[Path, Path]] = []
+        # built-in dict instances are ordered as of Python 3.7
+        self.__src_to_dest: Dict[Path, List[Path]] = {}
+        self.__dest_to_src: Dict[Path, List[Path]] = {}
 
-ArtifactPredicate = Callable[[Path, Path], bool]
+        # This dictionary accumulates keys for each directory or file to be created in
+        # the deploy root for any artifact mapping rule being processed. This includes
+        # children of directories that are copied to the deploy root. Having this
+        # information available is critical to detect possible clashes between rules.
+        self._dest_is_dir: Dict[Path, bool] = {}
+
+    def put(self, src: Path, dest: Path, dest_is_dir: bool) -> None:
+        """
+        Adds a new source-destination mapping pair to this map, if necessary. Note that
+        this is internal logic that assumes that src-dest pairs have already been preprocessed
+        by the enclosing BundleMap (for example, only file -> file and
+        directory -> directory mappings are possible here due to the preprocessing step).
+
+        Arguments:
+            src {Path} -- the source path, in canonical form.
+            dest {Path} -- the destination path, in canonical form.
+            dest_is_dir {bool} -- whether the destination path is a directory.
+        """
+        # Both paths should be in canonical form
+        assert not src.is_absolute()
+        assert not dest.is_absolute()
+
+        absolute_src = self._project_root / src
+
+        current_sources = self.__dest_to_src.get(dest, [])
+        src_is_dir = absolute_src.is_dir()
+        if dest_is_dir:
+            assert src_is_dir  # file -> directory is not possible here given how rules are processed
+
+            # directory -> directory
+            # Check that dest is currently unmapped
+            current_is_dir = self._dest_is_dir.get(dest, False)
+            if current_is_dir:
+                # mapping to an existing directory is not allowed
+                raise TooManyFilesError(dest)
+        else:
+            # file -> file
+            # Check that there is no previous mapping for the same file.
+            if current_sources and src not in current_sources:
+                # There is already a different source mapping to this destination
+                raise TooManyFilesError(dest)
+
+        if src_is_dir:
+            # mark all subdirectories of this source as directories so that we can
+            # detect accidental clobbering
+            for (root, _, files) in os.walk(absolute_src, followlinks=True):
+                canonical_subdir = Path(root).relative_to(absolute_src)
+                canonical_dest_subdir = dest / canonical_subdir
+                self._update_dest_is_dir(canonical_dest_subdir, is_dir=True)
+                for f in files:
+                    self._update_dest_is_dir(canonical_dest_subdir / f, is_dir=False)
+
+        # make sure we check for dest_is_dir consistency regardless of whether the
+        # insertion happened. This update can fail, so we need to do it first to
+        # avoid applying partial updates to the underlying data storage.
+        self._update_dest_is_dir(dest, dest_is_dir)
+
+        dests = self.__src_to_dest.setdefault(src, [])
+        srcs = self.__dest_to_src.setdefault(dest, [])
+        if dest not in dests:
+            dests.append(dest)
+            srcs.append(src)
+            self.__src_dest_pairs.append((src, dest))
+
+    def get_sources(self, dest: Path) -> Iterable[Path]:
+        """
+        Returns all source paths associated with the provided destination path, in insertion order.
+        """
+        return self.__dest_to_src.get(dest, [])
+
+    def get_destinations(self, src: Path) -> Iterable[Path]:
+        """
+        Returns all destination paths associated with the provided source path, in insertion order.
+        """
+        return self.__src_to_dest.get(src, [])
+
+    def __iter__(self) -> Iterator[Tuple[Path, Path]]:
+        """
+        Returns all (source, destination) pairs known to this map, in insertion order.
+        """
+        return iter(self.__src_dest_pairs)
+
+    def _update_dest_is_dir(self, dest: Path, is_dir: bool) -> None:
+        """
+        Recursively marks seen destination paths as either files or folders, raising an error if any inconsistencies
+        from previous invocations of this method are encountered.
+
+        Arguments:
+            dest {Path} -- the destination path, in canonical form.
+            is_dir {bool} -- whether the destination path is a directory.
+        """
+        assert not dest.is_absolute()  # dest must be in canonical relative form
+
+        current_is_dir = self._dest_is_dir.get(dest, None)
+        if current_is_dir is not None and current_is_dir != is_dir:
+            raise ArtifactError(
+                "Conflicting type for destination path: {canonical_dest}"
+            )
+
+        parent = dest.parent
+        if parent != dest:
+            self._update_dest_is_dir(parent, True)
+
+        self._dest_is_dir[dest] = is_dir
 
 
 class BundleMap:
@@ -107,9 +233,7 @@ class BundleMap:
     def __init__(self, *, project_root: Path, deploy_root: Path):
         self._project_root: Path = resolve_without_follow(project_root)
         self._deploy_root: Path = resolve_without_follow(deploy_root)
-        self._src_to_dest: Dict[Path, List[Path]] = {}
-        self._dest_to_src: Dict[Path, List[Path]] = {}
-        self._dest_is_dir: Dict[Path, bool] = {}
+        self._artifact_map = _ArtifactPathMap(project_root=self._project_root)
 
     def deploy_root(self) -> Path:
         return self._deploy_root
@@ -137,7 +261,7 @@ class BundleMap:
             )
 
         if self._deploy_root in absolute_src.parents:
-            # ignore this item since it's in the deploy root. This can happen if the bundle map is create
+            # ignore this item since it's in the deploy root. This can happen if the bundle map is created
             # after the bundle step and a project is using rules that are not sufficiently constrained.
             # Since the bundle step starts with deleting the deploy root, we wouldn't normally encounter this situation.
             return
@@ -145,28 +269,16 @@ class BundleMap:
         canonical_src = self._canonical_src(src)
         canonical_dest = self._canonical_dest(dest)
 
-        src_is_dir = absolute_src.is_dir()
         if map_as_child:
             # Make sure the destination is a child of the original, since this was requested
             canonical_dest = canonical_dest / canonical_src.name
-            dest_is_dir = src.is_dir()
-
-        # Verify that multiple files are not being mapped to a single file destination
-        current_sources = self._dest_to_src.setdefault(canonical_dest, [])
-        if not dest_is_dir:
-            # the destination is a file
-            if (canonical_src not in current_sources) and len(current_sources) > 0:
-                raise TooManyFilesError(dest)
+            dest_is_dir = absolute_src.is_dir()
 
-        # Perform all updates together we don't end up with inconsistent state
-        self._update_dest_is_dir(canonical_dest, dest_is_dir)
-        current_dests = self._src_to_dest.setdefault(canonical_src, [])
-        if canonical_dest not in current_dests:
-            current_dests.append(canonical_dest)
-        if canonical_src not in current_sources:
-            current_sources.append(canonical_src)
+        self._artifact_map.put(
+            src=canonical_src, dest=canonical_dest, dest_is_dir=dest_is_dir
+        )
 
-    def _add_mapping(self, src: str, dest: Optional[str]):
+    def _add_mapping(self, src: str, dest: Optional[str] = None):
         """
         Adds the specified artifact rule to this instance. The source should be relative to the project directory. It
         is interpreted as a file, directory or glob pattern. If the destination path is not specified, each source match
@@ -206,35 +318,42 @@ class BundleMap:
         if not match_found:
             raise SourceNotFoundError(src)
 
-    def add(self, mapping: Union[ArtifactMapping, PathMapping]) -> None:
+    def add(self, mapping: PathMapping) -> None:
         """
         Adds an artifact mapping rule to this instance.
         """
-        if isinstance(mapping, ArtifactMapping):
-            self._add_mapping(mapping.src, mapping.dest)
-        elif isinstance(mapping, PathMapping):
-            self._add_mapping(mapping.src, mapping.dest)
-        else:
-            raise RuntimeError(f"Unsupported mapping type: {type(mapping)}")
+        self._add_mapping(mapping.src, mapping.dest)
 
-    def _mappings_for_source(
+    def _expand_artifact_mapping(
         self,
         src: Path,
+        dest: Path,
         absolute: bool = False,
         expand_directories: bool = False,
         predicate: ArtifactPredicate = lambda src, dest: True,
     ) -> Iterator[Tuple[Path, Path]]:
+        """
+        Expands the specified source-destination mapping according to the provided options.
+        The original mapping is yielded, followed by any expanded mappings derived from
+        it.
+
+        Arguments:
+            src {Path} -- the source path
+            dest {Path} -- the destination path
+            absolute {bool} -- when True, all mappings will be yielded as absolute paths
+            expand_directories {bool} -- when True, child mappings are yielded if the source path is a directory.
+            predicate {ArtifactPredicate} -- when specified, only mappings satisfying this predicate will be yielded.
+        """
         canonical_src = self._canonical_src(src)
-        canonical_dests = self._src_to_dest.get(canonical_src)
-        assert canonical_dests is not None
+        canonical_dest = self._canonical_dest(dest)
 
         absolute_src = self._absolute_src(canonical_src)
+        absolute_dest = self._absolute_dest(canonical_dest)
         src_for_output = self._to_output_src(absolute_src, absolute)
-        dests_for_output = [self._to_output_dest(p, absolute) for p in canonical_dests]
+        dest_for_output = self._to_output_dest(absolute_dest, absolute)
 
-        for d in dests_for_output:
-            if predicate(src_for_output, d):
-                yield src_for_output, d
+        if predicate(src_for_output, dest_for_output):
+            yield src_for_output, dest_for_output
 
         if absolute_src.is_dir() and expand_directories:
             # both src and dest are directories, and expanding directories was requested. Traverse src, and map each
@@ -242,11 +361,10 @@ class BundleMap:
             for (root, subdirs, files) in os.walk(absolute_src, followlinks=True):
                 relative_root = Path(root).relative_to(absolute_src)
                 for name in itertools.chain(subdirs, files):
-                    for d in dests_for_output:
-                        src_file_for_output = src_for_output / relative_root / name
-                        dest_file_for_output = d / relative_root / name
-                        if predicate(src_file_for_output, dest_file_for_output):
-                            yield src_file_for_output, dest_file_for_output
+                    src_file_for_output = src_for_output / relative_root / name
+                    dest_file_for_output = dest_for_output / relative_root / name
+                    if predicate(src_file_for_output, dest_file_for_output):
+                        yield src_file_for_output, dest_file_for_output
 
     def all_mappings(
         self,
@@ -270,9 +388,10 @@ class BundleMap:
         Returns:
           An iterator over all matching deployed artifacts.
         """
-        for src in self._src_to_dest.keys():
-            for deployed_src, deployed_dest in self._mappings_for_source(
+        for src, dest in self._artifact_map:
+            for deployed_src, deployed_dest in self._expand_artifact_mapping(
                 src,
+                dest,
                 absolute=absolute,
                 expand_directories=expand_directories,
                 predicate=predicate,
@@ -300,8 +419,8 @@ class BundleMap:
 
         output_destinations: List[Path] = []
 
-        canonical_dests = self._src_to_dest.get(canonical_src)
-        if canonical_dests is not None:
+        canonical_dests = self._artifact_map.get_destinations(canonical_src)
+        if canonical_dests:
             for d in canonical_dests:
                 output_destinations.append(self._to_output_dest(d, is_absolute))
 
@@ -362,19 +481,6 @@ class BundleMap:
     def _to_output_src(self, src: Path, absolute: bool) -> Path:
         return self._absolute_src(src) if absolute else self._canonical_src(src)
 
-    def _update_dest_is_dir(self, canonical_dest: Path, is_dir: bool) -> None:
-        current_is_dir = self._dest_is_dir.get(canonical_dest, None)
-        if current_is_dir is not None and is_dir != current_is_dir:
-            raise ArtifactError(
-                "Conflicting type for destination path: {canonical_dest}"
-            )
-
-        parent = canonical_dest.parent
-        if parent != canonical_dest:
-            self._update_dest_is_dir(parent, True)
-
-        self._dest_is_dir[canonical_dest] = is_dir
-
 
 def specifies_directory(s: str) -> bool:
     """
@@ -402,49 +508,47 @@ def delete(path: Path) -> None:
         spath.rmdir(recursive=True)  # remove dir and all contains
 
 
-def symlink_or_copy(
-    src: Path, dst: Path, deploy_root: Path, makedirs=True, overwrite=True
-) -> None:
+def symlink_or_copy(src: Path, dst: Path, deploy_root: Path) -> None:
     """
-    Tries to create a symlink to src at dst; failing that (i.e. in Windows
-    without Administrator / Developer Mode) copies the file from src to dst instead.
-    If makedirs is True, the directory hierarchy above dst is created if any
-    of those directories do not exist.
+    Symlinks files from src to dst. If the src contains parent directories, then copies the empty directory shell to the deploy root.
+    The directory hierarchy above dst is created if any of those directories do not exist.
     """
     ssrc = SecurePath(src)
     sdst = SecurePath(dst)
-    if makedirs:
-        sdst.parent.mkdir(parents=True, exist_ok=True)
+    sdst.parent.mkdir(parents=True, exist_ok=True)
 
     # Verify that the mapping isn't accidentally trying to create a file in the project source through symlinks.
     # We need to ensure we're resolving symlinks for this check to be effective.
+    # We are unlikely to hit this if calling the function through bundle map, keeping it here for other future use cases outside bundle.
     resolved_dst = dst.resolve()
     resolved_deploy_root = deploy_root.resolve()
-    if resolved_deploy_root not in resolved_dst.parents:
+    dst_is_deploy_root = resolved_deploy_root == resolved_dst
+    if (not dst_is_deploy_root) and (resolved_deploy_root not in resolved_dst.parents):
         raise NotInDeployRootError(dest_path=dst, deploy_root=deploy_root, src_path=src)
 
-    if overwrite:
+    absolute_src = resolve_without_follow(src)
+    if absolute_src.is_file():
         delete(dst)
-    try:
-        os.symlink(src, dst)
-    except OSError:
-        ssrc.copy(dst)
-
-
-def translate_artifact(item: Union[PathMapping, str]) -> ArtifactMapping:
-    """
-    Builds an artifact mapping from a project definition value.
-    Validation is done later when we actually resolve files / folders.
-    """
-
-    if isinstance(item, PathMapping):
-        return ArtifactMapping(item.src, item.dest if item.dest else item.src)
-
-    elif isinstance(item, str):
-        return ArtifactMapping(item, item)
-
-    # XXX: validation should have caught this
-    raise ArtifactError("Item is not a valid artifact!")
+        try:
+            os.symlink(absolute_src, dst)
+        except OSError:
+            ssrc.copy(dst)
+    else:
+        # 1. Create a new directory in the deploy root
+        dst.mkdir(exist_ok=True)
+        # 2. For all children of src, create their counterparts in dst now that it exists
+        for root, _, files in sorted(os.walk(absolute_src, followlinks=True)):
+            relative_root = Path(root).relative_to(absolute_src)
+            absolute_root_in_deploy = Path(dst, relative_root)
+            absolute_root_in_deploy.mkdir(parents=True, exist_ok=True)
+            for file in files:
+                absolute_file_in_project = Path(absolute_src, relative_root, file)
+                absolute_file_in_deploy = Path(absolute_root_in_deploy, file)
+                symlink_or_copy(
+                    src=absolute_file_in_project,
+                    dst=absolute_file_in_deploy,
+                    deploy_root=deploy_root,
+                )
 
 
 def resolve_without_follow(path: Path) -> Path:
@@ -458,7 +562,7 @@ def resolve_without_follow(path: Path) -> Path:
 def build_bundle(
     project_root: Path,
     deploy_root: Path,
-    artifacts: List[ArtifactMapping],
+    artifacts: List[PathMapping],
 ) -> BundleMap:
     """
     Prepares a local folder (deploy_root) with configured app artifacts.

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

@@ -0,0 +1,13 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

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

@@ -1,3 +1,17 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from __future__ import annotations
 
 from abc import ABC, abstractmethod

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

@@ -1,3 +1,17 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from __future__ import annotations
 
 from pathlib import Path

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

@@ -1,3 +1,17 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from __future__ import annotations
 
 import os

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

@@ -2,7 +2,7 @@ import contextlib
 import functools
 import inspect
 import sys
-from typing import Callable
+from typing import Callable, Tuple
 
 try:
     import snowflake.snowpark
@@ -101,8 +101,7 @@ def __snowflake_internal_create_extension_fn_registration_callback():
 
 
     def __snowflake_internal_extension_fn_to_json(extension_fn):
-        if not isinstance(extension_fn.func, Callable):
-            # Unsupported case: extension function is a tuple
+        if not (isinstance(extension_fn.func, Callable) or isinstance(extension_fn.func, Tuple)):
             return
 
         if extension_fn.anonymous:
@@ -141,7 +140,8 @@ def __snowflake_internal_create_extension_fn_registration_callback():
         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)
+        if extension_fn_json: # Do not append if extension_fn_json is None
+            collected_extension_fn_json_list.append(extension_fn_json)
         return False
 
     return functools.partial(
@@ -169,7 +169,7 @@ 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_spec = importlib.util.spec_from_file_location("<string>", r"{{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

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

@@ -1,11 +1,25 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from __future__ import annotations
 
 import ast
 from typing import (
-    Any,
     List,
     Optional,
     Sequence,
+    Union,
 )
 
 from click.exceptions import ClickException
@@ -21,6 +35,8 @@ from snowflake.cli.plugins.nativeapp.codegen.snowpark.models import (
     NativeAppExtensionFunction,
 )
 
+ASTDefNode = Union[ast.FunctionDef, ast.ClassDef]
+
 
 class MalformedExtensionFunctionError(ClickException):
     """Required extension function attribute is missing."""
@@ -96,30 +112,35 @@ def ensure_all_string_literals(values: Sequence[str]) -> List[str]:
     return [ensure_string_literal(value) for value in values]
 
 
-class _FunctionDefAccumulator(ast.NodeVisitor):
+class _SnowparkHandlerAccumulator(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.
+    A NodeVisitor that collects AST nodes corresponding to a provided list of Snowpark external functions.
+    The returned nodes are filtered using the handlers provided for each of the Snowpark functions.
+    Returned definitions can be either function definition or class definition AST nodes.
     """
 
     def __init__(self, functions: Sequence[NativeAppExtensionFunction]):
-        self._wanted_functions_by_name = {
+        self._wanted_handlers_by_name = {
             fn.handler.split(".")[-1]: fn for fn in functions
         }
-        self.definitions: List[Any] = []
+        self.definitions: List[ASTDefNode] = []
 
     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:
+    def visit_ClassDef(self, node: ast.ClassDef):  # noqa: N802
+        if self._want(node):
+            self.definitions.append(node)
+        self.generic_visit(node)
+
+    def _want(self, node: ASTDefNode) -> bool:
         if not node.decorator_list:
             # No decorators for this definition, ignore it
             return False
 
-        return node.name in self._wanted_functions_by_name
+        return node.name in self._wanted_handlers_by_name
 
 
 def _get_decorator_id(node: ast.AST) -> Optional[str]:
@@ -136,10 +157,10 @@ def _get_decorator_id(node: ast.AST) -> Optional[str]:
         return None
 
 
-def _collect_ast_function_definitions(
+def _collect_ast_handler_definitions(
     tree: ast.AST, extension_functions: Sequence[NativeAppExtensionFunction]
-) -> Sequence[ast.FunctionDef]:
-    accumulator = _FunctionDefAccumulator(extension_functions)
+) -> Sequence[ASTDefNode]:
+    accumulator = _SnowparkHandlerAccumulator(extension_functions)
     accumulator.visit(tree)
     return accumulator.definitions
 
@@ -167,7 +188,7 @@ def deannotate_module_source(
 
     tree = ast.parse(module_source)
 
-    definitions = _collect_ast_function_definitions(tree, extension_functions)
+    definitions = _collect_ast_handler_definitions(tree, extension_functions)
     if not definitions:
         return module_source
 

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

@@ -1,3 +1,17 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from __future__ import annotations
 
 from enum import Enum