snowflake-cli 3.7.2__py3-none-any.whl → 3.8.0__py3-none-any.whl

snowflake/cli/_plugins/sql/lexer/keywords.py

@@ -0,0 +1,529 @@
+KEYWORDS = (
+    "ABORT",
+    "ACCESS",
+    "ACCOUNT",
+    "ACCOUNTS",
+    "ACTION",
+    "ACTIVATE",
+    "ADD",
+    "ADMIN",
+    "AFTER",
+    "ALL",
+    "ALTER",
+    "AND",
+    "ANY",
+    "API",
+    "APPEND_ONLY",
+    "APPLICATION",
+    "APPLICATIONS",
+    "APPLY",
+    "ARRAY",
+    "AS",
+    "ASC",
+    "AT",
+    "AUDITING",
+    "AUTHORIZATION",
+    "AUTHORIZATIONS",
+    "AUTOINCREMENT",
+    "BEFORE",
+    "BEGIN",
+    "BERNOULLI",
+    "BETWEEN",
+    "BIGINT",
+    "BINARY",
+    "BIT",
+    "BLOCK",
+    "BODY",
+    "BOOLEAN",
+    "BREAK",
+    "BUILD",
+    "BUILTIN",
+    "BY",
+    "BYTE",
+    "BYTEINT",
+    "CACHE",
+    "CALL",
+    "CALLED",
+    "CALLER",
+    "CASCADE",
+    "CASE",
+    "CAST",
+    "CATEGORIES",
+    "CHANGE",
+    "CHANGES",
+    "CHAR",
+    "CHARACTER",
+    "CHECK",
+    "CHILD",
+    "CLEAREPCACHE",
+    "CLICKLAG",
+    "CLONE",
+    "CLOSE",
+    "CLUSTER",
+    "CLUSTERING",
+    "CLUSTERS",
+    "COLLATE",
+    "COLUMN",
+    "COLUMNS",
+    "COMMENT",
+    "COMMIT",
+    "COMPACT",
+    "COMPRESSION",
+    "CONNECT",
+    "CONNECTION",
+    "CONNECTIONS",
+    "CONNECT_BY_ROOT",
+    "CONSTRAINT",
+    "CONSTRAINTS",
+    "CONSUMER",
+    "CONSUMERS",
+    "CONTINUE",
+    "COPY",
+    "CP",
+    "CREATE",
+    "CREDENTIALS",
+    "CROSS",
+    "CUBE",
+    "CURRENT",
+    "D",
+    "DATA",
+    "DATABASE",
+    "DATABASES",
+    "DATE",
+    "DATETIME",
+    "DATE_PART",
+    "DEC",
+    "DECIMAL",
+    "DECLARE",
+    "DEFAULT",
+    "DEFERRABLE",
+    "DEFERRED",
+    "DEFRAGMENT",
+    "DELEGATED",
+    "DELETE",
+    "DELIMITED",
+    "DELTA",
+    "DEPLOYMENT",
+    "DEPLOYMENTS",
+    "DESC",
+    "DESCRIBE",
+    "DIRECTORY",
+    "DISABLE",
+    "DISCONNECT",
+    "DISTINCT",
+    "DIV",
+    "DO",
+    "DOUBLE",
+    "DROP",
+    "DUAL",
+    "ELSE",
+    "ELSEIF",
+    "EMITNULL",
+    "ENABLE",
+    "END",
+    "ENFORCED",
+    "ERROR",
+    "ESCAPE",
+    "EXCEPT",
+    "EXCEPTION",
+    "EXCHANGE",
+    "EXCHANGES",
+    "EXCLUSIVE",
+    "EXECUTE",
+    "EXECUTE_SDL",
+    "EXECUTION",
+    "EXISTS",
+    "EXIT",
+    "EXPLAIN",
+    "EXPORTED",
+    "EXTERNAL",
+    "EXTRACT",
+    "FAIL",
+    "FAILOVER",
+    "FALSE",
+    "FETCH",
+    "FIELDS",
+    "FILE",
+    "FILES",
+    "FIRST",
+    "FIRST_VALUE",
+    "FLOAT",
+    "FN",
+    "FOLLOWING",
+    "FOR",
+    "FORALL",
+    "FORCE",
+    "FOREIGN",
+    "FORMAT",
+    "FORMATS",
+    "FREE",
+    "FROM",
+    "FULL",
+    "FUNCTION",
+    "FUNCTIONS",
+    "FUTURE",
+    "GET",
+    "GLOBAL",
+    "GOTO",
+    "GRANT",
+    "GRANTS",
+    "GROUP",
+    "GROUPING",
+    "GSCLUSTER",
+    "GSCLUSTERS",
+    "GSINSTANCE",
+    "GSINSTANCES",
+    "GSMAPPINGS",
+    "GSPACKAGES",
+    "HAVING",
+    "HELP",
+    "HIBERNATE",
+    "HINTS",
+    "HISTORY",
+    "IDENTIFIED",
+    "IDENTIFIER",
+    "IDENTITY",
+    "IF",
+    "IGNORE",
+    "ILIKE",
+    "ILIKE",
+    "IMMEDIATE",
+    "IMMUTABLE",
+    "IMPORT",
+    "IMPORTED",
+    "IN",
+    "INBOUND",
+    "INBOUND_FDN",
+    "INCIDENT",
+    "INCIDENTS",
+    "INCOMING",
+    "INCREMENT",
+    "INDEX",
+    "INDEXES",
+    "INFO",
+    "INFORMATION",
+    "INITIALLY",
+    "INNER",
+    "INPUT",
+    "INSERT",
+    "INSERT",
+    "INSERT_ONLY",
+    "INT",
+    "INTEGER",
+    "INTEGRATION",
+    "INTEGRATIONS",
+    "INTERLEAVED",
+    "INTERSECT",
+    "INTERVAL",
+    "INTO",
+    "IS",
+    "ISSUE",
+    "ISSUES",
+    "ITERATE",
+    "JOBS",
+    "JOIN",
+    "KEEP",
+    "KEY",
+    "KEYS",
+    "KIND",
+    "LAG",
+    "LANGUAGE",
+    "LAST",
+    "LAST_VALUE",
+    "LATERAL",
+    "LEAD",
+    "LEAVE",
+    "LEFT",
+    "LIKE",
+    "LIKE",
+    "LIMIT",
+    "LINEAR",
+    "LINK",
+    "LIST",
+    "LISTING",
+    "LISTINGS",
+    "LOAD",
+    "LOCAL",
+    "LOCATION",
+    "LOCK",
+    "LOCKS",
+    "LOOP",
+    "LS",
+    "MANAGE",
+    "MANAGED",
+    "MAP",
+    "MASKING",
+    "MATCH",
+    "MATCHED",
+    "MATERIALIZED",
+    "MEMBERS",
+    "MERGE",
+    "METASTORE",
+    "METASTORES",
+    "MINUS",
+    "MODIFY",
+    "MONITOR",
+    "MONITORS",
+    "NAME",
+    "NATURAL",
+    "NETWORK",
+    "NEXT",
+    "NO",
+    "NORELY",
+    "NOT",
+    "NOTIFICATION",
+    "NOTIFICATIONS",
+    "NOTIFY",
+    "NOVALIDATE",
+    "NTH_VALUE",
+    "NULL",
+    "NULLS",
+    "NUMBER",
+    "NUMERIC",
+    "OBJECT",
+    "OBJECTS",
+    "OF",
+    "OFFSET",
+    "OJ",
+    "ON",
+    "ONLY",
+    "OPEN",
+    "OPERATE",
+    "OPERATION",
+    "OPTION",
+    "OR",
+    "ORDER",
+    "ORGANIZATION",
+    "ORGANIZATIONS",
+    "OUT",
+    "OUTBOUND",
+    "OUTBOUND_FDN",
+    "OUTER",
+    "OUTGOING",
+    "OVER",
+    "OVERWRITE",
+    "OWNER",
+    "OWNERSHIP",
+    "PACREAD",
+    "PACWRITE",
+    "PARAMETER",
+    "PARAMETERS",
+    "PARTIAL",
+    "PARTITION",
+    "PASSWORD",
+    "PERCENT",
+    "PIPE",
+    "PIPES",
+    "PIVOT",
+    "PLAN",
+    "PLUS",
+    "POLICIES",
+    "POLICY",
+    "POOL",
+    "POSITION",
+    "PRECEDING",
+    "PRECISION",
+    "PRESERVE",
+    "PRIMARY",
+    "PRIOR",
+    "PRIVILEGES",
+    "PROCEDURE",
+    "PROCEDURES",
+    "PROVIDER",
+    "PROVIDERS",
+    "PROVISION",
+    "PROVISIONINGREQUESTS",
+    "PROVISIONINGSERVERS",
+    "PUBLIC",
+    "PULL",
+    "PURGE",
+    "PUT",
+    "QUALIFY",
+    "QUERIES",
+    "QUIESCE",
+    "RAISE",
+    "RANGE",
+    "READ",
+    "READS",
+    "REAL",
+    "REBUILD",
+    "RECLUSTER",
+    "RECURSIVE",
+    "REFERENCE",
+    "REFERENCES",
+    "REFERENCE_USAGE",
+    "REFRESH",
+    "REGEXP",
+    "REGIONS",
+    "RELEASE",
+    "RELY",
+    "REMOVE",
+    "RENAME",
+    "REPEAT",
+    "REPEATABLE",
+    "REPLACE",
+    "REPLACE_FAIL",
+    "REPLICA",
+    "REPLICATION",
+    "REQUESTS",
+    "RESET",
+    "RESOURCE",
+    "RESPECT",
+    "RESTART",
+    "RESTORE",
+    "RESTRICT",
+    "RESULT",
+    "RESUME",
+    "RETURN",
+    "RETURNS",
+    "REVERSE",
+    "REVOKE",
+    "REWRITE",
+    "RIGHT",
+    "RLIKE",
+    "RLIKE",
+    "RM",
+    "ROLE",
+    "ROLES",
+    "ROLLBACK",
+    "ROLLUP",
+    "ROW",
+    "ROWS",
+    "RULE",
+    "SAMPLE",
+    "SCHEMA",
+    "SCHEMAS",
+    "SEARCH",
+    "SECURE",
+    "SECURITY",
+    "SEED",
+    "SELECT",
+    "SEMI",
+    "SEQUENCE",
+    "SEQUENCES",
+    "SERVER",
+    "SERVERS",
+    "SESSION",
+    "SESSIONIZE",
+    "SET",
+    "SETLOGLEVEL",
+    "SETS",
+    "SETUP",
+    "SFC",
+    "SHARE",
+    "SHARED",
+    "SHARES",
+    "SHOW",
+    "SHUTDOWN",
+    "SIMPLE",
+    "SMALLINT",
+    "SOME",
+    "SORT",
+    "STAGE",
+    "STAGES",
+    "START",
+    "STARTS",
+    "STATEMENT",
+    "STATISTICS",
+    "STOP",
+    "STORAGE",
+    "STORED",
+    "STREAM",
+    "STREAMS",
+    "STRICT",
+    "STRING",
+    "STRUCT",
+    "SUBSCRIPTION",
+    "SUBSCRIPTIONS",
+    "SUPPORT",
+    "SUSPEND",
+    "SUSPENDED",
+    "SUSPEND_IMMEDIATE",
+    "SWAP",
+    "SWITCH",
+    "SYNC",
+    "SYNONYM",
+    "SYSTEM",
+    "T",
+    "TABLE",
+    "TABLES",
+    "TABLESAMPLE",
+    "TASK",
+    "TASKS",
+    "TEMP",
+    "TEMPLATE",
+    "TEMPORARY",
+    "TERSE",
+    "TEXT",
+    "THEN",
+    "TIME",
+    "TIMECOLUMN",
+    "TIMEOUT",
+    "TIMESTAMP",
+    "TIMESTAMPLTZ",
+    "TIMESTAMPNTZ",
+    "TIMESTAMPTZ",
+    "TIMESTAMP_LTZ",
+    "TIMESTAMP_NTZ",
+    "TIMESTAMP_TZ",
+    "TINYINT",
+    "TO",
+    "TOP",
+    "TRANSACTION",
+    "TRANSACTIONS",
+    "TRANSIENT",
+    "TRIGGER",
+    "TRIGGERS",
+    "TRUE",
+    "TRUNCATE",
+    "TRY_CAST",
+    "TS",
+    "TYPE",
+    "TYPES",
+    "UNBOUNDED",
+    "UNDO",
+    "UNDROP",
+    "UNION",
+    "UNIQUE",
+    "UNLOCK",
+    "UNPIVOT",
+    "UNSET",
+    "UNTIL",
+    "UPDATE",
+    "UPGRADE",
+    "USAGE",
+    "USE",
+    "USER",
+    "USERS",
+    "USE_ANY_ROLE",
+    "USING",
+    "UTC",
+    "UTCTIMESTAMP",
+    "VALIDATE",
+    "VALUES",
+    "VARBINARY",
+    "VARCHAR",
+    "VARIABLES",
+    "VARIANT",
+    "VARYING",
+    "VERSION",
+    "VIEW",
+    "VIEWS",
+    "VOLATILE",
+    "VOLUME",
+    "VOLUMES",
+    "WAREHOUSE",
+    "WAREHOUSES",
+    "WARN",
+    "WHEN",
+    "WHENEVER",
+    "WHERE",
+    "WHILE",
+    "WITH",
+    "WITHIN",
+    "WITHOUT",
+    "WMSQUEUEDOPS",
+    "WORK",
+    "WRITE",
+    "ZONE",
+)

