adding sql into json

lineagex/ColumnLineageNoConn.py

@@ -249,14 +249,15 @@ def _sub_shared_col_conds(self, sql_ast: expressions = None) -> None:
 
     def _sub_shared_col_conds_cte(
         self, sql_ast: expressions = None
-    ) -> Tuple[List, List, List]:
+    ) -> Tuple[List, List, List, str]:
         """
         Run the subquery inside the cte first that is with the shared conditions(WHERE, GROUP BY, etc)
         :param sql_ast: the ast for the cte
         """
         all_cte_sub_table = []
         potential_cte_sub_table = []
         all_cte_sub_cols = []
+        sub_name = ""
         # add in more conditions, including FROM/JOIN
         for cond in shared_conditions_with_table:
             for cond_sql in sql_ast.find_all(cond):
@@ -309,7 +310,7 @@ def _sub_shared_col_conds_cte(
                         self.cte_dict[sub_name] = temp_sub_dict
                         sub_ast.replace(exp.Table(this=sub_name))
                     sub_ast.pop()
-        return all_cte_sub_table, all_cte_sub_cols, potential_cte_sub_table
+        return all_cte_sub_table, all_cte_sub_cols, potential_cte_sub_table, sub_name
 
     def _run_cte_lineage(self):
         """
@@ -320,13 +321,16 @@ def _run_cte_lineage(self):
                 all_cte_sub_table,
                 all_cte_sub_cols,
                 potential_cte_sub_table,
+                sub_name
             ) = self._sub_shared_col_conds_cte(sql_ast=cte)
             self.all_used_col = []
             temp_cte_dict = {}
             temp_cte_table = self._resolve_table(part_ast=cte)
             if len(temp_cte_table) == 0:
                 temp_cte_table = potential_cte_sub_table
             cte_name = cte.find(exp.TableAlias).alias_or_name
+            if sub_name in temp_cte_table:
+                temp_cte_table.remove(sub_name)
             self.cte_table_dict[cte_name] = list(
                 set(
                     self._find_all_tables(temp_table_list=temp_cte_table)

lineagex/LineageXNoConn.py

@@ -95,6 +95,7 @@ def _run_lineage_no_conn(self, name: Optional[str] = "", sql: Optional[str] = ""
             "tables": col_lineage.table_list,
             "columns": col_lineage.column_dict,
             "table_name": name,
+            "sql": sql,
         }
         # add to the dict with the already parsed tables
         self.input_table_dict[self.target_schema + "." + name] = list(

lineagex/LineageXWithConn.py

@@ -103,6 +103,7 @@ def _run_table_lineage(self) -> None:
                         "tables": col_lineage.table_list,
                         "columns": col_lineage.column_dict,
                         "table_name": name,
+                        "sql": sql,
                     }
                     self.parsed += 1
                 except Exception as e:
@@ -268,6 +269,7 @@ def _explain_sql(self, name: Optional[str] = "", sql: Optional[str] = "") -> Non
                 "tables": col_lineage.table_list,
                 "columns": col_lineage.column_dict,
                 "table_name": table_name,
+                "sql": sql,
             }
             self.finished_list.append(name)
             self.parsed += 1

lineagex/utils.py

@@ -163,10 +163,12 @@ def produce_json(
         for i in cols:
             base_table_dict[t]["columns"][i] = [""]
         base_table_dict[t]["table_name"] = str(t)
+        base_table_dict[t]["sql"] = "this is a base table"
     base_table_dict.update(output_dict)
     with open("output.json", "w") as outfile:
         json.dump(base_table_dict, outfile)
-    _produce_html(output_json=str(base_table_dict).replace("'", '"'))
+    #_produce_html(output_json=str(base_table_dict).replace("'", '"'))
+    _produce_html(output_json=base_table_dict)
     return base_table_dict
 
 
@@ -189,7 +191,7 @@ def _guess_base_table(output_dict: Optional[dict] = None) -> dict:
     return base_table_noconn_dict
 
 
-def _produce_html(output_json: Optional[str] = "") -> None:
+def _produce_html(output_json: Optional[dict] = "") -> None:
     """
     Produce the html file for viewing
     :param output_json: the final output.json file
@@ -212,7 +214,7 @@ def _produce_html(output_json: Optional[str] = "") -> None:
           <div id="main"></div>
         <script type="text/javascript" src="vendor.js"></script><script type="text/javascript" src="app.js"></script></body>
         </html>""".format(
-                output_json
+                json.dumps(output_json)
             )
         )