fix: preserve lineage within query blocks

- fixes #41

Signed-off-by: Andreas Reichel <[email protected]>

src/main/java/ai/starlake/transpiler/schema/JdbcMetaData.java

@@ -237,7 +237,7 @@ public JdbcTable put(JdbcResultSetMetaData rsMetaData, String name, String error
       JdbcTable t = new JdbcTable(currentCatalogName, currentSchemaName, name);
       int columnCount = rsMetaData.getColumnCount();
       for (int i = 1; i <= columnCount; i++) {
-        t.add(t.tableCatalog, t.tableSchema, t.tableName,
+        JdbcColumn col = t.add(t.tableCatalog, t.tableSchema, t.tableName,
             rsMetaData.getColumnLabel(i) != null && !rsMetaData.getColumnLabel(i).isEmpty()
                 ? rsMetaData.getColumnLabel(i)
                 : rsMetaData.getColumnName(i),
@@ -254,6 +254,10 @@ public JdbcTable put(JdbcResultSetMetaData rsMetaData, String name, String error
                 ? rsMetaData.getScopeTable(i)
                 : rsMetaData.getTableName(i),
             rsMetaData.getColumnName(i), null, "", "");
+
+        // add the Lineage Information, 0-Indexed
+        col.add(rsMetaData.columns.get(i - 1).getChildren());
+        col.setExpression(rsMetaData.columns.get(i - 1).getExpression());
       }
       put(t);
       return t;

src/main/java/ai/starlake/transpiler/schema/JdbcTable.java

@@ -378,7 +378,8 @@ public JdbcColumn add(JdbcColumn jdbcColumn) {
       jdbcColumn.columnName = jdbcColumn.columnName + "_" + i;
 
     }
-    return columns.put(jdbcColumn.columnName, jdbcColumn);
+    columns.put(jdbcColumn.columnName, jdbcColumn);
+    return jdbcColumn;
   }
 
   public JdbcColumn add(String tableCatalog, String tableSchema, String tableName,

src/test/java/ai/starlake/transpiler/AbstractColumnResolverTest.java

@@ -121,6 +121,17 @@ void resolve(File f, int idx, SQLTest t) throws Exception {
 
   }
 
+  static String assertLineage(String[][] schemaDefinition, String sqlStr, String expected)
+      throws SQLException, InvocationTargetException, NoSuchMethodException, InstantiationException,
+      IllegalAccessException, JSQLParserException {
+
+    JSQLColumResolver resolver = new JSQLColumResolver(new JdbcMetaData(schemaDefinition));
+    String actual = resolver.getLineage(AsciiTreeBuilder.class, sqlStr);
+    Assertions.assertThat(actual).isEqualToIgnoringWhitespace(expected);
+
+    return actual;
+  }
+
   static String assertLineage(JdbcMetaData metaData, String sqlStr, String expected)
       throws SQLException, InvocationTargetException, NoSuchMethodException, InstantiationException,
       IllegalAccessException, JSQLParserException {

src/test/java/ai/starlake/transpiler/JSQLColumnResolverTest.java

@@ -527,11 +527,11 @@ void testWithBQProjectIdAndQuotes()
 
     //@formatter:off
     String lineage =
-            "SELECT\n"
-            + " ├─mycte.id → sales.customers.id : Other\n"
-            + " ├─sum AS Function sum\n"
-            + " │  └─mycte.amount → sales.orders.amount : Other\n"
-            + " └─mycte.timestamp1 → timestamp1 : Other\n";
+            "SELECT\n" +
+            " ├─mycte.id → sales.customers.id : Other\n" +
+            " ├─sum AS Function sum\n" +
+            " │  └─mycte.amount → sales.orders.amount : Other\n" +
+            " └─mycte.timestamp1 AS TimeKeyExpression: CURRENT_TIMESTAMP()\n";
     //@formatter:on
     assertLineage(JdbcMetaData.copyOf(metaData), sqlStr, lineage);
   }
@@ -565,11 +565,11 @@ void testWithAllQuoted() throws JSQLParserException, SQLException, InvocationTar
 
     //@formatter:off
     String lineage =
-            "SELECT\n"
-            + " ├─mycte.id → sales.customers.id : Other\n"
-            + " ├─sum AS Function sum\n"
-            + " │  └─mycte.amount → sales.orders.amount : Other\n"
-            + " └─mycte.timestamp → timestamp : Other\n";
+            "SELECT\n" +
+            " ├─mycte.id → sales.customers.id : Other\n" +
+            " ├─sum AS Function sum\n" +
+            " │  └─mycte.amount → sales.orders.amount : Other\n" +
+            " └─mycte.timestamp AS TimeKeyExpression: CURRENT_TIMESTAMP()\n";
     //@formatter:on
     assertLineage(JdbcMetaData.copyOf(metaData), sqlStr, lineage);
   }
@@ -616,11 +616,11 @@ public void call() throws Throwable {
 
     //@formatter:off
     String lineage =
-            "SELECT\n"
-            + " ├─mycte.id → sales.customers.id : Other\n"
-            + " ├─sum AS Function sum\n"
-            + " │  └─unresolvable\n"
-            + " └─mycte.timestamp → timestamp : Other\n";
+            "SELECT\n" +
+            " ├─mycte.id → sales.customers.id : Other\n" +
+            " ├─sum AS Function sum\n" +
+            " │  └─unresolvable\n" +
+            " └─mycte.timestamp AS TimeKeyExpression: CURRENT_TIMESTAMP()\n";
     //@formatter:on
     assertLineage(JdbcMetaData.copyOf(metaData), sqlStr, lineage);
 
@@ -633,10 +633,10 @@ public void call() throws Throwable {
 
     //@formatter:off
     lineage =
-            "SELECT\n"
-            + " ├─mycte.id → sales.customers.id : Other\n"
-            + " ├─sum AS Function sum\n"
-            + " └─mycte.timestamp → timestamp : Other\n";
+        "SELECT\n" +
+        " ├─mycte.id → sales.customers.id : Other\n" +
+        " ├─sum AS Function sum\n" +
+        " └─mycte.timestamp AS TimeKeyExpression: CURRENT_TIMESTAMP()\n";
     //@formatter:on
     assertLineage(JdbcMetaData.copyOf(metaData), sqlStr, lineage);
   }
@@ -673,11 +673,11 @@ void testWithAllBackTickQuoted()
 
     //@formatter:off
     String lineage =
-            "SELECT\n"
-            + " ├─mycte.id → sales.customers.id : Other\n"
-            + " ├─sum AS Function sum\n"
-            + " │  └─mycte.amount → sales.orders.amount : Other\n"
-            + " └─mycte.timestamp → timestamp : Other\n";
+            "SELECT\n" +
+            " ├─mycte.id → sales.customers.id : Other\n" +
+            " ├─sum AS Function sum\n" +
+            " │  └─mycte.amount → sales.orders.amount : Other\n" +
+            " └─mycte.timestamp AS TimeKeyExpression: CURRENT_TIMESTAMP()\n";
     //@formatter:on
     assertLineage(JdbcMetaData.copyOf(metaData), sqlStr, lineage);
   }
@@ -712,11 +712,11 @@ void testWithBigQuerySingleQuotePair()
 
     //@formatter:off
     String lineage =
-            "SELECT\n"
-            + " ├─mycte.id → sales.customers.id : Other\n"
-            + " ├─sum AS Function sum\n"
-            + " │  └─mycte.amount → sales.orders.amount : Other\n"
-            + " └─mycte.timestamp → timestamp : Other\n";
+            "SELECT\n" +
+            " ├─mycte.id → sales.customers.id : Other\n" +
+            " ├─sum AS Function sum\n" +
+            " │  └─mycte.amount → sales.orders.amount : Other\n" +
+            " └─mycte.timestamp AS TimeKeyExpression: CURRENT_TIMESTAMP()\n";
     //@formatter:on
     assertLineage(JdbcMetaData.copyOf(metaData), sqlStr, lineage);
   }
@@ -768,12 +768,12 @@ void testWithWith() throws JSQLParserException, SQLException, InvocationTargetEx
 
     //@formatter:off
     String lineage =
-            "SELECT\n"
-            + " ├─yourcte.id → sales.customers.id : Other\n"
-            + " ├─sum AS Function Sum\n"
-            + " │  └─yourcte.amount → sales.orders.amount : Other\n"
-            + " ├─yourcte.timestamp1 → mycte.timestamp1 : Other\n"
-            + " └─amount2 AS yourcte.amount → sales.orders.amount : Other\n";
+            "SELECT\n" +
+            " ├─yourcte.id → sales.customers.id : Other\n" +
+            " ├─sum AS Function Sum\n" +
+            " │  └─yourcte.amount → sales.orders.amount : Other\n" +
+            " ├─yourcte.timestamp1 AS TimeKeyExpression: CURRENT_TIMESTAMP()\n" +
+            " └─amount2 AS yourcte.amount → sales.orders.amount : Other\n";
     //@formatter:on
     assertLineage(JdbcMetaData.copyOf(metaData), sqlStr, lineage);
   }
