snowflake-cli-labs 2.8.0rc1__py3-none-any.whl → 3.0.0rc1__py3-none-any.whl

snowflake/cli/api/project/schemas/entities/snowpark_entity.py

@@ -0,0 +1,176 @@
+# 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
+from typing import List, Literal, Optional, Union
+
+from pydantic import Field, field_validator
+from snowflake.cli.api.identifiers import FQN
+from snowflake.cli.api.project.schemas.entities.common import (
+    EntityModelBase,
+    ExternalAccessBaseModel,
+)
+from snowflake.cli.api.project.schemas.snowpark.argument import Argument
+from snowflake.cli.api.project.schemas.updatable_model import (
+    DiscriminatorField,
+    UpdatableModel,
+)
+
+
+class PathMapping(UpdatableModel):
+    class Config:
+        frozen = True
+
+    src: Path = Field(title="Source path (relative to project root)", default=None)
+
+    dest: Optional[str] = Field(
+        title="Destination path on stage",
+        description="Paths are relative to stage root; paths ending with a slash indicate that the destination is a directory which source files should be copied into.",
+        default=None,
+    )
+
+
+class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel):
+    handler: str = Field(
+        title="Function’s or procedure’s implementation of the object inside source module",
+        examples=["functions.hello_function"],
+    )
+    returns: str = Field(
+        title="Type of the result"
+    )  # TODO: again, consider Literal/Enum
+    signature: Union[str, List[Argument]] = Field(
+        title="The signature parameter describes consecutive arguments passed to the object"
+    )
+    runtime: Optional[Union[str, float]] = Field(
+        title="Python version to use when executing ", default=None
+    )
+    imports: Optional[List[str]] = Field(
+        title="Stage and path to previously uploaded files you want to import",
+        default=[],
+    )
+    stage: str = Field(title="Stage in which artifacts will be stored")
+    artifacts: List[Union[PathMapping, str]] = Field(title="List of required sources")
+
+    @field_validator("artifacts")
+    @classmethod
+    def _convert_artifacts(cls, artifacts: Union[dict, str]):
+        _artifacts = []
+        for artefact in artifacts:
+            if isinstance(artefact, PathMapping):
+                _artifacts.append(artefact)
+            else:
+                _artifacts.append(PathMapping(src=artefact))
+        return _artifacts
+
+    @field_validator("runtime")
+    @classmethod
+    def convert_runtime(cls, runtime_input: Union[str, float]) -> str:
+        if isinstance(runtime_input, float):
+            return str(runtime_input)
+        return runtime_input
+
+    @field_validator("artifacts")
+    @classmethod
+    def validate_artifacts(cls, artifacts: List[Path]) -> List[Path]:
+        for artefact in artifacts:
+            if "*" in str(artefact):
+                raise ValueError("Glob patterns not supported for Snowpark artifacts.")
+        return artifacts
+
+    @property
+    def udf_sproc_identifier(self) -> UdfSprocIdentifier:
+        return UdfSprocIdentifier.from_definition(self)
+
+
+class ProcedureEntityModel(SnowparkEntityModel):
+    type: Literal["procedure"] = DiscriminatorField()  # noqa: A003
+    execute_as_caller: Optional[bool] = Field(
+        title="Determine whether the procedure is executed with the privileges of "
+        "the owner (you) or with the privileges of the caller",
+        default=False,
+    )
+
+
+class FunctionEntityModel(SnowparkEntityModel):
+    type: Literal["function"] = DiscriminatorField()  # noqa: A003
+
+
+class UdfSprocIdentifier:
+    def __init__(self, identifier: FQN, arg_names, arg_types, arg_defaults):
+        self._identifier = identifier
+        self._arg_names = arg_names
+        self._arg_types = arg_types
+        self._arg_defaults = arg_defaults
+
+    def _identifier_from_signature(self, sig: List[str], for_sql: bool = False):
+        signature = self._comma_join(sig)
+        id_ = self._identifier.sql_identifier if for_sql else self._identifier
+        return f"{id_}({signature})"
+
+    @staticmethod
+    def _comma_join(*args):
+        return ", ".join(*args)
+
+    @property
+    def identifier_with_arg_names(self):
+        return self._identifier_from_signature(self._arg_names)
+
+    @property
+    def identifier_with_arg_types(self):
+        return self._identifier_from_signature(self._arg_types)
+
+    @property
+    def identifier_with_arg_names_types(self):
+        sig = [f"{n} {t}" for n, t in zip(self._arg_names, self._arg_types)]
+        return self._identifier_from_signature(sig)
+
+    @property
+    def identifier_with_arg_names_types_defaults(self):
+        return self._identifier_from_signature(self._full_signature())
+
+    def _is_signature_type_a_string(self, sig_type: str) -> bool:
+        return sig_type.lower() in ["string", "varchar"]
+
+    def _full_signature(self):
+        sig = []
+        for name, _type, _default in zip(
+            self._arg_names, self._arg_types, self._arg_defaults
+        ):
+            s = f"{name} {_type}"
+            if _default:
+                if self._is_signature_type_a_string(_type):
+                    _default = f"'{_default}'"
+                s += f" default {_default}"
+            sig.append(s)
+        return sig
+
+    @property
+    def identifier_for_sql(self):
+        return self._identifier_from_signature(self._full_signature(), for_sql=True)
+
+    @classmethod
+    def from_definition(cls, udf_sproc: SnowparkEntityModel):
+        names = []
+        types = []
+        defaults = []
+        if udf_sproc.signature and udf_sproc.signature != "null":
+            for arg in udf_sproc.signature:
+                names.append(arg.name)  # type:ignore
+                types.append(arg.arg_type)  # type:ignore
+                defaults.append(arg.default)  # type:ignore
+
+        identifier = udf_sproc.fqn.using_context()
+        return cls(identifier, names, types, defaults)

