feat: support create <object> clone, close #313

tconbeer · tconbeer · Jan 13, 2023 · Jan 13, 2023 · Jan 13, 2023 · Jan 13, 2023
commit 03512242d8653bf6a29eb2e95855116261d01027
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+### Formatting Changes + Bug Fixes
+
+-   sqlfmt now supports `create <object> ... clone` statements ([#313](https://github.com/tconbeer/sqlfmt/issues/313)).
+
 ### Features
 
 -   by default, sqlfmt now runs an additional safety check that parses the formatted output to ensure it contains all of the same content as the raw input. This incurs a slight (~20%) performance penalty. To bypass this safety check, you can use the command line option `--fast`, the corresponding TOML or environment variable config, or pass `Mode(fast=True)` to any API method. The safety check is automatically bypassed if sqlfmt is run with the `--check` or `--diff` options. If the safety check fails, the CLI will include an error in the report, and the `format_string` API will raise a `SqlfmtEquivalenceError`, which is a subclass of `SqlfmtError`.

diff --git a/src/sqlfmt/rules/__init__.py b/src/sqlfmt/rules/__init__.py
@@ -2,9 +2,11 @@
 
 from sqlfmt import actions
 from sqlfmt.rule import Rule
+from sqlfmt.rules.clone import CLONE as CLONE
 from sqlfmt.rules.common import (
     ALTER_DROP_FUNCTION,
     ALTER_WAREHOUSE,
+    CREATE_CLONABLE,
     CREATE_FUNCTION,
     CREATE_WAREHOUSE,
     group,
@@ -195,6 +197,18 @@
             action=partial(actions.lex_ruleset, new_ruleset=GRANT),
         ),
     ),
+    Rule(
+        name="create_clone",
+        priority=2015,
+        pattern=group(CREATE_CLONABLE + r"\s+\w+\s+clone") + group(r"\W", r"$"),
+        action=partial(
+            actions.handle_nonreserved_keyword,
+            action=partial(
+                actions.lex_ruleset,
+                new_ruleset=CLONE,
+            ),
+        ),
+    ),
     Rule(
         name="create_function",
         priority=2020,

diff --git a/src/sqlfmt/rules/clone.py b/src/sqlfmt/rules/clone.py
@@ -0,0 +1,27 @@
+from functools import partial
+
+from sqlfmt import actions
+from sqlfmt.rule import Rule
+from sqlfmt.rules.common import CREATE_CLONABLE, group
+from sqlfmt.rules.core import CORE
+from sqlfmt.token import TokenType
+
+CLONE = [
+    *CORE,
+    Rule(
+        name="unterm_keyword",
+        priority=1300,
+        pattern=group(CREATE_CLONABLE, r"clone") + group(r"\W", r"$"),
+        action=partial(actions.add_node_to_buffer, token_type=TokenType.UNTERM_KEYWORD),
+    ),
+    Rule(
+        name="word_operator",
+        priority=1500,
+        pattern=group(
+            r"at",
+            r"before",
+        )
+        + group(r"\W", r"$"),
+        action=partial(actions.add_node_to_buffer, token_type=TokenType.WORD_OPERATOR),
+    ),
+]
diff --git a/src/sqlfmt/rules/common.py b/src/sqlfmt/rules/common.py
@@ -37,3 +37,18 @@ def group(*choices: str) -> str:
 
 CREATE_WAREHOUSE = r"create(\s+or\s+replace)?\s+warehouse(\s+if\s+not\s+exists)?"
 ALTER_WAREHOUSE = r"alter\s+warehouse(\s+if\s+exists)?"
+
+CREATE_CLONABLE = (
+    r"create(\s+or\s+replace)?\s+"
+    + group(
+        r"database",
+        r"schema",
+        r"table",
+        r"stage",
+        r"file\s+format",
+        r"sequence",
+        r"stream",
+        r"task",
+    )
+    + r"(\s+if\s+not\s+exists)?"
+)
diff --git a/tests/data/unformatted/411_create_clone.sql b/tests/data/unformatted/411_create_clone.sql
@@ -0,0 +1,19 @@
+create database mytestdb_clone clone mytestdb;
+create schema mytestschema_clone_restore clone testschema before (timestamp => to_timestamp(40*365*86400));
+create table orders_clone_restore clone orders at (timestamp => to_timestamp_tz('04/05/2013 01:02:03', 'mm/dd/yyyy hh24:mi:ss'));
+CREATE TABLE ORDERS_CLONE_RESTORE clone orders before (statement => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726');
+)))))__SQLFMT_OUTPUT__(((((
+create database mytestdb_clone
+clone mytestdb
+;
+create schema mytestschema_clone_restore
+clone testschema before (timestamp => to_timestamp(40 * 365 * 86400))
+;
+create table orders_clone_restore
+clone
+    orders
+    at (timestamp => to_timestamp_tz('04/05/2013 01:02:03', 'mm/dd/yyyy hh24:mi:ss'))
+;
+create table orders_clone_restore
+clone orders before (statement => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726')
+;
diff --git a/tests/functional_tests/test_general_formatting.py b/tests/functional_tests/test_general_formatting.py
@@ -69,6 +69,7 @@
         "unformatted/408_alter_function_snowflake_examples.sql",
         "unformatted/409_create_external_function.sql",
         "unformatted/410_create_warehouse.sql",
+        "unformatted/411_create_clone.sql",
         "unformatted/999_unsupported_ddl.sql",
     ],
 )

diff --git a/tests/unit_tests/test_rule.py b/tests/unit_tests/test_rule.py
@@ -5,9 +5,9 @@
 import pytest
 
 from sqlfmt.rule import Rule
-from sqlfmt.rules import CORE, FUNCTION, GRANT, JINJA, MAIN, WAREHOUSE
+from sqlfmt.rules import CLONE, CORE, FUNCTION, GRANT, JINJA, MAIN, WAREHOUSE
 
-ALL_RULESETS = [CORE, FUNCTION, GRANT, JINJA, MAIN, WAREHOUSE]
+ALL_RULESETS = [CLONE, CORE, FUNCTION, GRANT, JINJA, MAIN, WAREHOUSE]
 
 
 def get_rule(ruleset: List[Rule], rule_name: str) -> Rule:
@@ -311,6 +311,17 @@ def get_rule(ruleset: List[Rule], rule_name: str) -> Rule:
         (WAREHOUSE, "unterm_keyword", "set tag"),
         (WAREHOUSE, "unterm_keyword", "resume if suspended"),
         (WAREHOUSE, "unterm_keyword", "unset scaling_policy"),
+        (MAIN, "create_clone", "create table foo clone"),
+        (MAIN, "create_clone", "create or replace database foo clone"),
+        (MAIN, "create_clone", "create stage if not exists foo clone"),
+        (CLONE, "unterm_keyword", "create table"),
+        (CLONE, "unterm_keyword", "create database"),
+        (CLONE, "unterm_keyword", "create schema"),
+        (CLONE, "unterm_keyword", "create\nfile\nformat"),
+        (CLONE, "unterm_keyword", "clone"),
+        (CLONE, "name", "foo"),
+        (CLONE, "word_operator", "at"),
+        (CLONE, "word_operator", "before"),
     ],
 )
 def test_regex_exact_match(
@@ -346,6 +357,7 @@ def test_regex_exact_match(
         (MAIN, "unterm_keyword", "delete"),
         (MAIN, "unsupported_ddl", "insert('abc', 1, 2, 'Z')"),
         (MAIN, "unsupported_ddl", "get(foo, 'bar')"),
+        (MAIN, "create_clone", "create table"),
         (JINJA, "jinja_set_block_start", "{% set foo = 'baz' %}"),
         (GRANT, "unterm_keyword", "select"),
         (FUNCTION, "unterm_keyword", "secure"),