diff --git a/src/main/java/ai/starlake/transpiler/schema/JdbcMetaData.java b/src/main/java/ai/starlake/transpiler/schema/JdbcMetaData.java index 0b17a46..b46e3dd 100644 --- a/src/main/java/ai/starlake/transpiler/schema/JdbcMetaData.java +++ b/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; diff --git a/src/main/java/ai/starlake/transpiler/schema/JdbcTable.java b/src/main/java/ai/starlake/transpiler/schema/JdbcTable.java index 991c4ef..cd60869 100644 --- a/src/main/java/ai/starlake/transpiler/schema/JdbcTable.java +++ b/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, diff --git a/src/test/java/ai/starlake/transpiler/AbstractColumnResolverTest.java b/src/test/java/ai/starlake/transpiler/AbstractColumnResolverTest.java index f5d7cd3..670e9c7 100644 --- a/src/test/java/ai/starlake/transpiler/AbstractColumnResolverTest.java +++ b/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 { diff --git a/src/test/java/ai/starlake/transpiler/JSQLColumnResolverTest.java b/src/test/java/ai/starlake/transpiler/JSQLColumnResolverTest.java index 71881ba..2ee54c8 100644 --- a/src/test/java/ai/starlake/transpiler/JSQLColumnResolverTest.java +++ b/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); + } + } diff --git a/src/test/java/ai/starlake/transpiler/snowflake/AsciiTreeBuilder.java b/src/test/java/ai/starlake/transpiler/snowflake/AsciiTreeBuilder.java index f3f3c4d..963e466 100644 --- a/src/test/java/ai/starlake/transpiler/snowflake/AsciiTreeBuilder.java +++ b/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");