Merge branch 'comment-on-support' into 'main'

Add support for `COMMENT ON column IS` statements + fix some issues with pgdump format

See merge request !1

simple_ddl_parser/ddl_parser.py

@@ -97,6 +97,11 @@ class DDLParser(
         elif t.type == "CHECK":
             self.lexer.check = True
 
+    def t_TIMESTAMP_TIMEZONE(self, t: LexToken):
+        r"with\ time\ zone"
+        t.type = "TIMESTAMP_TIMEZONE"
+        return self.set_last_token(t)
+
     def t_DOT(self, t: LexToken) -> LexToken:
         r"\."
         t.type = "DOT"
@@ -116,13 +121,14 @@ class DDLParser(
         """many of reserved words can be used as column name,
         to decide is it a column name or not we need do some checks"""
         skip_id_tokens = ["(", ")", ","]
+        key_as_column_name = "key"
         return (
             t.value not in skip_id_tokens
             and self.lexer.is_table
             and self.lexer.lp_open
             and not self.lexer.is_like
             and (self.lexer.last_token == "COMMA" or self.lexer.last_token == "LP")
-            and t.value.upper() not in tok.first_liners
+            and (t.value.upper() not in tok.first_liners or t.value == key_as_column_name)
         )
 
     def is_creation_name(self, t: LexToken) -> bool:
@@ -173,6 +179,11 @@ class DDLParser(
         else:
             t = self.tokens_not_columns_names(t)
 
+        if self.lexer.is_comment_on:
+            _type = tok.comment_on_tokens.get(t.value)
+            if _type:
+                t.type = _type
+
         self.capitalize_tokens(t)
         self.commat_type(t)
 
@@ -199,6 +210,8 @@ class DDLParser(
 
         if t.type == "ALTER":
             self.lexer.is_alter = True
+        if t.type == "COMMENT":
+            self.lexer.is_comment_on = True
         if t.type == "LIKE":
             self.lexer.is_like = True
         elif t.type in ["TYPE", "DOMAIN", "TABLESPACE"]:

simple_ddl_parser/dialects/sql.py

@@ -193,6 +193,7 @@ class Column:
         | id id
         | id id id id
         | id id id
+        | id TIMESTAMP_TIMEZONE
         | id DOT id
         | tid
         | ARRAY
@@ -1587,6 +1588,14 @@ class BaseSQL(
         p_list = remove_par(list(p))
         p[0] = {"comment": check_spec(p_list[-1])}
 
+    def p_comment_on(self, p: List) -> None:
+        """comment_on : COMMENT ON COLUMN id DOT id DOT id id STRING"""
+        p[0] = {"comment_stmt": p[10], "comment_colref": ".".join((p[4], p[6], p[8]))}
+
+    def p_comment_on_expr(self, p: List) -> None:
+        """expr : comment_on"""
+        p[0] = p[1]
+
     def p_tablespace(self, p: List) -> None:
         """tablespace : TABLESPACE id
         | TABLESPACE id properties

simple_ddl_parser/output/common.py

@@ -181,6 +181,8 @@ def result_format(
             tables_dict = process_alter_and_index_result(
                 tables_dict, table, output_mode
             )
+        elif "comment_colref" in table:
+            process_comment_on_result(tables_dict, table)
         else:
             # process tables, types, sequence and etc. data
             table_data = process_entities(tables_dict, table, output_mode)
@@ -253,6 +255,13 @@ def set_unique_columns(table_data: Dict) -> Dict:
         del table_data["unique_statement"]
     return table_data
 
+def process_comment_on_result(tables_dict: Dict, statement: Dict):
+    schema, table_name, column_name = statement.get("comment_colref").split(".")
+    table = get_table_from_tables_data(tables_dict, (table_name, schema))
+    for column in table.get("columns"):
+        if column.get("name") == column_name:
+            column["comment"] = statement.get("comment_stmt")
+
 
 def group_by_type_result(final_result: List[Dict]) -> Dict[str, List]:
     result_as_dict = {

simple_ddl_parser/parser.py

@@ -313,6 +313,7 @@ class Parser:
             "last_par",
             "lp_open",
             "is_alter",
+            "is_comment_on",
             "is_like",
         ]
         for attr in attrs:

simple_ddl_parser/tokens.py

@@ -13,6 +13,7 @@ defenition_statements = {
     "CLUSTERED": "CLUSTERED",
     "SEQUENCE": "SEQUENCE",
     "TABLESPACE": "TABLESPACE",
+    "COMMENT": "COMMENT",
 }
 common_statements = {
     "INDEX": "INDEX",
@@ -100,14 +101,18 @@ sequence_reserved = {
 }
 
 
+comment_on_tokens = {"COLUMN"}
+comment_on_tokens = {value: value for value in comment_on_tokens}
+
 tokens = tuple(
     set(
-        ["ID", "DOT", "STRING", "DQ_STRING", "LP", "RP", "LT", "RT", "COMMAT", "AUTOINCREMENT"]
+        ["ID", "DOT", "STRING", "DQ_STRING", "LP", "RP", "LT", "RT", "COMMAT", "AUTOINCREMENT", "TIMESTAMP_TIMEZONE"]
         + list(defenition_statements.values())
         + list(common_statements.values())
         + list(columns_defenition.values())
         + list(sequence_reserved.values())
         + list(after_columns_tokens.values())
+        + list(comment_on_tokens.values())
     )
 )
 

tests/dialects/test_psql.py

+from simple_ddl_parser import DDLParser
+
+
+def test_timestamp_with_time_zone():
+    ddl = """
+    CREATE TABLE test (
+        id int NOT NULL,
+        testtime timestamp with time zone NOT NULL
+    );
+    """
+
+    expected = [
+        {
+            "alter": {},
+            "checks": [],
+            "table_name": "test",
+            "tablespace": None,
+            "primary_key": [],
+            "index": [],
+            "schema": None,
+            "partitioned_by": [],
+            "columns": [
+                {
+                    "name": "id",
+                    "type": "int",
+                    "size": None,
+                    "references": None,
+                    "unique": False,
+                    "nullable": False,
+                    "default": None,
+                    "check": None,
+                },
+                {
+                    "name": "testtime",
+                    "type": "timestamp with time zone",
+                    "size": None,
+                    "references": None,
+                    "unique": False,
+                    "nullable": False,
+                    "default": None,
+                    "check": None,
+                },
+            ]
+        }
+    ]
+
+    assert DDLParser(ddl).run() == expected
+
+
+def test_column_named_key():
+    ddl = """
+    CREATE TABLE test (
+        id int NOT NULL,
+        key int NOT NULL
+    );
+    """
+
+    expected = [
+        {
+            "alter": {},
+            "checks": [],
+            "table_name": "test",
+            "tablespace": None,
+            "primary_key": [],
+            "index": [],
+            "schema": None,
+            "partitioned_by": [],
+            "columns": [
+                {
+                    "name": "id",
+                    "type": "int",
+                    "size": None,
+                    "references": None,
+                    "unique": False,
+                    "nullable": False,
+                    "default": None,
+                    "check": None,
+                },
+                {
+                    "name": "key",
+                    "type": "int",
+                    "size": None,
+                    "references": None,
+                    "unique": False,
+                    "nullable": False,
+                    "default": None,
+                    "check": None,
+                },
+            ]
+        }
+    ]
\ No newline at end of file

tests/test_comment_on_statements.py

+from simple_ddl_parser import DDLParser
+
+
+def test_comment_on():
+    ddl = """
+    CREATE schema test;
+    CREATE TABLE test.some_table (
+        some_column int
+    );
+    COMMENT ON COLUMN test.some_table.some_column IS 'WOW';
+    """
+
+    expected = [
+        {"schema_name": "test"},
+        {
+            "table_name": "some_table",
+            "schema": "test",
+            "primary_key": [],
+            "columns": [
+                {
+                    "name": "some_column",
+                    "type": "int",
+                    "size": None,
+                    "references": None,
+                    "unique": False,
+                    "nullable": True,
+                    "default": None,
+                    "check": None,
+                    "comment": "'WOW'",
+                }
+            ],
+            "alter": {},
+            "checks": [],
+            "index": [],
+            "partitioned_by": [],
+            "tablespace": None,
+        }
+    ]
+
+    assert DDLParser(ddl).run() == expected