snowflake/cli/_plugins/sql/lexer/lexer.py

@@ -0,0 +1,56 @@
+import re
+
+from pygments.lexer import bygroups, include
+from pygments.lexers.sql import RegexLexer
+from pygments.token import (
+    Comment,
+    Keyword,
+    Name,
+    Number,
+    Operator,
+    Punctuation,
+    String,
+    Text,
+)
+from snowflake.cli._plugins.sql.lexer.functions import FUNCTIONS
+from snowflake.cli._plugins.sql.lexer.keywords import KEYWORDS
+from snowflake.cli._plugins.sql.lexer.types import TYPES
+
+
+class CliLexer(RegexLexer):
+    name = "Snowflake-CLI"
+    aliases = ("cli", "snowflake-cli")
+    mimetype = ("text/x-snowflake-cli",)
+    flags = re.IGNORECASE
+    tokens = {
+        "comments": [
+            (r"--.*?$", Comment.Single),
+            (r"/\*", Comment.Multiline, "multiline-comments"),
+        ],
+        "multiline-comments": [
+            (r"\*/", Comment.Multiline, "#pop"),
+            (r"[^/*]+", Comment.Multiline),
+            (r"[/*]", Comment.Multiline),
+        ],
+        "root": [
+            include("comments"),
+            (r"\s+", Text),
+            (r"[0-9]+", Number.Integer),
+            (r"[0-9]*\.[0-9]+(e[+-][0-9]+)", Number.Float),
+            (r"'(\\\\|\\'|''|[^'])*'", String.Single),
+            (r'"(\\\\|\\"|""|[^"])*"', String.Double),
+            (r"`(\\\\|\\`|``|[^`])*`", String.Symbol),
+            (r"[+*/<>=~!@#%^&$|`?-]", Operator),
+            (
+                r"\b({0})(\b\s*)(\()?".format("|".join(TYPES)),
+                bygroups(Keyword.Type, Text, Punctuation),
+            ),
+            (r"\b({0})\b".format("|".join(KEYWORDS + FUNCTIONS)), Keyword),
+            (r"\b(auto_increment|engine|charset|tables)\b", Keyword.Pseudo),
+            (r"(true|false|null)", Name.Constant),
+            (r"([a-z_]\w*)(\s*)(\()", bygroups(Name.Function, Text, Punctuation)),
+            (r"[a-z_]\w*", Name),
+            (r"\$[a-z0-9]*[._]*[a-z0-9]*", Name.Variable),
+            (r"[;:()\[\],.]", Punctuation),
+        ],
+    }