snowflake/cli/api/project/schemas/entities/streamlit_entity_model.py

@@ -0,0 +1,73 @@
+# 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
+from typing import List, Literal, Optional
+
+from pydantic import Field, model_validator
+from snowflake.cli.api.project.schemas.entities.common import (
+    EntityModelBase,
+    ExternalAccessBaseModel,
+)
+from snowflake.cli.api.project.schemas.updatable_model import (
+    DiscriminatorField,
+)
+
+
+class StreamlitEntityModel(EntityModelBase, ExternalAccessBaseModel):
+    type: Literal["streamlit"] = DiscriminatorField()  # noqa: A003
+    title: Optional[str] = Field(
+        title="Human-readable title for the Streamlit dashboard", default=None
+    )
+    query_warehouse: str = Field(
+        title="Snowflake warehouse to host the app", default=None
+    )
+    main_file: Optional[str] = Field(
+        title="Entrypoint file of the Streamlit app", default="streamlit_app.py"
+    )
+    pages_dir: Optional[str] = Field(title="Streamlit pages", default=None)
+    stage: Optional[str] = Field(
+        title="Stage in which the app’s artifacts will be stored", default="streamlit"
+    )
+    # Possibly can be PathMapping
+    artifacts: Optional[List[Path]] = Field(
+        title="List of files which should be deployed. Each file needs to exist locally. "
+        "Main file needs to be included in the artifacts.",
+        default=None,
+    )
+
+    @model_validator(mode="after")
+    def main_file_must_be_in_artifacts(self):
+        if not self.artifacts:
+            return self
+
+        if Path(self.main_file) not in self.artifacts:
+            raise ValueError(
+                f"Specified main file {self.main_file} is not included in artifacts."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def artifacts_must_exists(self):
+        if not self.artifacts:
+            return self
+
+        for artifact in self.artifacts:
+            if not artifact.exists():
+                raise ValueError(
+                    f"Specified artifact {artifact} does not exist locally."
+                )
+
+        return self

snowflake/cli/api/project/schemas/identifier_model.py

@@ -17,7 +17,16 @@ from __future__ import annotations
 from typing import Optional, cast
 
 from pydantic import Field
-from snowflake.cli.api.project.schemas.updatable_model import IdentifierField
+from snowflake.cli.api.project.schemas.updatable_model import (
+    IdentifierField,
+    UpdatableModel,
+)
+
+
+class Identifier(UpdatableModel):
+    name: str = Field(title="Entity name")
+    schema_: str = Field(title="Entity schema", alias="schema", default=None)
+    database: str = Field(title="Entity database", default=None)
 
 
 class ObjectIdentifierBaseModel:

snowflake/cli/api/project/schemas/native_app/application.py

@@ -16,19 +16,13 @@ from __future__ import annotations
 
 from typing import List, Optional
 
-from pydantic import Field
+from pydantic import Field, field_validator
+from snowflake.cli.api.project.schemas.entities.common import PostDeployHook
 from snowflake.cli.api.project.schemas.updatable_model import (
     IdentifierField,
     UpdatableModel,
 )
-
-
-class SqlScriptHookType(UpdatableModel):
-    sql_script: str = Field(title="SQL file path relative to the project root")
-
-
-# Currently sql_script is the only supported hook type. Change to a Union once other hook types are added
-PostDeployHook = SqlScriptHookType
+from snowflake.cli.api.project.util import append_test_resource_suffix
 
 
 class Application(UpdatableModel):
@@ -53,6 +47,11 @@ class Application(UpdatableModel):
         default=None,
     )
 
