Copilot commented on code in PR #4827: URL: https://github.com/apache/polaris/pull/4827#discussion_r3449811649
########## runtime/service/src/main/java/org/apache/polaris/service/lineage/DefaultLineageService.java: ########## @@ -0,0 +1,77 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.polaris.service.lineage; + +import jakarta.enterprise.context.RequestScoped; +import jakarta.inject.Inject; +import java.time.Instant; +import org.apache.polaris.core.config.FeatureConfiguration; +import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; +import org.apache.polaris.core.lineage.LineageGraph; +import org.apache.polaris.core.lineage.LineageIngestRequest; +import org.apache.polaris.core.lineage.LineagePersistence; +import org.apache.polaris.core.lineage.LineageQueryRequest; +import org.apache.polaris.core.lineage.LineageService; + +@RequestScoped +public class DefaultLineageService implements LineageService { + private final CallContext callContext; + private final LineageConfiguration configuration; + private final LineagePersistence persistence; + + @Inject + public DefaultLineageService( + CallContext callContext, LineageConfiguration configuration, LineagePersistence persistence) { + this.callContext = callContext; + this.configuration = configuration; + this.persistence = persistence; + } Review Comment: `DefaultLineageService` injects `LineagePersistence`, but the only CDI bean in this PR appears to be `DisabledLineagePersistence` (`@DefaultBean`). Without a producer that adapts the per-request `BasePersistence` session (from `MetaStoreManagerFactory.getOrCreateSession(...)`) to `LineagePersistence`, lineage ingest/query will keep throwing even when JDBC supports it. Also, `LineageConfiguration.persistence().enabled/type` are currently unused here, so the documented config knobs won’t have any effect. Please wire `LineagePersistence` similarly to how `MetricsPersistence` is wired in `ServiceProducers` (reuse the session when it implements the SPI), and gate usage on the intended config flag(s). ########## 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(); Review Comment: These lineage methods derive the `realmId` from the *passed* `RealmContext`. `JdbcBasePersistenceImpl` instances are created per-realm (constructor already receives `realmId`), and most other persistence operations use the instance field. Using the caller-provided realm here makes it easy for an internal misuse to write/query the wrong tenant with a session created for another realm. Consider using `this.realmId` (and optionally asserting it matches `realmContext.getRealmIdentifier()`) for consistency and to harden tenant isolation. ########## 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: With `includeColumns` enabled, this stream mapping calls `loadFieldMappings(...)` once per adjacent dataset (N+1 query pattern). For nodes with many direct upstream/downstream neighbors, a single lineage request can fan out into many extra DB queries. Consider batching field-mapping loads (single query for all adjacent dataset ids + grouping in-memory) to keep query count bounded. ########## 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: `lookupLineageDataset(...)` (and the table uniqueness in schema-v5) identify datasets by `(realm_id, namespace, name)` only, but the public node id format includes `catalog` (`dataset:<catalog>:<namespace>.<name>`). This makes node IDs effectively non-canonical (queries with the “wrong” catalog will still resolve), and it prevents representing two datasets that differ only by catalog/platform. Please align identity across: (1) node-id parsing/formatting, (2) uniqueness/indexes in `lineage_datasets`, and (3) all lookup/upsert predicates (either include `catalog` everywhere, or remove it from the node id and schema). -- 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]
