flink-sql-language-server
Version:
A LSP-based language server for Apache Flink SQL
227 lines (226 loc) • 8.88 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.LineageVisitor = void 0;
const flinksql_relation_visitor_1 = require("./flinksql-relation.visitor");
class LineageVisitor extends flinksql_relation_visitor_1.FlinkSQLRelationVisitor {
constructor(options, getTableMetadata) {
super(getTableMetadata);
this.options = options;
this.getTableMetadata = getTableMetadata;
this.relations = [];
this.edges = [];
this.edgeIdSeq = 1;
this.queryIdSeq = 1;
this.viewIdSeq = 1;
}
findColumnByEdgeIdentity(identity) {
return this.relations
.find(r => r.relation.id === identity.relationId)
?.relation.columns.find(c => c.id === identity.columnId);
}
defaultResult() {
return;
}
getLineage() {
const cleanedNodes = [];
const edges = [...this.edges];
if (this.options.mergeLeaves) {
const deduplicateTable = new Map();
const usedRelations = new Map();
this.relations.forEach(r => {
if (r.relationPrimary === undefined) {
cleanedNodes.push(r.relation);
return;
}
const key = (0, flinksql_relation_visitor_1.relationPrimaryToString)(r.relationPrimary);
const entry = usedRelations.get(key);
if (entry !== undefined) {
if (r.relation.type === 'view') {
usedRelations.delete(key);
usedRelations.set(key, r.relation.id);
deduplicateTable.set(entry, r.relation.id);
}
else {
deduplicateTable.set(r.relation.id, entry);
}
}
else {
usedRelations.set(key, r.relation.id);
r.relation.label = r.relationPrimary.relationName;
cleanedNodes.push(r.relation);
}
});
const usedColumns = new Map();
edges.forEach(e => {
const remapEdge = (identity) => {
const originalIdentity = { ...identity };
const remappedSourceTable = deduplicateTable.get(identity.relationId);
if (remappedSourceTable !== undefined) {
identity.relationId = remappedSourceTable;
}
if (identity.columnId !== undefined) {
const column = this.findColumnByEdgeIdentity(originalIdentity);
if (column !== undefined) {
const columns = usedColumns.get(identity.relationId);
if (columns !== undefined) {
const sameColumn = columns.find(c => c.label === column.label);
if (sameColumn !== undefined) {
identity.columnId = sameColumn.id;
return;
}
columns.push(column);
}
else {
usedColumns.set(identity.relationId, [column]);
}
}
}
};
remapEdge(e.source);
remapEdge(e.target);
});
cleanedNodes.forEach(n => {
if (n.data !== undefined) {
const tableColumns = usedColumns.get(n.id);
if (tableColumns !== undefined) {
n.columns = [...new Set([...n.columns, ...tableColumns])];
}
}
});
}
else {
this.relations.forEach(t => cleanedNodes.push(t.relation));
}
const sourceNodes = new Set();
const targetNodes = new Set();
const isStarOrCountStar = (label) => {
return label === flinksql_relation_visitor_1.STAR_LABEL || label === flinksql_relation_visitor_1.COUNT_STAR_LABEL;
};
edges.forEach(e => {
sourceNodes.add(e.source.relationId);
targetNodes.add(e.target.relationId);
});
cleanedNodes.forEach(t => {
if (!sourceNodes.has(t.id)) {
t.isTargetOnly = true;
}
if (!targetNodes.has(t.id)) {
t.isSourceOnly = true;
}
if (t.data !== undefined) {
const ddlColumns = this.relationDDLColumns.get((0, flinksql_relation_visitor_1.relationPrimaryToString)(t.data)) ?? this.getTableMetadata?.(t.data)?.columns;
if (!this.options.ignoreDDLs) {
t.columns = t.columns.map(col => ({
...col,
inconsistent: !isStarOrCountStar(col.label) && !!ddlColumns && ddlColumns.every(d => d.label !== col.label)
}));
}
}
});
return {
nodes: cleanedNodes,
edges
};
}
aggregateResult(_aggregate, _nextResult) {
return;
}
onColumnReference(source, target, currentClause) {
this.edges.push({
id: `${this.edgeIdSeq++}`,
type: 'edge',
edgeType: currentClause ?? this.currentQueryRelation?.currentClause,
source: {
relationId: source.relationId,
columnId: source.columnId
},
target: {
relationId: target.relationId,
columnId: target.columnId
}
});
}
onRelation(relation, alias) {
const old = this.relations.findIndex(rel => rel.relation.id === relation.id);
if (old !== -1) {
this.relations.splice(old, 1);
}
const columns = relation.columns.map(c => {
return {
id: c.id,
label: c.label,
range: c.range,
isAssumed: c.isAssumed
};
});
if (relation instanceof flinksql_relation_visitor_1.TableRelation) {
this.relations.push({
relationPrimary: relation.tablePrimary,
relation: {
type: 'table',
id: relation.id,
label: alias !== undefined && alias != relation.tablePrimary.relationName
? `${relation.tablePrimary.relationName} -> ${alias}`
: relation.tablePrimary.relationName,
range: relation.range,
isSourceOnly: false,
isTargetOnly: false,
data: relation.tablePrimary,
columns
}
});
}
else if (relation instanceof flinksql_relation_visitor_1.QueryRelation) {
let label = `query_${this.queryIdSeq++}`;
if (alias !== undefined) {
label += ` -> ${alias}`;
}
this.relations.push({
relation: {
type: 'query',
id: relation.id,
label,
isSourceOnly: false,
isTargetOnly: false,
range: relation.range,
columns
}
});
}
else if (relation instanceof flinksql_relation_visitor_1.InsertRelation) {
this.relations.push({
relationPrimary: relation.sinkTableRelation?.tablePrimary,
relation: {
type: 'insert',
id: relation.id,
label: relation.sinkTableRelation?.tablePrimary.relationName || relation.id,
range: relation.range,
isSourceOnly: false,
isTargetOnly: false,
data: relation.tablePrimary,
columns
}
});
}
else {
let label = `view_${this.viewIdSeq++}`;
if (alias !== undefined) {
label = alias;
}
this.relations.push({
relationPrimary: relation.tablePrimary,
relation: {
type: 'view',
id: relation.id,
label,
range: relation.range,
isSourceOnly: false,
isTargetOnly: false,
data: relation.tablePrimary,
columns
}
});
}
}
}
exports.LineageVisitor = LineageVisitor;