+    @field_validator("name")
+    @classmethod
+    def append_test_resource_suffix_to_name(cls, input_value: str) -> str:
+        return append_test_resource_suffix(input_value)
+
 
 class ApplicationV11(Application):
     # Templated defaults only supported in v1.1+

snowflake/cli/api/project/schemas/native_app/package.py

@@ -17,11 +17,12 @@ from __future__ import annotations
 from typing import List, Literal, Optional
 
 from pydantic import Field, field_validator, model_validator
-from snowflake.cli.api.project.schemas.native_app.application import PostDeployHook
+from snowflake.cli.api.project.schemas.entities.common import PostDeployHook
 from snowflake.cli.api.project.schemas.updatable_model import (
     IdentifierField,
     UpdatableModel,
 )
+from snowflake.cli.api.project.util import append_test_resource_suffix
 
 DistributionOptions = Literal["internal", "external", "INTERNAL", "EXTERNAL"]
 
@@ -50,6 +51,11 @@ class Package(UpdatableModel):
         default=None,
     )
 
+    @field_validator("name")
+    @classmethod
+    def append_test_resource_suffix_to_name(cls, input_value: str) -> str:
+        return append_test_resource_suffix(input_value)
+
     @field_validator("scripts")
     @classmethod
     def validate_scripts(cls, input_list):

snowflake/cli/api/project/schemas/project_definition.py

@@ -15,22 +15,21 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from packaging.version import Version
 from pydantic import Field, ValidationError, field_validator, model_validator
-from snowflake.cli.api.feature_flags import FeatureFlag
 from snowflake.cli.api.project.errors import SchemaValidationError
