iting0321 commented on code in PR #4827:
URL: https://github.com/apache/polaris/pull/4827#discussion_r3487974396
##########
persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java:
##########
@@ -1354,4 +1392,500 @@ private void writeCommitMetricsReport(@NonNull
ModelCommitMetricsReport report)
String.format("Failed to write commit metrics report due to %s",
e.getMessage()), e);
}
}
+
+ //
============================================================================
+ // LineagePersistence Implementation
+ //
============================================================================
+
+ @Override
+ public void upsertDatasets(RealmContext realmContext, List<LineageDataset>
datasets) {
+ verifyLineagePersistenceSupported();
+ String realmId = realmContext.getRealmIdentifier();
+ long nowMillis = Instant.now().toEpochMilli();
+ for (LineageDataset dataset : datasets) {
+ long datasetId =
+ lookupLineageDataset(realmId, dataset)
+ .map(ModelLineageDataset::getDatasetId)
+ .orElseGet(IdGenerator.getIdGenerator()::nextId);
+ ModelLineageDataset model =
+ ModelLineageDataset.fromDataset(dataset, realmId, datasetId,
nowMillis);
+ try {
+
datasourceOperations.executeUpdate(generateLineageDatasetUpsert(model));
+ } catch (SQLException e) {
+ throw new RuntimeException(
+ String.format("Failed to upsert lineage dataset due to %s",
e.getMessage()), e);
+ }
+ }
+ }
+
+ @Override
+ public void replaceDatasetEdges(
+ RealmContext realmContext,
+ List<LineageDataset> targetDatasets,
+ List<LineageEdge> edges,
+ Instant lastEventAt) {
+ verifyLineagePersistenceSupported();
+ String realmId = realmContext.getRealmIdentifier();
+ long lastEventAtMillis = lastEventAt.toEpochMilli();
+ Map<Long, List<ModelLineageEdge>> edgesByTargetDatasetId = new
LinkedHashMap<>();
+ for (LineageDataset targetDataset : targetDatasets) {
+ edgesByTargetDatasetId.putIfAbsent(
+ requireLineageDatasetId(realmId, targetDataset), new ArrayList<>());
+ }
+ for (LineageEdge edge : edges) {
+ ModelLineageEdge model =
+ ModelLineageEdge.fromIds(
+ realmId,
+ requireLineageDatasetId(realmId, edge.source()),
+ requireLineageDatasetId(realmId, edge.target()),
+ lastEventAtMillis);
+ edgesByTargetDatasetId
+ .computeIfAbsent(model.getTargetDatasetId(), ignored -> new
ArrayList<>())
+ .add(model);
+ }
+
+ try {
+ datasourceOperations.runWithinTransaction(
+ connection -> {
+ for (Map.Entry<Long, List<ModelLineageEdge>> entry :
+ edgesByTargetDatasetId.entrySet()) {
+ long targetDatasetId = entry.getKey();
+ Optional<Long> currentLastEventAt =
+ loadLineageDatasetLastEventAt(connection, realmId,
targetDatasetId);
+ if (currentLastEventAt.isPresent() && currentLastEventAt.get() >
lastEventAtMillis) {
+ continue;
+ }
+
+ List<Long> sourceDatasetIds =
+ entry.getValue().stream()
+ .map(ModelLineageEdge::getSourceDatasetId)
+ .distinct()
+ .toList();
+ datasourceOperations.execute(
+ connection,
+ generateDeleteStaleLineageEdges(realmId, targetDatasetId,
sourceDatasetIds));
+ datasourceOperations.execute(
+ connection,
+ generateDeleteStaleLineageColumnEdges(
+ realmId, targetDatasetId, sourceDatasetIds));
+ for (ModelLineageEdge model : entry.getValue()) {
+ datasourceOperations.execute(connection,
generateLineageEdgeUpsert(model));
+ }
+ datasourceOperations.execute(
+ connection,
+ generateLineageDatasetLastEventAtUpdate(
+ realmId, targetDatasetId, lastEventAtMillis));
+ }
+ return true;
+ });
+ } catch (SQLException e) {
+ throw new RuntimeException(
+ String.format("Failed to upsert lineage edge due to %s",
e.getMessage()), e);
+ }
+ }
+
+ private Optional<Long> loadLineageDatasetLastEventAt(
+ Connection connection, String realmId, long targetDatasetId) throws
SQLException {
+ String table =
QueryGenerator.getFullyQualifiedTableName(ModelLineageDataset.TABLE_NAME);
+ try (PreparedStatement statement =
+ connection.prepareStatement(
+ "SELECT last_lineage_event_at FROM "
+ + table
+ + " WHERE realm_id = ? AND dataset_id = ?")) {
+ statement.setString(1, realmId);
+ statement.setLong(2, targetDatasetId);
+ try (ResultSet resultSet = statement.executeQuery()) {
+ if (!resultSet.next()) {
+ return Optional.empty();
+ }
+ long value = resultSet.getLong(1);
+ return resultSet.wasNull() ? Optional.empty() : Optional.of(value);
+ }
+ }
+ }
+
+ @Override
+ public void upsertColumnEdges(
+ RealmContext realmContext, List<LineageColumnEdge> columnEdges, Instant
lastEventAt) {
+ verifyLineagePersistenceSupported();
+ String realmId = realmContext.getRealmIdentifier();
+ long lastEventAtMillis = lastEventAt.toEpochMilli();
+ for (LineageColumnEdge columnEdge : columnEdges) {
+ ModelLineageDataset targetDataset =
+ requireLineageDataset(realmId, columnEdge.target().dataset());
+ Long currentLastEventAt = targetDataset.getLastLineageEventAt();
+ if (currentLastEventAt != null && currentLastEventAt >
lastEventAtMillis) {
+ continue;
+ }
+
+ ModelLineageColumnEdge model =
+ ModelLineageColumnEdge.fromIds(
+ realmId,
+ requireLineageDatasetId(realmId, columnEdge.source().dataset()),
+ columnEdge.source().field(),
+ targetDataset.getDatasetId(),
+ columnEdge.target().field(),
+ lastEventAtMillis);
+ try {
+
datasourceOperations.executeUpdate(generateLineageColumnEdgeUpsert(model));
+ } catch (SQLException e) {
+ throw new RuntimeException(
+ String.format("Failed to upsert lineage column edge due to %s",
e.getMessage()), e);
+ }
+ }
+ }
+
+ @Override
+ public LineageGraph loadLineage(RealmContext realmContext,
LineageQueryRequest request) {
+ verifyLineagePersistenceSupported();
+ String realmId = realmContext.getRealmIdentifier();
+ Optional<ModelLineageDataset> requested =
+ lookupLineageDatasetByNodeId(realmId, request.nodeId());
+ if (requested.isEmpty()) {
+ return new LineageGraph(
+ new LineageNode(request.nodeId(), LineageNodeType.DATASET, null,
true),
+ List.of(),
+ List.of());
+ }
+
+ ModelLineageDataset dataset = requested.get();
+ boolean includeColumns = request.granularity() ==
LineageGranularity.COLUMN;
+ List<LineageNode> upstream =
+ request.direction() == LineageDirection.UPSTREAM
+ || request.direction() == LineageDirection.BOTH
+ ? loadAdjacentLineageNodes(realmId, dataset.getDatasetId(), true,
includeColumns)
+ : List.of();
+ List<LineageNode> downstream =
+ request.direction() == LineageDirection.DOWNSTREAM
+ || request.direction() == LineageDirection.BOTH
+ ? loadAdjacentLineageNodes(realmId, dataset.getDatasetId(), false,
includeColumns)
+ : List.of();
+
+ return new LineageGraph(toLineageNode(dataset, List.of()), upstream,
downstream);
+ }
+
+ private Optional<ModelLineageDataset> lookupLineageDataset(
+ String realmId, LineageDataset dataset) {
+ String table =
QueryGenerator.getFullyQualifiedTableName(ModelLineageDataset.TABLE_NAME);
+ PreparedQuery query =
+ new PreparedQuery(
+ "SELECT realm_id, dataset_id, catalog, namespace, name,
polaris_entity_id, last_lineage_event_at, created_at, updated_at "
+ + "FROM "
+ + table
+ + " WHERE realm_id = ? AND namespace = ? AND name = ?",
+ List.of(realmId, dataset.namespace(), dataset.name()));
Review Comment:
fixed
##########
persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java:
##########
@@ -1354,4 +1392,500 @@ private void writeCommitMetricsReport(@NonNull
ModelCommitMetricsReport report)
String.format("Failed to write commit metrics report due to %s",
e.getMessage()), e);
}
}
+
+ //
============================================================================
+ // LineagePersistence Implementation
+ //
============================================================================
+
+ @Override
+ public void upsertDatasets(RealmContext realmContext, List<LineageDataset>
datasets) {
+ verifyLineagePersistenceSupported();
+ String realmId = realmContext.getRealmIdentifier();
+ long nowMillis = Instant.now().toEpochMilli();
+ for (LineageDataset dataset : datasets) {
+ long datasetId =
+ lookupLineageDataset(realmId, dataset)
+ .map(ModelLineageDataset::getDatasetId)
+ .orElseGet(IdGenerator.getIdGenerator()::nextId);
+ ModelLineageDataset model =
+ ModelLineageDataset.fromDataset(dataset, realmId, datasetId,
nowMillis);
+ try {
+
datasourceOperations.executeUpdate(generateLineageDatasetUpsert(model));
+ } catch (SQLException e) {
+ throw new RuntimeException(
+ String.format("Failed to upsert lineage dataset due to %s",
e.getMessage()), e);
+ }
+ }
+ }
+
+ @Override
+ public void replaceDatasetEdges(
+ RealmContext realmContext,
+ List<LineageDataset> targetDatasets,
+ List<LineageEdge> edges,
+ Instant lastEventAt) {
+ verifyLineagePersistenceSupported();
+ String realmId = realmContext.getRealmIdentifier();
+ long lastEventAtMillis = lastEventAt.toEpochMilli();
+ Map<Long, List<ModelLineageEdge>> edgesByTargetDatasetId = new
LinkedHashMap<>();
+ for (LineageDataset targetDataset : targetDatasets) {
+ edgesByTargetDatasetId.putIfAbsent(
+ requireLineageDatasetId(realmId, targetDataset), new ArrayList<>());
+ }
+ for (LineageEdge edge : edges) {
+ ModelLineageEdge model =
+ ModelLineageEdge.fromIds(
+ realmId,
+ requireLineageDatasetId(realmId, edge.source()),
+ requireLineageDatasetId(realmId, edge.target()),
+ lastEventAtMillis);
+ edgesByTargetDatasetId
+ .computeIfAbsent(model.getTargetDatasetId(), ignored -> new
ArrayList<>())
+ .add(model);
+ }
+
+ try {
+ datasourceOperations.runWithinTransaction(
+ connection -> {
+ for (Map.Entry<Long, List<ModelLineageEdge>> entry :
+ edgesByTargetDatasetId.entrySet()) {
+ long targetDatasetId = entry.getKey();
+ Optional<Long> currentLastEventAt =
+ loadLineageDatasetLastEventAt(connection, realmId,
targetDatasetId);
+ if (currentLastEventAt.isPresent() && currentLastEventAt.get() >
lastEventAtMillis) {
+ continue;
+ }
+
+ List<Long> sourceDatasetIds =
+ entry.getValue().stream()
+ .map(ModelLineageEdge::getSourceDatasetId)
+ .distinct()
+ .toList();
+ datasourceOperations.execute(
+ connection,
+ generateDeleteStaleLineageEdges(realmId, targetDatasetId,
sourceDatasetIds));
+ datasourceOperations.execute(
+ connection,
+ generateDeleteStaleLineageColumnEdges(
+ realmId, targetDatasetId, sourceDatasetIds));
+ for (ModelLineageEdge model : entry.getValue()) {
+ datasourceOperations.execute(connection,
generateLineageEdgeUpsert(model));
+ }
+ datasourceOperations.execute(
+ connection,
+ generateLineageDatasetLastEventAtUpdate(
+ realmId, targetDatasetId, lastEventAtMillis));
+ }
+ return true;
+ });
+ } catch (SQLException e) {
+ throw new RuntimeException(
+ String.format("Failed to upsert lineage edge due to %s",
e.getMessage()), e);
+ }
+ }
+
+ private Optional<Long> loadLineageDatasetLastEventAt(
+ Connection connection, String realmId, long targetDatasetId) throws
SQLException {
+ String table =
QueryGenerator.getFullyQualifiedTableName(ModelLineageDataset.TABLE_NAME);
+ try (PreparedStatement statement =
+ connection.prepareStatement(
+ "SELECT last_lineage_event_at FROM "
+ + table
+ + " WHERE realm_id = ? AND dataset_id = ?")) {
+ statement.setString(1, realmId);
+ statement.setLong(2, targetDatasetId);
+ try (ResultSet resultSet = statement.executeQuery()) {
+ if (!resultSet.next()) {
+ return Optional.empty();
+ }
+ long value = resultSet.getLong(1);
+ return resultSet.wasNull() ? Optional.empty() : Optional.of(value);
+ }
+ }
+ }
+
+ @Override
+ public void upsertColumnEdges(
+ RealmContext realmContext, List<LineageColumnEdge> columnEdges, Instant
lastEventAt) {
+ verifyLineagePersistenceSupported();
+ String realmId = realmContext.getRealmIdentifier();
+ long lastEventAtMillis = lastEventAt.toEpochMilli();
+ for (LineageColumnEdge columnEdge : columnEdges) {
+ ModelLineageDataset targetDataset =
+ requireLineageDataset(realmId, columnEdge.target().dataset());
+ Long currentLastEventAt = targetDataset.getLastLineageEventAt();
+ if (currentLastEventAt != null && currentLastEventAt >
lastEventAtMillis) {
+ continue;
+ }
+
+ ModelLineageColumnEdge model =
+ ModelLineageColumnEdge.fromIds(
+ realmId,
+ requireLineageDatasetId(realmId, columnEdge.source().dataset()),
+ columnEdge.source().field(),
+ targetDataset.getDatasetId(),
+ columnEdge.target().field(),
+ lastEventAtMillis);
+ try {
+
datasourceOperations.executeUpdate(generateLineageColumnEdgeUpsert(model));
+ } catch (SQLException e) {
+ throw new RuntimeException(
+ String.format("Failed to upsert lineage column edge due to %s",
e.getMessage()), e);
+ }
+ }
+ }
+
+ @Override
+ public LineageGraph loadLineage(RealmContext realmContext,
LineageQueryRequest request) {
+ verifyLineagePersistenceSupported();
+ String realmId = realmContext.getRealmIdentifier();
+ Optional<ModelLineageDataset> requested =
+ lookupLineageDatasetByNodeId(realmId, request.nodeId());
+ if (requested.isEmpty()) {
+ return new LineageGraph(
+ new LineageNode(request.nodeId(), LineageNodeType.DATASET, null,
true),
+ List.of(),
+ List.of());
+ }
+
+ ModelLineageDataset dataset = requested.get();
+ boolean includeColumns = request.granularity() ==
LineageGranularity.COLUMN;
+ List<LineageNode> upstream =
+ request.direction() == LineageDirection.UPSTREAM
+ || request.direction() == LineageDirection.BOTH
+ ? loadAdjacentLineageNodes(realmId, dataset.getDatasetId(), true,
includeColumns)
+ : List.of();
+ List<LineageNode> downstream =
+ request.direction() == LineageDirection.DOWNSTREAM
+ || request.direction() == LineageDirection.BOTH
+ ? loadAdjacentLineageNodes(realmId, dataset.getDatasetId(), false,
includeColumns)
+ : List.of();
+
+ return new LineageGraph(toLineageNode(dataset, List.of()), upstream,
downstream);
+ }
+
+ private Optional<ModelLineageDataset> lookupLineageDataset(
+ String realmId, LineageDataset dataset) {
+ String table =
QueryGenerator.getFullyQualifiedTableName(ModelLineageDataset.TABLE_NAME);
+ PreparedQuery query =
+ new PreparedQuery(
+ "SELECT realm_id, dataset_id, catalog, namespace, name,
polaris_entity_id, last_lineage_event_at, created_at, updated_at "
+ + "FROM "
+ + table
+ + " WHERE realm_id = ? AND namespace = ? AND name = ?",
+ List.of(realmId, dataset.namespace(), dataset.name()));
+ try {
+ List<ModelLineageDataset> results =
+ datasourceOperations.executeSelect(query,
ModelLineageDataset.CONVERTER);
+ return results.stream().findFirst();
+ } catch (SQLException e) {
+ throw new RuntimeException(
+ String.format("Failed to load lineage dataset due to %s",
e.getMessage()), e);
+ }
+ }
+
+ private Optional<ModelLineageDataset> lookupLineageDatasetByNodeId(
+ String realmId, String nodeId) {
+ return parseLineageDatasetNodeId(nodeId)
+ .flatMap(dataset -> lookupLineageDataset(realmId, dataset));
+ }
+
+ private Optional<LineageDataset> parseLineageDatasetNodeId(String nodeId) {
+ String prefix = "dataset:";
+ if (!nodeId.startsWith(prefix)) {
+ return Optional.empty();
+ }
+ String identity = nodeId.substring(prefix.length());
+ int catalogSeparator = identity.lastIndexOf(':');
+ int nameSeparator = identity.lastIndexOf('.');
+ if (catalogSeparator < 0 || nameSeparator <= catalogSeparator + 1) {
+ return Optional.empty();
+ }
+ String catalog = identity.substring(0, catalogSeparator);
+ String namespace = identity.substring(catalogSeparator + 1, nameSeparator);
+ String name = identity.substring(nameSeparator + 1);
+ return Optional.of(new LineageDataset(catalog, namespace, name));
+ }
+
+ private long requireLineageDatasetId(String realmId, LineageDataset dataset)
{
+ return requireLineageDataset(realmId, dataset).getDatasetId();
+ }
+
+ private ModelLineageDataset requireLineageDataset(String realmId,
LineageDataset dataset) {
+ return lookupLineageDataset(realmId, dataset)
+ .orElseThrow(
+ () ->
+ new IllegalArgumentException(
+ String.format(
+ "Lineage dataset '%s.%s' does not exist in realm '%s'.
Call upsertDatasets before writing edges.",
+ dataset.namespace(), dataset.name(), realmId)));
+ }
+
+ private List<LineageNode> loadAdjacentLineageNodes(
+ String realmId, long datasetId, boolean upstream, boolean
includeColumns) {
+ String edgesTable =
QueryGenerator.getFullyQualifiedTableName(ModelLineageEdge.TABLE_NAME);
+ String datasetsTable =
+
QueryGenerator.getFullyQualifiedTableName(ModelLineageDataset.TABLE_NAME);
+ String adjacentEdgeColumn = upstream ? "source_dataset_id" :
"target_dataset_id";
+ String requestedEdgeColumn = upstream ? "target_dataset_id" :
"source_dataset_id";
+ PreparedQuery query =
+ new PreparedQuery(
+ "SELECT d.realm_id, d.dataset_id, d.catalog, d.namespace, d.name,
d.polaris_entity_id, d.last_lineage_event_at, d.created_at, d.updated_at "
+ + "FROM "
+ + edgesTable
+ + " e JOIN "
+ + datasetsTable
+ + " d ON d.realm_id = e.realm_id AND d.dataset_id = e."
+ + adjacentEdgeColumn
+ + " WHERE e.realm_id = ? AND e."
+ + requestedEdgeColumn
+ + " = ?",
+ List.of(realmId, datasetId));
+ try {
+ List<ModelLineageDataset> datasets =
+ datasourceOperations.executeSelect(query,
ModelLineageDataset.CONVERTER);
+ return datasets.stream()
+ .map(
+ dataset ->
+ toLineageNode(
+ dataset,
+ includeColumns
+ ? loadFieldMappings(realmId, datasetId,
dataset.getDatasetId(), upstream)
+ : List.of()))
Review Comment:
fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
