Skip to content

Commit e11af5f

Browse files
iting0321iting0321
iting0321
authored and
iting0321
committed
remove duplicate lineage persistence record types
1 parent 998465a commit e11af5f

11 files changed

Lines changed: 394 additions & 435 deletions

File tree

‎persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java‎

Lines changed: 237 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import jakarta.annotation.Nullable;
2626
import java.sql.Connection;
2727
import java.sql.SQLException;
28+
import java.time.Instant;
2829
import java.util.ArrayList;
2930
import java.util.Collections;
3031
import java.util.HashMap;
@@ -33,12 +34,14 @@
3334
import java.util.Map;
3435
import java.util.Objects;
3536
import java.util.Optional;
37+
import java.util.OptionalLong;
3638
import java.util.concurrent.atomic.AtomicReference;
3739
import java.util.function.Function;
3840
import java.util.function.Predicate;
3941
import java.util.stream.Collectors;
4042
import org.apache.polaris.core.PolarisCallContext;
4143
import org.apache.polaris.core.PolarisDiagnostics;
44+
import org.apache.polaris.core.context.RealmContext;
4245
import org.apache.polaris.core.entity.EntityNameLookupRecord;
4346
import org.apache.polaris.core.entity.LocationBasedEntity;
4447
import org.apache.polaris.core.entity.PolarisBaseEntity;
@@ -51,16 +54,25 @@
5154
import org.apache.polaris.core.entity.PolarisEvent;
5255
import org.apache.polaris.core.entity.PolarisGrantRecord;
5356
import org.apache.polaris.core.entity.PolarisPrincipalSecrets;
57+
import org.apache.polaris.core.lineage.LineageColumnEdge;
58+
import org.apache.polaris.core.lineage.LineageData;
59+
import org.apache.polaris.core.lineage.LineageDataset;
60+
import org.apache.polaris.core.lineage.LineageDirection;
61+
import org.apache.polaris.core.lineage.LineageEdge;
62+
import org.apache.polaris.core.lineage.LineageFieldMapping;
63+
import org.apache.polaris.core.lineage.LineageGranularity;
64+
import org.apache.polaris.core.lineage.LineageGraph;
65+
import org.apache.polaris.core.lineage.LineageNode;
66+
import org.apache.polaris.core.lineage.LineageNodeType;
67+
import org.apache.polaris.core.lineage.LineagePersistence;
68+
import org.apache.polaris.core.lineage.LineageQueryRequest;
5469
import org.apache.polaris.core.persistence.BaseMetaStoreManager;
5570
import org.apache.polaris.core.persistence.BasePersistence;
5671
import org.apache.polaris.core.persistence.EntityAlreadyExistsException;
5772
import org.apache.polaris.core.persistence.IntegrationPersistence;
5873
import org.apache.polaris.core.persistence.PolicyMappingAlreadyExistsException;
5974
import org.apache.polaris.core.persistence.PrincipalSecretsGenerator;
6075
import org.apache.polaris.core.persistence.RetryOnConcurrencyException;
61-
import org.apache.polaris.core.persistence.lineage.LineageColumnEdgeRecord;
62-
import org.apache.polaris.core.persistence.lineage.LineageDatasetRecord;
63-
import org.apache.polaris.core.persistence.lineage.LineageEdgeRecord;
6476
import org.apache.polaris.core.persistence.metrics.CommitMetricsRecord;
6577
import org.apache.polaris.core.persistence.metrics.ScanMetricsRecord;
6678
import org.apache.polaris.core.persistence.pagination.EntityIdToken;
@@ -73,6 +85,7 @@
7385
import org.apache.polaris.core.storage.PolarisStorageIntegration;
7486
import org.apache.polaris.core.storage.PolarisStorageIntegrationProvider;
7587
import org.apache.polaris.core.storage.StorageLocation;
88+
import org.apache.polaris.persistence.relational.jdbc.models.Converter;
7689
import org.apache.polaris.persistence.relational.jdbc.models.EntityNameLookupRecordConverter;
7790
import org.apache.polaris.persistence.relational.jdbc.models.ModelCommitMetricsReport;
7891
import org.apache.polaris.persistence.relational.jdbc.models.ModelEntity;
@@ -88,7 +101,8 @@
88101
import org.slf4j.Logger;
89102
import org.slf4j.LoggerFactory;
90103