-from snowflake.cli.api.project.schemas.entities.application_entity import (
-    ApplicationEntity,
+from snowflake.cli.api.project.schemas.entities.application_entity_model import (
+    ApplicationEntityModel,
 )
 from snowflake.cli.api.project.schemas.entities.common import (
     DefaultsField,
     TargetField,
 )
 from snowflake.cli.api.project.schemas.entities.entities import (
-    Entity,
-    v2_entity_types_map,
+    EntityModel,
+    v2_entity_model_types_map,
 )
 from snowflake.cli.api.project.schemas.native_app.native_app import (
     NativeApp,
@@ -42,6 +41,8 @@ from snowflake.cli.api.project.schemas.updatable_model import UpdatableModel
 from snowflake.cli.api.utils.types import Context
 from typing_extensions import Annotated
 
+AnnotatedEntity = Annotated[EntityModel, Field(discriminator="type")]
+
 
 @dataclass
 class ProjectProperties:
@@ -111,9 +112,7 @@ class DefinitionV11(DefinitionV10):
 
 
 class DefinitionV20(_ProjectDefinitionBase):
-    entities: Dict[str, Annotated[Entity, Field(discriminator="type")]] = Field(
-        title="Entity definitions."
-    )
+    entities: Dict[str, AnnotatedEntity] = Field(title="Entity definitions.")
 
     @model_validator(mode="before")
     @classmethod
@@ -123,34 +122,52 @@ class DefinitionV20(_ProjectDefinitionBase):
         """
         if "defaults" in data and "entities" in data:
             for key, entity in data["entities"].items():
-                entity_type = entity["type"]
-                if entity_type not in v2_entity_types_map:
+                entity_fields = get_allowed_fields_for_entity(entity)
+                if not entity_fields:
                     continue
-                entity_model = v2_entity_types_map[entity_type]
                 for default_key, default_value in data["defaults"].items():
-                    if (
-                        default_key in entity_model.model_fields
-                        and default_key not in entity
-                    ):
+                    if default_key in entity_fields and default_key not in entity:
                         entity[default_key] = default_value
         return data
 
     @field_validator("entities", mode="after")
     @classmethod
-    def validate_entities(cls, entities: Dict[str, Entity]) -> Dict[str, Entity]:
+    def validate_entities_identifiers(
+        cls, entities: Dict[str, EntityModel]
+    ) -> Dict[str, EntityModel]:
+        for key, entity in entities.items():
+            entity.set_entity_id(key)
+            entity.validate_identifier()
+        return entities
+
+    @field_validator("entities", mode="after")
+    @classmethod
+    def validate_entities(
+        cls, entities: Dict[str, AnnotatedEntity]
+    ) -> Dict[str, AnnotatedEntity]:
         for key, entity in entities.items():
             # TODO Automatically detect TargetFields to validate
-            if entity.type == ApplicationEntity.get_type():
-                if isinstance(entity.from_, TargetField):
-                    target_key = entity.from_.target
-                    target_object = entity.from_
-                    target_type = target_object.get_type()
-                    cls._validate_target_field(target_key, target_type, entities)
+            if isinstance(entity, list):
+                for e in entity:
+                    cls._validate_single_entity(e, entities)
+            else:
+                cls._validate_single_entity(entity, entities)
         return entities
 
+    @classmethod
+    def _validate_single_entity(
+        cls, entity: EntityModel, entities: Dict[str, AnnotatedEntity]
+    ):
+        if entity.type == ApplicationEntityModel.get_type():
+            if isinstance(entity.from_, TargetField):
+                target_key = entity.from_.target
+                target_object = entity.from_
+                target_type = target_object.get_type()
+                cls._validate_target_field(target_key, target_type, entities)
+
     @classmethod
     def _validate_target_field(
-        cls, target_key: str, target_type: Entity, entities: Dict[str, Entity]
+        cls, target_key: str, target_type: EntityModel, entities: Dict[str, EntityModel]
     ):
         if target_key not in entities:
             raise ValueError(f"No such target: {target_key}")
@@ -172,6 +189,39 @@ class DefinitionV20(_ProjectDefinitionBase):
         default=None,
     )
 
+    mixins: Optional[Dict[str, Dict]] = Field(
+        title="Mixins to apply to entities",
+        default=None,
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def apply_mixins(cls, data: Dict) -> Dict:
+        """
+        Applies mixins to those entities, whose meta field contains the mixin name.
+        """
+        if "mixins" not in data or "entities" not in data:
+            return data
+
+        for entity in data["entities"].values():
+            entity_mixins = entity_mixins_to_list(
+                entity.get("meta", {}).get("use_mixins")
+            )
+
+            entity_fields = get_allowed_fields_for_entity(entity)
+            if entity_fields and entity_mixins:
+                for mixin_name in entity_mixins:
+                    if mixin_name in data["mixins"]:
+                        for key, value in data["mixins"][mixin_name].items():
+                            if key in entity_fields:
+                                entity[key] = value
+                    else:
+                        raise ValueError(f"Mixin {mixin_name} not found in mixins")
+        return data
+
+    def get_entities_by_type(self, entity_type: str):
+        return {i: e for i, e in self.entities.items() if e.get_type() == entity_type}
+
 
 def build_project_definition(**data) -> ProjectDefinition:
     """
@@ -193,7 +243,28 @@ ProjectDefinition = Union[ProjectDefinitionV1, ProjectDefinitionV2]
 
 
 def get_version_map():
-    version_map = {"1": DefinitionV10, "1.1": DefinitionV11}
-    if FeatureFlag.ENABLE_PROJECT_DEFINITION_V2.is_enabled():
-        version_map["2"] = DefinitionV20
+    version_map = {"1": DefinitionV10, "1.1": DefinitionV11, "2": DefinitionV20}
     return version_map
+
+
+def entity_mixins_to_list(entity_mixins: Optional[str | List[str]]) -> List[str]:
+    """
+    Convert an optional string or a list of strings to a list of strings.
+    """
+    if entity_mixins is None:
+        return []
+    if isinstance(entity_mixins, str):
+        return [entity_mixins]
+    return entity_mixins
+
+
+def get_allowed_fields_for_entity(entity: Dict[str, Any]) -> List[str]:
+    """
+    Get the allowed fields for the given entity.
+    """
+    entity_type = entity.get("type")
+    if entity_type not in v2_entity_model_types_map:
+        return []
+
+    entity_model = v2_entity_model_types_map[entity_type]
+    return entity_model.model_fields

snowflake/cli/api/project/schemas/updatable_model.py

@@ -117,6 +117,7 @@ class UpdatableModel(BaseModel):
         # all the values of the class attributes. We go in reverse order so that
         # values in subclasses overrides values from parent classes in case of field overrides.
 
+        private_attrs = set()
         for class_ in reversed(cls.__mro__):
             class_dict = class_.__dict__
             field_annotations.update(class_dict.get("__annotations__", {}))
@@ -128,10 +129,15 @@ class UpdatableModel(BaseModel):
             else:
                 # If Pydantic did not process this class yet, get the values from class_dict directly
                 field_values.update(class_dict)
+            for pa in class_dict.get("__private_attributes__", []):
+                private_attrs.add(pa)
 
         # Add Pydantic validation wrapper around all fields except `DiscriminatorField`s
         for field_name in field_annotations:
-            if not cls._is_entity_type_field(field_values.get(field_name)):
+            if field_name in private_attrs:
+                continue
+            field = field_values.get(field_name)
+            if not cls._is_entity_type_field(field):
                 cls._add_validator(field_name)
 
     @classmethod
@@ -154,7 +160,8 @@ class UpdatableModel(BaseModel):
 
         setattr(
             cls,
-            f"_field_validator_with_verbose_name_to_avoid_name_conflict_{field_name}",
+            # Unique name so that subclasses get a unique instance of this validator
+            f"_{cls.__module__}.{cls.__name__}_validate_{field_name}",
             field_validator(field_name, mode="wrap")(validator_skipping_templated_str),
         )
 
@@ -187,7 +194,8 @@ def DiscriminatorField(*args, **kwargs):  # noqa N802
     When this `DiscriminatorField` is used on a pydantic attribute,
     we will not allow templating on it.
     """
-    return Field(is_discriminator_field=True, *args, **kwargs)
+    extra = dict(is_discriminator_field=True)
+    return Field(json_schema_extra=extra, *args, **kwargs)
 
 
 def IdentifierField(*args, **kwargs):  # noqa N802

snowflake/cli/api/project/util.py

@@ -34,6 +34,9 @@ UNQUOTED_IDENTIFIER_REGEX = r"([a-zA-Z_])([a-zA-Z0-9_$]{0,254})"
 QUOTED_IDENTIFIER_REGEX = r'"((""|[^"]){0,255})"'
 VALID_IDENTIFIER_REGEX = f"(?:{UNQUOTED_IDENTIFIER_REGEX}|{QUOTED_IDENTIFIER_REGEX})"
 
+# An env var that is used to suffix the names of some account-level resources
+TEST_RESOURCE_SUFFIX_VAR = "SNOWFLAKE_CLI_TEST_RESOURCE_SUFFIX"
+
 
 def encode_uri_component(s: str) -> str:
     """
@@ -191,12 +194,6 @@ def extract_schema(qualified_name: str):
     return None
 
 
-def generate_user_env(username: str) -> dict:
-    return {
-        "USER": username,
-    }
-
-
 def first_set_env(*keys: str):
     for k in keys:
         v = os.getenv(k)
@@ -259,3 +256,23 @@ def identifier_to_show_like_pattern(identifier: str) -> str:
     matching this identifier
     """
     return f"'{escape_like_pattern(unquote_identifier(identifier))}'"
+
+
+def append_test_resource_suffix(identifier: str) -> str:
+    """
+    Append a suffix that should be added to specified account-level resources.
+
+    This is an internal concern that is currently only used in tests
+    to isolate concurrent runs and to add the test name to resources.
+    """
+    suffix = os.environ.get(TEST_RESOURCE_SUFFIX_VAR, "")
+    if identifier_to_str(identifier).endswith(identifier_to_str(suffix)):
+        # If the suffix has already been added, don't add it again
+        return identifier
+    if is_valid_quoted_identifier(identifier) or is_valid_quoted_identifier(suffix):
+        # If either identifier is already quoted, use concat_identifier
+        # to add the suffix inside the quotes
+        return concat_identifiers([identifier, suffix])
+    # Otherwise just append the string, don't add quotes
+    # in case the user doesn't want them
+    return f"{identifier}{suffix}"

snowflake/cli/api/rendering/jinja.py

@@ -17,7 +17,7 @@ from __future__ import annotations
 
 from pathlib import Path
 from textwrap import dedent
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 import jinja2
 from jinja2 import Environment, StrictUndefined, loaders
@@ -82,8 +82,18 @@ class IgnoreAttrEnvironment(Environment):
             return self.undefined(obj=obj, name=argument)
 
 
+def get_basic_jinja_env(loader: Optional[loaders.BaseLoader] = None) -> Environment:
+    return env_bootstrap(
+        IgnoreAttrEnvironment(
+            loader=loader or loaders.BaseLoader(),
+            keep_trailing_newline=True,
+            undefined=StrictUndefined,
+        )
+    )
+
+
 def jinja_render_from_file(
-    template_path: Path, data: Dict, output_file_path: Optional[Path] = None
+    template_path: Path, data: Dict[str, Any], output_file_path: Optional[Path] = None
 ) -> Optional[str]:
     """
     Renders a jinja template and outputs either the rendered contents as string or writes to a file.
@@ -96,12 +106,8 @@ def jinja_render_from_file(
     Returns:
         None if file path is provided, else returns the rendered string.
     """
-    env = env_bootstrap(
-        IgnoreAttrEnvironment(
-            loader=loaders.FileSystemLoader(template_path.parent),
-            keep_trailing_newline=True,
-            undefined=StrictUndefined,
-        )
+    env = get_basic_jinja_env(
+        loader=loaders.FileSystemLoader(template_path.parent.as_posix())
     )
     loaded_template = env.get_template(template_path.name)
     rendered_result = loaded_template.render(**data)

snowflake/cli/api/rendering/project_definition_templates.py

@@ -24,7 +24,7 @@ _YML_TEMPLATE_START = "<%"
 _YML_TEMPLATE_END = "%>"
 
 
-def get_project_definition_cli_jinja_env() -> Environment:
+def get_client_side_jinja_env() -> Environment:
     _random_block = "___very___unique___block___to___disable___logic___blocks___"
     return env_bootstrap(
         IgnoreAttrEnvironment(