add support for set with = and ; in them

Author: xnuinside

simple_ddl_parser/dialects/sql.py

@@ -162,6 +162,8 @@ def parse_complex_type(p_list: List[str]) -> str:
     def p_c_type(self, p: List) -> None:
         """c_type : id
         | id id
+        | id id id id
+        | id id id
         | id DOT id
         | tid
         | ARRAY

simple_ddl_parser/output/dialects.py

@@ -49,9 +49,6 @@ def update_bigquery_output(table_data: Dict) -> Dict:
     if table_data.get("schema"):
         table_data["dataset"] = table_data["schema"]
         del table_data["schema"]
-    else:
-        if not table_data.get("schema_name"):
-            table_data["dataset"] = None
     return table_data
 
 

simple_ddl_parser/parser.py

@@ -106,24 +106,29 @@ def pre_process_data(self, data):
     @staticmethod
     def process_set(tables: List, set_line: str) -> None:
         set_line = set_line.split()
-        name = set_line[-2]
-        value = set_line[-1]
+        if set_line[-2] == "=":
+            name = set_line[1]
+        else:
+            name = set_line[-2]
+        value = set_line[-1].replace(";", "")
         tables.append({"name": name, "value": value})
 
     def parse_set_statement(
-        self, tables: List, line: str, set_line: Optional[str]
+        self, tables: List, line: str, set_line: Optional[str], set_was_in_line: bool
     ) -> Optional[str]:
-        set_was_in_line = False
         if re.match(r"SET", line):
             set_was_in_line = True
             if not set_line:
                 set_line = line
             else:
                 self.process_set(tables, set_line)
                 set_line = line
-        elif set_line and len(set_line.split()) == 3:
+        elif (set_line and len(set_line.split()) == 3) or (
+            set_line and set_was_in_line
+        ):
             self.process_set(tables, set_line)
             set_line = None
+            set_was_in_line = False
         return set_line, set_was_in_line
 
     def check_new_statement_start(self, line: str, statement: str) -> bool:
@@ -142,7 +147,11 @@ def parse_data(self):  # noqa: C901 need to refactor this
         data = self.pre_process_data(self.data)
         lines = data.replace("\\t", "").split("\\n")
         skip_line_words = ["USE", "GO"]
+
         set_line = None
+
+        set_was_in_line = False
+
         for num, line in enumerate(lines):
             line, block_comments = self.pre_process_line(line, block_comments)
             line = line.strip().replace("\n", "").replace("\t", "")
@@ -151,19 +160,21 @@ def parse_data(self):  # noqa: C901 need to refactor this
                 if line.startswith(word):
                     skip = True
                     break
-            set_line, set_was_in_line = self.parse_set_statement(tables, line, set_line)
+            set_line, set_was_in_line = self.parse_set_statement(
+                tables, line, set_line, set_was_in_line
+            )
             if line or num == len(lines) - 1:
                 # to avoid issues when comma or parath are glued to column name
                 new_statement = self.check_new_statement_start(line, statement)
 
-                final_line = line.endswith(";")
+                final_line = line.endswith(";") and not set_was_in_line
 
                 if not skip and not set_was_in_line and not new_statement:
                     if statement is None:
                         statement = line
                     else:
                         statement += f" {line}"
-
+                print(statement)
                 if final_line or new_statement:
                     # end of sql operation, remove ; from end of line
                     statement = statement[:-1]