Merge pull request #242 from kliushnichenko/main

Support functions with schema prefix in `DEFAULT` and `CHECK` statements.
xnuinside · Mar 25, 2024 · 365e520 · 365e520
2 parents f640799 + 82f2cf2
commit 365e520
Show file tree

Hide file tree

Showing 5 changed files with 175 additions and 60 deletions.
diff --git a/CHANGELOG.txt b/CHANGELOG.txt
@@ -1,3 +1,7 @@
+**v1.0.4**
+### Improvements
+1. Support functions with schema prefix in `DEFAULT` and `CHECK` statements.
+
 **v1.0.3**
 ### Improvements
 1. Fixed bug with `CREATE OR REPLACE SCHEMA`.

diff --git a/simple_ddl_parser/dialects/sql.py b/simple_ddl_parser/dialects/sql.py
@@ -1429,6 +1429,7 @@ def p_default(self, p: List) -> None:
         | DEFAULT LP pid RP
         | DEFAULT LP funct_expr pid RP
         | default id
+        | DEFAULT id DOT funct_expr
         | default LP RP
         """
         p_list = remove_par(list(p))
@@ -1448,8 +1449,11 @@ def p_default(self, p: List) -> None:
 
     @staticmethod
     def pre_process_default(p_list: List) -> Any:
-        if len(p_list) == 5 and isinstance(p_list[3], list):
-            default = p_list[3][0]
+        if len(p_list) == 5:
+            if isinstance(p_list[3], list):
+                default = p_list[3][0]
+            else:
+                default = f"{''.join(p_list[2:5])}"
         elif "DEFAULT" in p_list and len(p_list) == 4:
             default = f"{p_list[2]} {p_list[3]}"
         else:
@@ -1523,20 +1527,33 @@ def p_check_st(self, p: List) -> None:
         | check_st id RP
         | check_st STRING RP
         | check_st funct_args
+        | CHECK LP id DOT id
         | check_st LP pid RP
         """
         p_list = remove_par(list(p))
         if isinstance(p[1], dict):
             p[0] = p[1]
         else:
             p[0] = {"check": []}
-        for item in p_list[2:]:
+
+        i = 0
+        items = p_list[2:]
+        items_num = len(items)
+
+        while i < items_num:
+            item = items[i]
+            # handle <schema>.<function>
+            if i + 1 < items_num and items[i + 1] == ".":
+                p[0]["check"].append(f"{''.join(items[i:i + 3])}")
+                i += 3
+                continue
             if isinstance(p_list[-1], dict) and p_list[-1].get("args"):
                 p[0]["check"][-1] += p_list[-1]["args"]
             elif isinstance(item, list):
                 p[0]["check"].append(f"({','.join(item)})")
             else:
                 p[0]["check"].append(item)
+            i += 1
 
     def p_using_tablespace(self, p: List) -> None:
         """using_tablespace : USING INDEX tablespace"""

diff --git a/simple_ddl_parser/parsetab.py b/simple_ddl_parser/parsetab.py
diff --git a/tests/test_checks.py b/tests/test_checks.py
@@ -165,3 +165,53 @@ def test_checks_with_in_works():
     }
 
     assert result == expected
+
+
+def test_check_function_with_schema():
+    ddl = """
+    CREATE TABLE foo
+    (
+    entity_id        UUID PRIMARY KEY DEFAULT public.getId(),
+    name             TEXT,
+    CONSTRAINT my_constraint  CHECK(v2.my_function(name) IS TRUE)
+    );
+    """
+
+    result = DDLParser(ddl).run(group_by_type=True)
+    expected = {
+        'tables': [
+            {'alter': {},
+             'checks': [{'constraint_name': 'my_constraint', 'statement': 'v2.my_function(name) IS TRUE'}],
+             'columns': [{'check': None,
+                          'default': 'public.getId()',
+                          'name': 'entity_id',
+                          'nullable': False,
+                          'references': None,
+                          'size': None,
+                          'type': 'UUID',
+                          'unique': False},
+                         {'check': None,
+                          'default': None,
+                          'name': 'name',
+                          'nullable': True,
+                          'references': None,
+                          'size': None,
+                          'type': 'TEXT',
+                          'unique': False}],
+             'constraints': {'checks': [{'constraint_name': 'my_constraint',
+                                         'statement': 'v2.my_function(name) IS '
+                                                      'TRUE'}]},
+             'index': [],
+             'partitioned_by': [],
+             'primary_key': ['entity_id'],
+             'schema': None,
+             'table_name': 'foo',
+             'tablespace': None}],
+        'types': [],
+        'ddl_properties': [],
+        'domains': [],
+        'schemas': [],
+        'sequences': [],
+    }
+
+    assert result == expected
diff --git a/tests/test_simple_ddl_parser.py b/tests/test_simple_ddl_parser.py
@@ -1184,6 +1184,48 @@ def test_default_expression():
     }
     assert expected == result
 
+def test_default_function_with_schema():
+    ddl = """
+    CREATE TABLE foo
+    (
+        entity_id        UUID PRIMARY KEY DEFAULT public.getId()
+    );
+    """
+
+    result = DDLParser(ddl).run(group_by_type=True)
+    expected = {
+        "tables": [
+            {
+                "columns": [
+                    {
+                        "name": "entity_id",
+                        "type": "UUID",
+                        "size": None,
+                        "references": None,
+                        "unique": False,
+                        "nullable": False,
+                        "default": "public.getId()",
+                        "check": None,
+                    }
+                ],
+                "primary_key": ["entity_id"],
+                "alter": {},
+                "checks": [],
+                "index": [],
+                "partitioned_by": [],
+                "tablespace": None,
+                "schema": None,
+                "table_name": "foo",
+            }
+        ],
+        "types": [],
+        "ddl_properties": [],
+        "sequences": [],
+        "domains": [],
+        "schemas": [],
+    }
+    assert result == expected
+
 
 def test_comments_in_columns():
     ddl = """