@@ -795,4 +795,33 @@ void testScopeColumn() throws JSQLParserException, SQLException, InvocationTarge
     System.out.println(resolver.getLineage(AsciiTreeBuilder.class, sqlStr));
   }
 
+  @Test
+  void testWithAndFunctionClauseIssue41()
+      throws JSQLParserException, SQLException, InvocationTargetException, NoSuchMethodException,
+      InstantiationException, IllegalAccessException {
+
+    //@formatter:off
+    String[][] schemaDefinition = {
+            {"a", "col1", "col2"}
+    };
+
+    String sqlStr =
+          " WITH d AS (\n"
+        + "        SELECT SUM(a.col1, a.col2) as colx\n"
+        + "        FROM a )\n" + " SELECT *\n"
+        + " FROM d\n"
+        + ";"
+        ;
+
+    String expected =
+          "SELECT\n"
+        + " └─d.colx AS Function SUM\n"
+        + "    ├─a.col1 : Other\n"
+        + "    └─a.col2 : Other"
+        ;
+    //@formatter:on
+
+    assertLineage(schemaDefinition, sqlStr, expected);
+  }
+
 }

src/test/java/ai/starlake/transpiler/snowflake/AsciiTreeBuilder.java

@@ -48,6 +48,17 @@ private String getNodeText(JdbcColumn column, String alias) {
 
     Expression expression = column.getExpression();
     if (expression instanceof Function) {
+      if (!StringUtils.isEmpty(column.columnName) && column.getChildren().size() > 1) {
+        if (column.tableCatalog != null && !column.tableCatalog.isEmpty()) {
+          b.append(column.tableCatalog).append(".")
+              .append(column.tableSchema != null ? column.tableSchema : "").append(".");
+        } else if (column.tableSchema != null && !column.tableSchema.isEmpty()) {
+          b.append(column.tableSchema).append(".");
+        }
+        b.append(column.tableName).append(".").append(column.columnName);
+        b.append(" AS ");
+      }
+
       Function f = (Function) expression;
       b.append("Function ").append(f.getName());
     } else if (expression instanceof Select) {
@@ -99,6 +110,17 @@ private String getNodeText(JdbcColumn column, String alias) {
 
       return b.toString();
     } else if (expression != null) {
+      if (!StringUtils.isEmpty(column.tableName)) {
+        if (column.tableCatalog != null && !column.tableCatalog.isEmpty()) {
+          b.append(column.tableCatalog).append(".")
+              .append(column.tableSchema != null ? column.tableSchema : "").append(".");
+        } else if (column.tableSchema != null && !column.tableSchema.isEmpty()) {
+          b.append(column.tableSchema).append(".");
+        }
+        b.append(column.tableName).append(".").append(column.columnName);
+        b.append(" AS ");
+      }
+
       b.append(expression.getClass().getSimpleName()).append(": ").append(expression);
     } else {
       b.append("unresolvable");