91-
public class JdbcBasePersistenceImpl implements BasePersistence, IntegrationPersistence {
104+
public class JdbcBasePersistenceImpl
105+
implements BasePersistence, IntegrationPersistence, LineagePersistence {
92106

93107
private static final Logger LOGGER = LoggerFactory.getLogger(JdbcBasePersistenceImpl.class);
94108

@@ -1365,41 +1379,232 @@ private void writeCommitMetricsReport(@Nonnull ModelCommitMetricsReport report)
13651379
// ============================================================================
13661380

13671381
@Override
1368-
public void upsertLineageDataset(@Nonnull LineageDatasetRecord record) {
1382+
public void upsertDatasets(RealmContext realmContext, List<LineageDataset> datasets) {
13691383
verifyLineagePersistenceSupported();
1370-
ModelLineageDataset model = ModelLineageDataset.fromRecord(record, realmId);
1384+
String realmId = realmContext.getRealmIdentifier();
1385+
long nowMillis = Instant.now().toEpochMilli();
1386+
for (LineageDataset dataset : datasets) {
1387+
long datasetId =
1388+
lookupLineageDataset(realmId, dataset)
1389+
.map(ModelLineageDataset::getDatasetId)
1390+
.orElseGet(IdGenerator.getIdGenerator()::nextId);
1391+
ModelLineageDataset model =
1392+
ModelLineageDataset.fromDataset(dataset, realmId, datasetId, nowMillis);
1393+
try {
1394+
datasourceOperations.executeUpdate(generateLineageDatasetUpsert(model));
1395+
} catch (SQLException e) {
1396+
throw new RuntimeException(
1397+
String.format("Failed to upsert lineage dataset due to %s", e.getMessage()), e);
1398+
}
1399+
}
1400+
}
1401+
1402+
@Override
1403+
public void upsertDatasetEdges(
1404+
RealmContext realmContext, List<LineageEdge> edges, Instant lastEventAt) {
1405+
verifyLineagePersistenceSupported();
1406+
String realmId = realmContext.getRealmIdentifier();
1407+
long lastEventAtMillis = lastEventAt.toEpochMilli();
1408+
for (LineageEdge edge : edges) {
1409+
ModelLineageEdge model =
1410+
ModelLineageEdge.fromIds(
1411+
realmId,
1412+
requireLineageDatasetId(realmId, edge.source()),
1413+
requireLineageDatasetId(realmId, edge.target()),
1414+
lastEventAtMillis);
1415+
try {
1416+
datasourceOperations.executeUpdate(generateLineageEdgeUpsert(model));
1417+
} catch (SQLException e) {
1418+
throw new RuntimeException(
1419+
String.format("Failed to upsert lineage edge due to %s", e.getMessage()), e);
1420+
}
1421+
}
1422+
}
1423+
1424+
@Override
1425+
public void upsertColumnEdges(
1426+
RealmContext realmContext, List<LineageColumnEdge> columnEdges, Instant lastEventAt) {
1427+
verifyLineagePersistenceSupported();
1428+
String realmId = realmContext.getRealmIdentifier();
1429+
long lastEventAtMillis = lastEventAt.toEpochMilli();
1430+
for (LineageColumnEdge columnEdge : columnEdges) {
1431+
ModelLineageColumnEdge model =
1432+
ModelLineageColumnEdge.fromIds(
1433+
realmId,
1434+
requireLineageDatasetId(realmId, columnEdge.source().dataset()),
1435+
columnEdge.source().field(),
1436+
requireLineageDatasetId(realmId, columnEdge.target().dataset()),
1437+
columnEdge.target().field(),
1438+
lastEventAtMillis);
1439+
try {
1440+
datasourceOperations.executeUpdate(generateLineageColumnEdgeUpsert(model));
1441+
} catch (SQLException e) {
1442+
throw new RuntimeException(
1443+
String.format("Failed to upsert lineage column edge due to %s", e.getMessage()), e);
1444+
}
1445+
}
1446+
}
1447+
1448+
@Override
1449+
public LineageGraph loadLineage(RealmContext realmContext, LineageQueryRequest request) {
1450+
verifyLineagePersistenceSupported();
1451+
String realmId = realmContext.getRealmIdentifier();
1452+
Optional<ModelLineageDataset> requested =
1453+
lookupLineageDatasetByNodeId(realmId, request.nodeId());
1454+
if (requested.isEmpty()) {
1455+
return new LineageGraph(
1456+
new LineageNode(request.nodeId(), LineageNodeType.DATASET, null, true),
1457+
List.of(),
1458+
List.of());
1459+
}
1460+
1461+
ModelLineageDataset dataset = requested.get();
1462+
boolean includeColumns = request.granularity() == LineageGranularity.COLUMN;
1463+
List<LineageNode> upstream =
1464+
request.direction() == LineageDirection.UPSTREAM
1465+
|| request.direction() == LineageDirection.BOTH
1466+
? loadAdjacentLineageNodes(realmId, dataset.getDatasetId(), true, includeColumns)
1467+
: List.of();
1468+
List<LineageNode> downstream =
1469+
request.direction() == LineageDirection.DOWNSTREAM
1470+
|| request.direction() == LineageDirection.BOTH
1471+
? loadAdjacentLineageNodes(realmId, dataset.getDatasetId(), false, includeColumns)
1472+
: List.of();
1473+
1474+
return new LineageGraph(toLineageNode(dataset, List.of()), upstream, downstream);
1475+
}
1476+
1477+
private Optional<ModelLineageDataset> lookupLineageDataset(
1478+
String realmId, LineageDataset dataset) {
1479+
String table = QueryGenerator.getFullyQualifiedTableName(ModelLineageDataset.TABLE_NAME);
1480+
PreparedQuery query =
1481+
new PreparedQuery(
1482+
"SELECT realm_id, dataset_id, catalog, namespace, name, polaris_entity_id, created_at, updated_at "
1483+
+ "FROM "
1484+
+ table
1485+
+ " WHERE realm_id = ? AND namespace = ? AND name = ?",
1486+
List.of(realmId, dataset.namespace(), dataset.name()));
13711487
try {
1372-
datasourceOperations.executeUpdate(generateLineageDatasetUpsert(model));
1488+
List<ModelLineageDataset> results =
1489+
datasourceOperations.executeSelect(query, ModelLineageDataset.CONVERTER);
1490+
return results.stream().findFirst();
13731491
} catch (SQLException e) {
13741492
throw new RuntimeException(
1375-
String.format("Failed to upsert lineage dataset due to %s", e.getMessage()), e);
1493+
String.format("Failed to load lineage dataset due to %s", e.getMessage()), e);
13761494
}
13771495
}
13781496

1379-
@Override
1380-
public void upsertLineageEdge(@Nonnull LineageEdgeRecord record) {
1381-
verifyLineagePersistenceSupported();
1382-
ModelLineageEdge model = ModelLineageEdge.fromRecord(record, realmId);
1497+
private Optional<ModelLineageDataset> lookupLineageDatasetByNodeId(
1498+
String realmId, String nodeId) {
1499+
return parseLineageDatasetNodeId(nodeId)
1500+
.flatMap(dataset -> lookupLineageDataset(realmId, dataset));
1501+
}
1502+
1503+
private Optional<LineageDataset> parseLineageDatasetNodeId(String nodeId) {
1504+
String prefix = "dataset:";
1505+
if (!nodeId.startsWith(prefix)) {
1506+
return Optional.empty();
1507+
}
1508+
String identity = nodeId.substring(prefix.length());
1509+
int catalogSeparator = identity.lastIndexOf(':');
1510+
int nameSeparator = identity.lastIndexOf('.');
1511+
if (catalogSeparator < 0 || nameSeparator <= catalogSeparator + 1) {
1512+
return Optional.empty();
1513+
}
1514+
String catalog = identity.substring(0, catalogSeparator);
1515+
String namespace = identity.substring(catalogSeparator + 1, nameSeparator);
1516+
String name = identity.substring(nameSeparator + 1);
1517+
return Optional.of(new LineageDataset(catalog, namespace, name));
1518+
}
1519+
1520+
private long requireLineageDatasetId(String realmId, LineageDataset dataset) {
1521+
return lookupLineageDataset(realmId, dataset)
1522+
.map(ModelLineageDataset::getDatasetId)
1523+
.orElseThrow(
1524+
() ->
1525+
new IllegalArgumentException(
1526+
String.format(
1527+
"Lineage dataset '%s.%s' does not exist in realm '%s'. Call upsertDatasets before writing edges.",
1528+
dataset.namespace(), dataset.name(), realmId)));
1529+
}
1530+
1531+
private List<LineageNode> loadAdjacentLineageNodes(
1532+
String realmId, long datasetId, boolean upstream, boolean includeColumns) {
1533+
String edgesTable = QueryGenerator.getFullyQualifiedTableName(ModelLineageEdge.TABLE_NAME);
1534+
String datasetsTable =
1535+
QueryGenerator.getFullyQualifiedTableName(ModelLineageDataset.TABLE_NAME);
1536+
String adjacentEdgeColumn = upstream ? "source_dataset_id" : "target_dataset_id";
1537+
String requestedEdgeColumn = upstream ? "target_dataset_id" : "source_dataset_id";
1538+
PreparedQuery query =
1539+
new PreparedQuery(
1540+
"SELECT d.realm_id, d.dataset_id, d.catalog, d.namespace, d.name, d.polaris_entity_id, d.created_at, d.updated_at "
1541+
+ "FROM "
1542+
+ edgesTable
1543+
+ " e JOIN "
1544+
+ datasetsTable
1545+
+ " d ON d.realm_id = e.realm_id AND d.dataset_id = e."
1546+
+ adjacentEdgeColumn
1547+
+ " WHERE e.realm_id = ? AND e."
1548+
+ requestedEdgeColumn
1549+
+ " = ?",
1550+
List.of(realmId, datasetId));
13831551
try {
1384-
datasourceOperations.executeUpdate(generateLineageEdgeUpsert(model));
1552+
List<ModelLineageDataset> datasets =
1553+
datasourceOperations.executeSelect(query, ModelLineageDataset.CONVERTER);
1554+
return datasets.stream()
1555+
.map(
1556+
dataset ->
1557+
toLineageNode(
1558+
dataset,
1559+
includeColumns
1560+
? loadFieldMappings(realmId, datasetId, dataset.getDatasetId(), upstream)
1561+
: List.of()))
1562+
.toList();
13851563
} catch (SQLException e) {
13861564
throw new RuntimeException(
1387-
String.format("Failed to upsert lineage edge due to %s", e.getMessage()), e);
1565+
String.format("Failed to load adjacent lineage nodes due to %s", e.getMessage()), e);
13881566
}
13891567
}
13901568

1391-
@Override
1392-
public void upsertLineageColumnEdge(@Nonnull LineageColumnEdgeRecord record) {
1393-
verifyLineagePersistenceSupported();
1394-
ModelLineageColumnEdge model = ModelLineageColumnEdge.fromRecord(record, realmId);
1569+
private List<LineageFieldMapping> loadFieldMappings(
1570+
String realmId, long requestedDatasetId, long adjacentDatasetId, boolean upstream) {
1571+
String table = QueryGenerator.getFullyQualifiedTableName(ModelLineageColumnEdge.TABLE_NAME);
1572+
long sourceDatasetId = upstream ? adjacentDatasetId : requestedDatasetId;
1573+
long targetDatasetId = upstream ? requestedDatasetId : adjacentDatasetId;
1574+
PreparedQuery query =
1575+
new PreparedQuery(
1576+
"SELECT source_field, target_field FROM "
1577+
+ table
1578+
+ " WHERE realm_id = ? AND source_dataset_id = ? AND target_dataset_id = ?",
1579+
List.of(realmId, sourceDatasetId, targetDatasetId));
13951580
try {
1396-
datasourceOperations.executeUpdate(generateLineageColumnEdgeUpsert(model));
1581+
return datasourceOperations.executeSelect(query, new LineageFieldMappingConverter());
13971582
} catch (SQLException e) {
13981583
throw new RuntimeException(
1399-
String.format("Failed to upsert lineage column edge due to %s", e.getMessage()), e);
1584+
String.format("Failed to load lineage field mappings due to %s", e.getMessage()), e);
14001585
}
14011586
}
14021587

1588+
private static LineageNode toLineageNode(
1589+
ModelLineageDataset dataset, List<LineageFieldMapping> fieldMappings) {
1590+
LineageData data =
1591+
new LineageData(
1592+
OptionalLong.empty(),
1593+
OptionalLong.of(dataset.getDatasetId()),
1594+
dataset.getNamespace(),
1595+
dataset.getName(),
1596+
null,
1597+
OptionalLong.of(dataset.getCreatedAt()),
1598+
OptionalLong.of(dataset.getUpdatedAt()));
1599+
return new LineageNode(
1600+
lineageNodeId(dataset), LineageNodeType.DATASET, data, false, fieldMappings);
1601+
}
1602+
1603+
private static String lineageNodeId(ModelLineageDataset dataset) {
1604+
return String.format(
1605+
"dataset:%s:%s.%s", dataset.getCatalog(), dataset.getNamespace(), dataset.getName());
1606+
}
1607+
14031608
private void verifyLineagePersistenceSupported() {
14041609
if (schemaVersion >= MIN_LINEAGE_SCHEMA_VERSION) {
14051610
return;
@@ -1516,4 +1721,16 @@ private PreparedQuery generateLineageColumnEdgeUpsert(ModelLineageColumnEdge mod
15161721
+ ".last_event_at, EXCLUDED.last_event_at)",
15171722
values);
15181723
}
1724+
1725+
private static class LineageFieldMappingConverter implements Converter<LineageFieldMapping> {
1726+
@Override
1727+
public LineageFieldMapping fromResultSet(java.sql.ResultSet rs) throws SQLException {
1728+
return new LineageFieldMapping(rs.getString("source_field"), rs.getString("target_field"));
1729+
}
1730+
1731+
@Override
1732+
public Map<String, Object> toMap(DatabaseType databaseType) {
1733+
return Map.of();
1734+
}
1735+
}
15191736
}

‎persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/models/ModelLineageColumnEdge.java‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import java.util.LinkedHashMap;
2424
import java.util.List;
2525
import java.util.Map;
26-
import org.apache.polaris.core.persistence.lineage.LineageColumnEdgeRecord;
2726
import org.apache.polaris.immutables.PolarisImmutable;
2827
import org.apache.polaris.persistence.relational.jdbc.DatabaseType;
2928

@@ -85,14 +84,20 @@ default Map<String, Object> toMap(DatabaseType databaseType) {
8584
return map;
8685
}
8786

88-
static ModelLineageColumnEdge fromRecord(LineageColumnEdgeRecord record, String realmId) {
87+
static ModelLineageColumnEdge fromIds(
88+
String realmId,
89+
long sourceDatasetId,
90+
String sourceField,
91+
long targetDatasetId,
92+
String targetField,
93+
long lastEventAtMillis) {
8994
return ImmutableModelLineageColumnEdge.builder()
9095
.realmId(realmId)
91-
.sourceDatasetId(record.sourceDatasetId())
92-
.sourceField(record.sourceField())
93-
.targetDatasetId(record.targetDatasetId())
94-
.targetField(record.targetField())
95-
.lastEventAt(record.lastEventAt())
96+
.sourceDatasetId(sourceDatasetId)
97+
.sourceField(sourceField)
98+
.targetDatasetId(targetDatasetId)
99+
.targetField(targetField)
100+
.lastEventAt(lastEventAtMillis)
96101
.build();
97102
}
98103

0 commit comments

Comments
 (0)