Handle CHECK and UNIQUE during column population for tables

jzmiller1 · Aug 15, 2024 · 9b16341 · 9b16341
1 parent 6904b27
commit 9b16341
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 1 deletion.
diff --git a/postnormalism/schema/table.py b/postnormalism/schema/table.py
@@ -51,9 +51,18 @@ def _extract_columns(self):
             if in_table_definition:
                 parts = line.split(',')
                 for part in parts:
-                    match = re.match(self._pattern_create, part.strip())
+                    part = part.strip()
+
+                    if part.upper().startswith("UNIQUE") or part.upper().startswith("CHECK"):
+                        continue
+                    match = re.match(self._pattern_create, part)
                     if match:
                         columns.append(match.group(1))
+                    else:
+                        column_and_constraint = re.match(r"^\s*(\w+)\s+.*(?:UNIQUE|CHECK)\s*\(.*\)", part,
+                                                         re.IGNORECASE)
+                        if column_and_constraint:
+                            columns.append(column_and_constraint.group(1))
 
         if self.alter:
             for line in self.alter.splitlines():

diff --git a/tests/items/test_table.py b/tests/items/test_table.py
@@ -217,3 +217,23 @@ def test_inherited_table_columns(self):
             "id", "created_at", "updated_at", "element"
         ]
         self.assertEqual(universe_db.process_element_material.columns, expected_columns)
+
+    def test_exclude_check_and_unique_constraints(self):
+        create_statement = """
+        CREATE TABLE process_action (
+            id uuid PRIMARY KEY DEFAULT uuid_generate_v4(),
+            process uuid REFERENCES process NOT NULL,
+            target uuid NOT NULL,
+            calculated uuid REFERENCES attribute NOT NULL,
+            description varchar(240),
+            action_type varchar(10) NOT NULL,
+            CHECK (action_type = 'linear' OR action_type = 'inplace'),
+            UNIQUE (process, target, calculated)
+        );
+        """
+
+        expected_columns = [
+            'id', 'process', 'target', 'calculated', 'description', 'action_type'
+        ]
+        process_action_table = Table(create=create_statement)
+        self.assertEqual(process_action_table.columns, expected_columns)