snowflake/cli/_plugins/sql/lexer/types.py

@@ -0,0 +1,37 @@
+TYPES = (
+    "ARRAY",
+    "BIGINT",
+    "BINARY",
+    "BIT",
+    "BOOLEAN",
+    "BYTEINT",
+    "CHAR",
+    "CHARACTER",
+    "DATE",
+    "DATETIME",
+    "DEC",
+    "DECIMAL",
+    "DOUBLE",
+    "FLOAT",
+    "INT",
+    "INTEGER",
+    "NUMBER",
+    "NUMERIC",
+    "OBJECT",
+    "REAL",
+    "SMALLINT",
+    "STRING",
+    "TEXT",
+    "TIME",
+    "TIMESTAMP",
+    "TIMESTAMPLTZ",
+    "TIMESTAMPNTZ",
+    "TIMESTAMPTZ",
+    "TIMESTAMP_LTZ",
+    "TIMESTAMP_NTZ",
+    "TIMESTAMP_TZ",
+    "TINYINT",
+    "VARBINARY",
+    "VARCHAR",
+    "VARIANT",
+)

snowflake/cli/_plugins/sql/manager.py

@@ -20,20 +20,24 @@ from functools import partial
 from pathlib import Path
 from typing import Dict, Iterable, List, Tuple
 
