Skip to content
Snippets Groups Projects
Commit 939d5e09 authored by Felipe de Oliveira Hargreaves's avatar Felipe de Oliveira Hargreaves
Browse files

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
parents d8d653bc 48c4d14c
Branches main
No related tags found
1 merge request!1Add support for `COMMENT ON column IS` statements + fix some issues with pgdump format
...@@ -97,6 +97,11 @@ class DDLParser( ...@@ -97,6 +97,11 @@ class DDLParser(
elif t.type == "CHECK": elif t.type == "CHECK":
self.lexer.check = True 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: def t_DOT(self, t: LexToken) -> LexToken:
r"\." r"\."
t.type = "DOT" t.type = "DOT"
...@@ -116,13 +121,14 @@ class DDLParser( ...@@ -116,13 +121,14 @@ class DDLParser(
"""many of reserved words can be used as column name, """many of reserved words can be used as column name,
to decide is it a column name or not we need do some checks""" to decide is it a column name or not we need do some checks"""
skip_id_tokens = ["(", ")", ","] skip_id_tokens = ["(", ")", ","]
key_as_column_name = "key"
return ( return (
t.value not in skip_id_tokens t.value not in skip_id_tokens
and self.lexer.is_table and self.lexer.is_table
and self.lexer.lp_open and self.lexer.lp_open
and not self.lexer.is_like and not self.lexer.is_like
and (self.lexer.last_token == "COMMA" or self.lexer.last_token == "LP") 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: def is_creation_name(self, t: LexToken) -> bool:
...@@ -173,6 +179,11 @@ class DDLParser( ...@@ -173,6 +179,11 @@ class DDLParser(
else: else:
t = self.tokens_not_columns_names(t) 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.capitalize_tokens(t)
self.commat_type(t) self.commat_type(t)
...@@ -199,6 +210,8 @@ class DDLParser( ...@@ -199,6 +210,8 @@ class DDLParser(
if t.type == "ALTER": if t.type == "ALTER":
self.lexer.is_alter = True self.lexer.is_alter = True
if t.type == "COMMENT":
self.lexer.is_comment_on = True
if t.type == "LIKE": if t.type == "LIKE":
self.lexer.is_like = True self.lexer.is_like = True
elif t.type in ["TYPE", "DOMAIN", "TABLESPACE"]: elif t.type in ["TYPE", "DOMAIN", "TABLESPACE"]:
......
...@@ -193,6 +193,7 @@ class Column: ...@@ -193,6 +193,7 @@ class Column:
| id id | id id
| id id id id | id id id id
| id id id | id id id
| id TIMESTAMP_TIMEZONE
| id DOT id | id DOT id
| tid | tid
| ARRAY | ARRAY
...@@ -1587,6 +1588,14 @@ class BaseSQL( ...@@ -1587,6 +1588,14 @@ class BaseSQL(
p_list = remove_par(list(p)) p_list = remove_par(list(p))
p[0] = {"comment": check_spec(p_list[-1])} 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: def p_tablespace(self, p: List) -> None:
"""tablespace : TABLESPACE id """tablespace : TABLESPACE id
| TABLESPACE id properties | TABLESPACE id properties
......
...@@ -181,6 +181,8 @@ def result_format( ...@@ -181,6 +181,8 @@ def result_format(
tables_dict = process_alter_and_index_result( tables_dict = process_alter_and_index_result(
tables_dict, table, output_mode tables_dict, table, output_mode
) )
elif "comment_colref" in table:
process_comment_on_result(tables_dict, table)
else: else:
# process tables, types, sequence and etc. data # process tables, types, sequence and etc. data
table_data = process_entities(tables_dict, table, output_mode) table_data = process_entities(tables_dict, table, output_mode)
...@@ -253,6 +255,13 @@ def set_unique_columns(table_data: Dict) -> Dict: ...@@ -253,6 +255,13 @@ def set_unique_columns(table_data: Dict) -> Dict:
del table_data["unique_statement"] del table_data["unique_statement"]
return table_data 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]: def group_by_type_result(final_result: List[Dict]) -> Dict[str, List]:
result_as_dict = { result_as_dict = {
......
...@@ -313,6 +313,7 @@ class Parser: ...@@ -313,6 +313,7 @@ class Parser:
"last_par", "last_par",
"lp_open", "lp_open",
"is_alter", "is_alter",
"is_comment_on",
"is_like", "is_like",
] ]
for attr in attrs: for attr in attrs:
......
Source diff could not be displayed: it is too large. Options to address this: view the blob.
...@@ -13,6 +13,7 @@ defenition_statements = { ...@@ -13,6 +13,7 @@ defenition_statements = {
"CLUSTERED": "CLUSTERED", "CLUSTERED": "CLUSTERED",
"SEQUENCE": "SEQUENCE", "SEQUENCE": "SEQUENCE",
"TABLESPACE": "TABLESPACE", "TABLESPACE": "TABLESPACE",
"COMMENT": "COMMENT",
} }
common_statements = { common_statements = {
"INDEX": "INDEX", "INDEX": "INDEX",
...@@ -100,14 +101,18 @@ sequence_reserved = { ...@@ -100,14 +101,18 @@ sequence_reserved = {
} }
comment_on_tokens = {"COLUMN"}
comment_on_tokens = {value: value for value in comment_on_tokens}
tokens = tuple( tokens = tuple(
set( 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(defenition_statements.values())
+ list(common_statements.values()) + list(common_statements.values())
+ list(columns_defenition.values()) + list(columns_defenition.values())
+ list(sequence_reserved.values()) + list(sequence_reserved.values())
+ list(after_columns_tokens.values()) + list(after_columns_tokens.values())
+ list(comment_on_tokens.values())
) )
) )
......
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
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
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment