public static Set extractColumnLevelLineage()

in zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/lineage/ColumnLineageExtractor.java [169:201]

29 lines of code
2 McCabe index (conditional complexity)


  public static Set<ColumnLineage> extractColumnLevelLineage(ResolvedInsertStmt insertStmt) {
    if (Objects.isNull(insertStmt.getQuery())) {
      // The statement is inserting rows manually using "INSERT INTO ... VALUES ..."
      // Since it does not query any tables, it does not produce lineage
      return ImmutableSet.of();
    }

    Table targetTable = insertStmt.getTableScan().getTable();
    ResolvedScan query = insertStmt.getQuery();
    List<ResolvedColumn> insertedColumns =
        insertStmt.getInsertColumnList().stream()
            .map(ColumnLineageExtractor::expandColumn)
            .flatMap(List::stream)
            .collect(Collectors.toList());
    List<ResolvedColumn> matchingColumnsInQuery =
        query.getColumnList().stream()
            .map(ColumnLineageExtractor::expandColumn)
            .flatMap(List::stream)
            .collect(Collectors.toList());

    return IntStream.range(0, insertedColumns.size())
        .mapToObj(
            index ->
                new SimpleEntry<>(
                    insertedColumns.get(index),
                    ParentColumnFinder.findParentsForColumn(
                        insertStmt, matchingColumnsInQuery.get(index))))
        .map(
            entry ->
                buildColumnLineage(
                    targetTable.getFullName(), entry.getKey().getName(), entry.getValue()))
        .collect(Collectors.toSet());
  }