snowflake-cli 2.8.2__py3-none-any.whl

snowflake/cli/api/sql_execution.py

@@ -0,0 +1,260 @@
+# 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 logging
+from contextlib import contextmanager
+from functools import cached_property
+from io import StringIO
+from textwrap import dedent
+from typing import Iterable, Optional, Tuple
+
+from snowflake.cli.api.cli_global_context import cli_context
+from snowflake.cli.api.console import cli_console
+from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.exceptions import (
+    DatabaseNotProvidedError,
+    SchemaNotProvidedError,
+    SnowflakeSQLExecutionError,
+)
+from snowflake.cli.api.identifiers import FQN
+from snowflake.cli.api.project.util import (
+    identifier_to_show_like_pattern,
+    unquote_identifier,
+)
+from snowflake.cli.api.utils.cursor import find_first_row
+from snowflake.connector.cursor import DictCursor, SnowflakeCursor
+from snowflake.connector.errors import ProgrammingError
+
+
+class SqlExecutionMixin:
+    def __init__(self):
+        self._snowpark_session = None
+
+    @property
+    def _conn(self):
+        return cli_context.connection
+
+    @property
+    def snowpark_session(self):
+        if not self._snowpark_session:
+            from snowflake.snowpark.session import Session
+
+            self._snowpark_session = Session.builder.configs(
+                {"connection": self._conn}
+            ).create()
+        return self._snowpark_session
+
+    @cached_property
+    def _log(self):
+        return logging.getLogger(__name__)
+
+    def _execute_string(
+        self,
+        sql_text: str,
+        remove_comments: bool = False,
+        return_cursors: bool = True,
+        cursor_class: SnowflakeCursor = SnowflakeCursor,
+        **kwargs,
+    ) -> Iterable[SnowflakeCursor]:
+        """
+        This is a custom implementation of SnowflakeConnection.execute_string that returns generator
+        instead of list. In case of executing multiple queries are executed one by one. This mean we can
+        access result of previous queries while evaluating next one. For example, we can print the results.
+        """
+        self._log.debug("Executing %s", sql_text)
+        stream = StringIO(sql_text)
+        stream_generator = self._conn.execute_stream(
+            stream, remove_comments=remove_comments, cursor_class=cursor_class, **kwargs
+        )
+        return stream_generator if return_cursors else list()
+
+    def _execute_query(self, query: str, **kwargs):
+        *_, last_result = self._execute_queries(query, **kwargs)
+        return last_result
+
+    def _execute_queries(self, queries: str, **kwargs):
+        return list(self._execute_string(dedent(queries), **kwargs))
+
+    def use(self, object_type: ObjectType, name: str):
+        try:
+            self._execute_query(f"use {object_type.value.sf_name} {name}")
+        except ProgrammingError:
+            # Rewrite the error to make the message more useful.
+            raise ProgrammingError(
+                f"Could not use {object_type} {name}. Object does not exist, or operation cannot be performed."
+            )
+
+    @contextmanager
+    def use_role(self, new_role: str):
+        """
+        Switches to a different role for a while, then switches back.
+        This is a no-op if the requested role is already active.
+        """
+        role_result = self._execute_query(
+            f"select current_role()", cursor_class=DictCursor
+        ).fetchone()
+        prev_role = role_result["CURRENT_ROLE()"]
+        is_different_role = new_role.lower() != prev_role.lower()
+        if is_different_role:
+            self._log.debug("Assuming different role: %s", new_role)
+            self._execute_query(f"use role {new_role}")
+        try:
+            yield
+        finally:
+            if is_different_role:
+                self._execute_query(f"use role {prev_role}")
+
+    @contextmanager
+    def use_warehouse(self, new_wh: str):
+        """
+        Switches to a different warehouse for a while, then switches back.
+        This is a no-op if the requested warehouse is already active.
+        If there is no default warehouse in the account, it will throw an error.
+        """
+
+        wh_result = self._execute_query(
+            f"select current_warehouse()", cursor_class=DictCursor
+        ).fetchone()
+        # If user has an assigned default warehouse, prev_wh will contain a value even if the warehouse is suspended.
+        try:
+            prev_wh = wh_result["CURRENT_WAREHOUSE()"]
+        except:
+            prev_wh = None
+
+        # new_wh is not None, and should already be a valid identifier, no additional check is performed here.
+        is_different_wh = new_wh != prev_wh
+        try:
+            if is_different_wh:
+                self._log.debug("Using warehouse: %s", new_wh)
+                self.use(object_type=ObjectType.WAREHOUSE, name=new_wh)
+            yield
+        finally:
+            if prev_wh and is_different_wh:
+                self._log.debug("Switching back to warehouse: %s", prev_wh)
+                self.use(object_type=ObjectType.WAREHOUSE, name=prev_wh)
+
+    def create_password_secret(
+        self, name: FQN, username: str, password: str
+    ) -> SnowflakeCursor:
+        return self._execute_query(
+            f"""
+            create secret {name.sql_identifier}
+            type = password
+            username = '{username}'
+            password = '{password}'
+            """
+        )
+
+    def create_api_integration(
+        self, name: FQN, api_provider: str, allowed_prefix: str, secret: Optional[str]
+    ) -> SnowflakeCursor:
+        return self._execute_query(
+            f"""
+            create api integration {name.sql_identifier}
+            api_provider = {api_provider}
+            api_allowed_prefixes = ('{allowed_prefix}')
+            allowed_authentication_secrets = ({secret if secret else ''})
+            enabled = true
+            """
+        )
+
+    def _execute_schema_query(self, query: str, name: Optional[str] = None, **kwargs):
+        """
+        Check that a database and schema are provided before executing the query. Useful for operating on schema level objects.
+        """
+        self.check_database_and_schema_provided(name)
+        return self._execute_query(query, **kwargs)
+
+    def check_database_and_schema_provided(self, name: Optional[str] = None) -> None:
+        """
+        Checks if a database and schema are provided, either through the connection context or a qualified name.
+        """
+        fqn = FQN.from_string(name).using_connection(self._conn)
+        if not fqn.database:
+            raise DatabaseNotProvidedError()
+        if not fqn.schema:
+            raise SchemaNotProvidedError()
+
+    @staticmethod
+    def _qualified_name_to_in_clause(identifier: FQN) -> Tuple[str, Optional[str]]:
+        if identifier.database:
+            schema = identifier.schema or "PUBLIC"
+            in_clause = f"in schema {identifier.database}.{schema}"
+        elif identifier.schema:
+            in_clause = f"in schema {identifier.schema}"
+        else:
+            in_clause = None
+        return identifier.name, in_clause
+
+    class InClauseWithQualifiedNameError(ValueError):
+        def __init__(self):
+            super().__init__("non-empty 'in_clause' passed with qualified 'name'")
+
+    def show_specific_object(
+        self,
+        object_type_plural: str,
+        name: str,
+        name_col: str = "name",
+        in_clause: str = "",
+        check_schema: bool = False,
+    ) -> Optional[dict]:
+        """
+        Executes a "show <objects> like" query for a particular entity with a
+        given (optionally qualified) name. This command is useful when the corresponding
+        "describe <object>" query does not provide the information you seek.
+
+        Note that this command is analogous to describe and should only return a single row.
+        If the target object type is a schema level object, then check_schema should be set to True
+        so that the function will verify that a database and schema are provided, either through
+        the connection or a qualified name, before executing the query.
+        """
+
+        unqualified_name, name_in_clause = self._qualified_name_to_in_clause(
+            FQN.from_string(name)
+        )
+        if in_clause and name_in_clause:
+            raise self.InClauseWithQualifiedNameError()
+        elif name_in_clause:
+            in_clause = name_in_clause
+        show_obj_query = f"show {object_type_plural} like {identifier_to_show_like_pattern(unqualified_name)} {in_clause}".strip()
+
+        if check_schema:
+            show_obj_cursor = self._execute_schema_query(  # type: ignore
+                show_obj_query, name=name, cursor_class=DictCursor
+            )
+        else:
+            show_obj_cursor = self._execute_query(  # type: ignore
+                show_obj_query, cursor_class=DictCursor
+            )
+
+        if show_obj_cursor.rowcount is None:
+            raise SnowflakeSQLExecutionError(show_obj_query)
+        elif show_obj_cursor.rowcount > 1:
+            raise ProgrammingError(
+                f"Received multiple rows from result of SQL statement: {show_obj_query}. Usage of 'show_specific_object' may not be properly scoped."
+            )
+
+        show_obj_row = find_first_row(
+            show_obj_cursor,
+            lambda row: row[name_col] == unquote_identifier(unqualified_name),
+        )
+        return show_obj_row
+
+
+class VerboseCursor(SnowflakeCursor):
+    def execute(self, command: str, *args, **kwargs):
+        cli_console.message(command)
+        super().execute(command, *args, **kwargs)