+from snowflake.cli._app.printing import print_result
 from snowflake.cli._plugins.sql.snowsql_templating import transpile_snowsql_templates
-from snowflake.cli._plugins.sql.source_reader import (
+from snowflake.cli._plugins.sql.statement_reader import (
+    CompiledStatement,
     compile_statements,
     files_reader,
     query_reader,
 )
+from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.exceptions import CliArgumentError, CliSqlError
+from snowflake.cli.api.output.types import CollectionResult
 from snowflake.cli.api.rendering.sql_templates import snowflake_sql_jinja_render
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.cli.api.sql_execution import SqlExecutionMixin, VerboseCursor
 from snowflake.connector.cursor import SnowflakeCursor
 
-IsSingleStatement = bool
+ExpectedResultsCount = int
 
 logger = logging.getLogger(__name__)
 
@@ -46,7 +50,8 @@ class SqlManager(SqlExecutionMixin):
         std_in: bool,
         data: Dict | None = None,
         retain_comments: bool = False,
-    ) -> Tuple[IsSingleStatement, Iterable[SnowflakeCursor]]:
+        single_transaction: bool = False,
+    ) -> Tuple[ExpectedResultsCount, Iterable[SnowflakeCursor]]:
         """Reads, transforms and execute statements from input.
 
         Only one input can be consumed at a time.
@@ -71,8 +76,10 @@ class SqlManager(SqlExecutionMixin):
         else:
             raise CliArgumentError("Use either query, filename or input option.")
 
-        errors, stmt_count, compiled_statements = compile_statements(stmt_reader)
-        if not any((errors, stmt_count, compiled_statements)):
+        errors, expected_results_cnt, compiled_statements = compile_statements(
+            stmt_reader
+        )
+        if not any((errors, expected_results_cnt, compiled_statements)):
             raise CliArgumentError("Use either query, filename or input option.")
 
         if errors:
@@ -81,8 +88,35 @@ class SqlManager(SqlExecutionMixin):
                 cli_console.warning(error)
             raise CliSqlError("SQL rendering error")
 
-        is_single_statement = not (stmt_count > 1)
-        return is_single_statement, self.execute_string(
-            "\n".join(compiled_statements),
-            cursor_class=VerboseCursor,
+        if single_transaction:
+            logger.info("disabling AUTOCOMMIT")
+            self.disable_autocommit()
+            compiled_statements = [
+                CompiledStatement(statement="BEGIN;"),
+                *compiled_statements,
+                CompiledStatement(statement="COMMIT;"),
+            ]
+            expected_results_cnt = len(compiled_statements)
+
+        cursor_class = SnowflakeCursor if get_cli_context().is_repl else VerboseCursor
+        return expected_results_cnt, self._execute_compiled_statements(
+            compiled_statements,
+            cursor_class=cursor_class,
         )
+
+    def _execute_compiled_statements(
+        self, compiled_statements: List[CompiledStatement], cursor_class
+    ) -> Iterable[SnowflakeCursor]:
+        for stmt in compiled_statements:
+            if stmt.execute_async:
+                cursor = self._conn.cursor(cursor_class=cursor_class)
+                cursor.execute(stmt.statement, _no_results=True)
+                # only log query ID for consistency with SnowSQL
+                logger.info("Async execution id: %s", cursor.sfqid)
+                print_result(CollectionResult([{"scheduled query ID": cursor.sfqid}]))
+            elif stmt.statement:
+                yield from self.execute_string(
+                    stmt.statement, cursor_class=cursor_class
+                )
+            if stmt.command:
+                stmt.command.execute(self._conn)