snowflake/cli/api/utils/__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/api/utils/cursor.py

@@ -0,0 +1,34 @@
+# 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 typing import Callable, List, Optional
+
+from snowflake.connector.cursor import DictCursor
+
+
+def _rows_generator(cursor: DictCursor, predicate: Callable[[dict], bool]):
+    return (row for row in cursor.fetchall() if predicate(row))
+
+
+def find_all_rows(cursor: DictCursor, predicate: Callable[[dict], bool]) -> List[dict]:
+    return list(_rows_generator(cursor, predicate))
+
+
+def find_first_row(
+    cursor: DictCursor, predicate: Callable[[dict], bool]
+) -> Optional[dict]:
+    """Returns the first row that matches the predicate, or None."""
+    return next(_rows_generator(cursor, predicate), None)

snowflake/cli/api/utils/definition_rendering.py

@@ -0,0 +1,383 @@
+# 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 copy
+from typing import Any, Optional
+
+from jinja2 import Environment, TemplateSyntaxError, nodes
+from packaging.version import Version
+from snowflake.cli.api.console import cli_console as cc
+from snowflake.cli.api.exceptions import CycleDetectedError, InvalidTemplate
+from snowflake.cli.api.project.schemas.project_definition import (
+    ProjectProperties,
+    build_project_definition,
+)
+from snowflake.cli.api.project.schemas.updatable_model import context
+from snowflake.cli.api.rendering.jinja import CONTEXT_KEY
+from snowflake.cli.api.rendering.project_definition_templates import (
+    get_project_definition_cli_jinja_env,
+)
+from snowflake.cli.api.utils.dict_utils import deep_merge_dicts, traverse
+from snowflake.cli.api.utils.graph import Graph, Node
+from snowflake.cli.api.utils.models import ProjectEnvironment
+from snowflake.cli.api.utils.templating_functions import get_templating_functions
+from snowflake.cli.api.utils.types import Context, Definition
+
+
+class TemplatedEnvironment:
+    """
+    This class is a utility class
+    that encapsulates some of the Jinja Templating functionality.
+    """
+
+    def __init__(self, env: Environment):
+        self._jinja_env: Environment = env
+
+    def render(self, template_value: Any, context: Context) -> Any:
+        if not self.get_referenced_vars(template_value):
+            return template_value
+
+        jinja_template = self._jinja_env.from_string(str(template_value))
+        return jinja_template.render(context)
+
+    def get_referenced_vars(self, template_value: Any) -> set[TemplateVar]:
+        template_str = str(template_value)
+        try:
+            ast = self._jinja_env.parse(template_str)
+        except TemplateSyntaxError as e:
+            raise InvalidTemplate(
+                f"Error parsing template from project definition file. Value: '{template_str}'. Error: {e}"
+            ) from e
+
+        return self._get_referenced_vars(ast, template_str)
+
+    def _get_referenced_vars(
+        self,
+        ast_node: nodes.Template,
+        template_value: str,
+        current_attr_chain: Optional[list[str]] = None,
+    ) -> set[TemplateVar]:
+        """
+        Traverse Jinja AST to find the variable chain referenced by the template.
+        A variable like ctx.env.test is internally represented in the AST tree as
+        Getattr Node (attr='test') -> Getattr Node (attr='env') -> Name Node (name='ctx')
+        """
+        all_referenced_vars: set[TemplateVar] = set()
+        if isinstance(ast_node, nodes.Getattr):
+            current_attr_chain = [ast_node.attr] + (current_attr_chain or [])  # type: ignore[attr-defined]
+        elif isinstance(ast_node, nodes.Name):
+            current_attr_chain = [ast_node.name] + (current_attr_chain or [])  # type: ignore[attr-defined]
+            all_referenced_vars.add(TemplateVar(current_attr_chain))
+            current_attr_chain = None
+        elif (
+            not isinstance(
+                ast_node,
+                (
+                    nodes.Template,
+                    nodes.TemplateData,
+                    nodes.Output,
+                    nodes.Call,
+                    nodes.Const,
+                    nodes.Filter,
+                ),
+            )
+            or current_attr_chain is not None
+        ):
+            raise InvalidTemplate(f"Unexpected templating syntax in {template_value}")
+
+        for child_node in ast_node.iter_child_nodes():
+            all_referenced_vars.update(
+                self._get_referenced_vars(
+                    child_node, template_value, current_attr_chain
+                )
+            )
+
+        return all_referenced_vars
+
+
+class TemplateVar:
+    """
+    This class tracks template variable information.
+    For a variable like ctx.env.var, this class will track
+    the chain of keys referenced by this variable (ctx, env, var),
+    as well as the value of this variable. (e.g. ctx.env.var = "hello_<% ctx.definition_version %>")
+
+    The value of this variable is divided into 2 parts.
+    The templated value (e.g. "hello_<% ctx.definition %>"),
+    as well as the rendered_value (e.g. "hello_1.1")
+    """
+
+    def __init__(self, vars_chain):
+        self._vars_chain: list[str] = list(vars_chain)
+        self.templated_value: Optional[Any] = None
+        self.rendered_value: Optional[Any] = None
+
+    @property
+    def key(self) -> str:
+        return ".".join(self._vars_chain)
+
+    @property
+    def is_env_var(self) -> bool:
+        return (
+            len(self._vars_chain) == 3
+            and self._vars_chain[0] == CONTEXT_KEY
+            and self._vars_chain[1] == "env"
+        )
+
+    def get_env_var_name(self) -> str:
+        if not self.is_env_var:
+            raise KeyError(
+                f"Referenced variable {self.key} is not an environment variable"
+            )
+        return self._vars_chain[2]
+
+    def add_to_context(self, context: Context) -> None:
+        """
+        Takes a multi-level context dict as input. Modifies the context dict with the rendered value of this variable.
+
+        If the variable has multi-levels (e.g. ctx.env), recursively traverse the dictionary
+        to set the inner level's key to the rendered value of this variable.
+
+        Example: vars chain contains ['ctx', 'env', 'x'], and context is {}, and rendered_value is 'val'.
+        At the end of this call, context content will be: {'ctx': {'env': {'x': 'val'}}}
+        """
+        current_dict_level = context
+        last_element_index = len(self._vars_chain) - 1
+        for index, var in enumerate(self._vars_chain):
+            if index == last_element_index:
+                current_dict_level[var] = self.rendered_value
+            else:
+                current_dict_level.setdefault(var, {})
+                current_dict_level = current_dict_level[var]
+
+    def read_from_context(self, context: Context) -> Any:
+        """
+        Takes a multi-level context dict as input.
+
+        If the variable has multi-levels (e.g. ctx.env), recursively traverse the dictionary
+        to find the key that the variable points to.
+
+        Returns the value in that location.
+
+        Raise InvalidTemplate if the variable is None or not found.
+        """
+        current_dict_level = context
+        for key in self._vars_chain:
+            if (
+                not isinstance(current_dict_level, dict)
+                or key not in current_dict_level
+            ):
+                raise InvalidTemplate(f"Could not find template variable {self.key}")
+            current_dict_level = current_dict_level[key]
+
+        value = current_dict_level
+
+        if value is None:
+            raise InvalidTemplate(f"Template variable {self.key} does not have a value")
+
+        if isinstance(value, (dict, list)):
+            raise InvalidTemplate(
+                f"Template variable {self.key} does not have a scalar value"
+            )
+
+        return value
+
+    def __hash__(self):
+        return hash(self.key)
+
+    def __eq__(self, other):
+        return self.key == other.key
+
+
+def _build_dependency_graph(
+    env: TemplatedEnvironment,
+    all_vars: set[TemplateVar],
+    context: Context,
+    environment_overrides: ProjectEnvironment,
+) -> Graph[TemplateVar]:
+    dependencies_graph = Graph[TemplateVar]()
+    for variable in all_vars:
+        dependencies_graph.add(Node[TemplateVar](key=variable.key, data=variable))
+    for variable in all_vars:
+        # If variable is found in os.environ or from cli override, then use the value as is
+        # skip rendering by pre-setting the rendered_value attribute
+        if variable.is_env_var and variable.get_env_var_name() in environment_overrides:
+            env_value = environment_overrides.get(variable.get_env_var_name())
+            variable.rendered_value = env_value
+            variable.templated_value = env_value
+        else:
+            variable.templated_value = variable.read_from_context(context)
+            dependencies_vars = env.get_referenced_vars(variable.templated_value)
+
+            for referenced_var in dependencies_vars:
+                dependencies_graph.add_directed_edge(variable.key, referenced_var.key)
+
+    return dependencies_graph
+
+
+def _render_graph_node(env: TemplatedEnvironment, node: Node[TemplateVar]) -> None:
+    if node.data.rendered_value is not None:
+        # Do not re-evaluate resolved nodes like env variable nodes
+        # which might contain template-like values, or non-string nodes
+        return
+
+    current_context: Context = {}
+    for dep_node in node.neighbors:
+        dep_node.data.add_to_context(current_context)
+
+    node.data.rendered_value = env.render(node.data.templated_value, current_context)
+
+
+def _validate_env_section(env_section: dict):
+    if not isinstance(env_section, dict):
+        raise InvalidTemplate(
+            "env section in project definition file should be a mapping"
+        )
+    for variable, value in env_section.items():
+        if value is None or isinstance(value, (dict, list)):
+            raise InvalidTemplate(
+                f"Variable {variable} in env section of project definition file should be a scalar"
+            )
+
+
+def _get_referenced_vars_in_definition(
+    template_env: TemplatedEnvironment, definition: Definition
+):
+    referenced_vars = set()
+
+    def find_any_template_vars(element):
+        referenced_vars.update(template_env.get_referenced_vars(element))
+
+    traverse(definition, visit_action=find_any_template_vars)
+
+    return referenced_vars
+
+
+def _template_version_warning():
+    cc.warning(
+        "Ignoring template pattern in project definition file. "
+        "Update 'definition_version' to 1.1 or later in snowflake.yml to enable template expansion."
+    )
+
+
+def _add_defaults_to_definition(original_definition: Definition) -> Definition:
+    with context({"skip_validation_on_templates": True}):
+        # pass a flag to Pydantic to skip validation for templated scalars
+        # populate the defaults
+        project_definition = build_project_definition(**original_definition)
+
+    definition_with_defaults = project_definition.model_dump(
+        exclude_none=True, warnings=False, by_alias=True
+    )
+    # The main purpose of the above operation was to populate defaults from Pydantic.
+    # By merging the original definition back in, we ensure that any transformations
+    # that Pydantic would have performed are undone.
+    deep_merge_dicts(definition_with_defaults, original_definition)
+    return definition_with_defaults
+
+
+def render_definition_template(
+    original_definition: Optional[Definition], context_overrides: Context
+) -> ProjectProperties:
+    """
+    Takes a definition file as input. An arbitrary structure containing dict|list|scalars,
+    with the top level being a dictionary.
+    Requires item 'definition_version' to be set to a version of 1.1 and higher.
+
+    Searches for any templating in all of the scalar fields, and attempts to resolve them
+    from the definition structure itself or from the environment variable.
+
+    Environment variables take precedence during the rendering process.
+    """
+
+    # copy input to protect it from update
+    definition = copy.deepcopy(original_definition)
+
+    # collect all the override --env variables passed through CLI input
+    override_env = context_overrides.get(CONTEXT_KEY, {}).get("env", {})
+
+    # set up Project Environment with empty default_env because
+    # default env section from project definition file is still templated at this time
+    environment_overrides = ProjectEnvironment(
+        default_env={}, override_env=override_env
+    )
+
+    if definition is None:
+        return ProjectProperties(None, {CONTEXT_KEY: {"env": environment_overrides}})
+
+    template_env = TemplatedEnvironment(get_project_definition_cli_jinja_env())
+
+    if "definition_version" not in definition or Version(
+        definition["definition_version"]
+    ) < Version("1.1"):
+        try:
+            referenced_vars = _get_referenced_vars_in_definition(
+                template_env, definition
+            )
+            if referenced_vars:
+                _template_version_warning()
+        except Exception:
+            # also warn on Exception, as it means the user is incorrectly attempting to use templating
+            _template_version_warning()
+
+        project_definition = build_project_definition(**definition)
+        project_context = {CONTEXT_KEY: definition}
+        project_context[CONTEXT_KEY]["env"] = environment_overrides
+        return ProjectProperties(project_definition, project_context)
+
+    definition = _add_defaults_to_definition(definition)
+    project_context = {CONTEXT_KEY: definition}
+
+    _validate_env_section(definition.get("env", {}))
+
+    # add available templating functions
+    project_context["fn"] = get_templating_functions()
+
+    referenced_vars = _get_referenced_vars_in_definition(template_env, definition)
+
+    dependencies_graph = _build_dependency_graph(
+        template_env, referenced_vars, project_context, environment_overrides
+    )
+
+    def on_cycle_action(node: Node[TemplateVar]):
+        raise CycleDetectedError(
+            f"Cycle detected in templating variable {node.data.key}"
+        )
+
+    dependencies_graph.dfs(
+        visit_action=lambda node: _render_graph_node(template_env, node),
+        on_cycle_action=on_cycle_action,
+    )
+
+    # now that we determined the values of all templated vars,
+    # use these resolved values as a fresh context to resolve definition
+    final_context: Context = {}
+    for node in dependencies_graph.get_all_nodes():
+        node.data.add_to_context(final_context)
+
+    traverse(
+        definition,
+        update_action=lambda val: template_env.render(val, final_context),
+    )
+
+    project_definition = build_project_definition(**definition)
+    project_context[CONTEXT_KEY] = definition
+    # Use `ProjectEnvironment` in project context in order to
+    # handle env variables overrides from OS env and from CLI arguments.
+    project_context[CONTEXT_KEY]["env"] = ProjectEnvironment(
+        default_env=project_context[CONTEXT_KEY].get("env"), override_env=override_env
+    )
+    return ProjectProperties(project_definition, project_context)