[FLINK-34467] add lineage integration for jdbc connector

apache · Dec 1, 2024 · 76f3e11 · 76f3e11
1 parent 134d858
commit 76f3e11
Show file tree

Hide file tree

Showing 29 changed files with 1,236 additions and 13 deletions.
diff --git a/.java-version b/.java-version
@@ -0,0 +1 @@
+11
diff --git a/flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/JdbcInputFormat.java b/flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/JdbcInputFormat.java
@@ -24,17 +24,23 @@
 import org.apache.flink.api.common.io.InputFormat;
 import org.apache.flink.api.common.io.RichInputFormat;
 import org.apache.flink.api.common.io.statistics.BaseStatistics;
+import org.apache.flink.api.connector.source.Boundedness;
 import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
 import org.apache.flink.api.java.typeutils.RowTypeInfo;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.connector.jdbc.core.datastream.source.JdbcSource;
 import org.apache.flink.connector.jdbc.core.datastream.source.JdbcSourceBuilder;
 import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider;
 import org.apache.flink.connector.jdbc.datasource.connections.SimpleJdbcConnectionProvider;
+import org.apache.flink.connector.jdbc.lineage.DefaultTypeDatasetFacet;
+import org.apache.flink.connector.jdbc.lineage.LineageUtils;
 import org.apache.flink.connector.jdbc.split.JdbcParameterValuesProvider;
 import org.apache.flink.core.io.GenericInputSplit;
 import org.apache.flink.core.io.InputSplit;
 import org.apache.flink.core.io.InputSplitAssigner;
+import org.apache.flink.streaming.api.lineage.LineageDataset;
+import org.apache.flink.streaming.api.lineage.LineageVertex;
+import org.apache.flink.streaming.api.lineage.LineageVertexProvider;
 import org.apache.flink.types.Row;
 import org.apache.flink.util.Preconditions;
 
@@ -53,6 +59,7 @@
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.util.Arrays;
+import java.util.Collections;
 
 /**
  * InputFormat to read data from a database and generate Rows. The InputFormat has to be configured
@@ -107,7 +114,7 @@
 @Deprecated
 @Experimental
 public class JdbcInputFormat extends RichInputFormat<Row, InputSplit>
-        implements ResultTypeQueryable<Row> {
+        implements LineageVertexProvider, ResultTypeQueryable<Row> {
 
     protected static final long serialVersionUID = 2L;
     protected static final Logger LOG = LoggerFactory.getLogger(JdbcInputFormat.class);
@@ -344,6 +351,18 @@ public static JdbcInputFormatBuilder buildJdbcInputFormat() {
         return new JdbcInputFormatBuilder();
     }
 
+    @Override
+    public LineageVertex getLineageVertex() {
+        DefaultTypeDatasetFacet defaultTypeDatasetFacet =
+                new DefaultTypeDatasetFacet(getProducedType());
+        String name = LineageUtils.nameOf(connectionProvider);
+        String namespace = LineageUtils.namespaceOf(connectionProvider);
+        LineageDataset dataset =
+                LineageUtils.datasetOf(name, namespace, Arrays.asList(defaultTypeDatasetFacet));
+        return LineageUtils.sourceLineageVertexOf(
+                Boundedness.BOUNDED, Collections.singleton(dataset));
+    }
+
     /** Builder for {@link JdbcInputFormat}. */
     public static class JdbcInputFormatBuilder {
         private final JdbcConnectionOptions.JdbcConnectionOptionsBuilder connOptionsBuilder;

diff --git a/...dbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/sink/JdbcSink.java b/...dbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/sink/JdbcSink.java
@@ -34,7 +34,11 @@
 import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider;
 import org.apache.flink.connector.jdbc.datasource.statements.JdbcQueryStatement;
 import org.apache.flink.connector.jdbc.internal.JdbcOutputSerializer;
+import org.apache.flink.connector.jdbc.lineage.LineageUtils;
 import org.apache.flink.core.io.SimpleVersionedSerializer;
+import org.apache.flink.streaming.api.lineage.LineageDataset;
+import org.apache.flink.streaming.api.lineage.LineageVertex;
+import org.apache.flink.streaming.api.lineage.LineageVertexProvider;
 
 import java.io.IOException;
 import java.util.Collection;
@@ -47,7 +51,9 @@
  */
 @PublicEvolving
 public class JdbcSink<IN>
-        implements StatefulSink<IN, JdbcWriterState>, TwoPhaseCommittingSink<IN, JdbcCommitable> {
+        implements LineageVertexProvider,
+                StatefulSink<IN, JdbcWriterState>,
+                TwoPhaseCommittingSink<IN, JdbcCommitable> {
 
     private final DeliveryGuarantee deliveryGuarantee;
     private final JdbcConnectionProvider connectionProvider;
@@ -113,4 +119,12 @@ public SimpleVersionedSerializer<JdbcCommitable> getCommittableSerializer() {
     public SimpleVersionedSerializer<JdbcWriterState> getWriterStateSerializer() {
         return new JdbcWriterStateSerializer();
     }
+
+    @Override
+    public LineageVertex getLineageVertex() {
+        String name = LineageUtils.nameOf(queryStatement);
+        String namespace = LineageUtils.namespaceOf(connectionProvider);
+        LineageDataset dataset = LineageUtils.datasetOf(name, namespace, Collections.emptyList());
+        return LineageUtils.lineageVertexOf(Collections.singleton(dataset));
+    }
 }
diff --git a/...core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/JdbcSource.java b/...core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/JdbcSource.java
@@ -40,20 +40,28 @@
 import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit;
 import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplitSerializer;
 import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider;
+import org.apache.flink.connector.jdbc.lineage.DefaultTypeDatasetFacet;
+import org.apache.flink.connector.jdbc.lineage.LineageUtils;
 import org.apache.flink.connector.jdbc.utils.ContinuousUnBoundingSettings;
 import org.apache.flink.core.io.SimpleVersionedSerializer;
+import org.apache.flink.streaming.api.lineage.LineageDataset;
+import org.apache.flink.streaming.api.lineage.LineageVertex;
+import org.apache.flink.streaming.api.lineage.LineageVertexProvider;
 import org.apache.flink.util.Preconditions;
 
 import javax.annotation.Nullable;
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.Objects;
 
 /** JDBC source. */
 @PublicEvolving
 public class JdbcSource<OUT>
-        implements Source<OUT, JdbcSourceSplit, JdbcSourceEnumeratorState>,
+        implements LineageVertexProvider,
+                Source<OUT, JdbcSourceSplit, JdbcSourceEnumeratorState>,
                 ResultTypeQueryable<OUT> {
 
     private final Boundedness boundedness;
@@ -195,4 +203,15 @@ public boolean equals(Object o) {
                 && deliveryGuarantee == that.deliveryGuarantee
                 && Objects.equals(continuousUnBoundingSettings, that.continuousUnBoundingSettings);
     }
+
+    @Override
+    public LineageVertex getLineageVertex() {
+        DefaultTypeDatasetFacet defaultTypeDatasetFacet =
+                new DefaultTypeDatasetFacet(getTypeInformation());
+        String name = LineageUtils.nameOf(connectionProvider);
+        String namespace = LineageUtils.namespaceOf(connectionProvider);
+        LineageDataset dataset =
+                LineageUtils.datasetOf(name, namespace, Arrays.asList(defaultTypeDatasetFacet));
+        return LineageUtils.sourceLineageVertexOf(boundedness, Collections.singleton(dataset));
+    }
 }
diff --git a/...c/main/java/org/apache/flink/connector/jdbc/core/table/source/JdbcRowDataInputFormat.java b/...c/main/java/org/apache/flink/connector/jdbc/core/table/source/JdbcRowDataInputFormat.java
@@ -24,16 +24,22 @@
 import org.apache.flink.api.common.io.RichInputFormat;
 import org.apache.flink.api.common.io.statistics.BaseStatistics;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.connector.source.Boundedness;
 import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.connector.jdbc.JdbcConnectionOptions;
 import org.apache.flink.connector.jdbc.core.database.dialect.JdbcDialectConverter;
 import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider;
 import org.apache.flink.connector.jdbc.datasource.connections.SimpleJdbcConnectionProvider;
+import org.apache.flink.connector.jdbc.lineage.DefaultTypeDatasetFacet;
+import org.apache.flink.connector.jdbc.lineage.LineageUtils;
 import org.apache.flink.connector.jdbc.split.JdbcParameterValuesProvider;
 import org.apache.flink.core.io.GenericInputSplit;
 import org.apache.flink.core.io.InputSplit;
 import org.apache.flink.core.io.InputSplitAssigner;
+import org.apache.flink.streaming.api.lineage.LineageDataset;
+import org.apache.flink.streaming.api.lineage.LineageVertex;
+import org.apache.flink.streaming.api.lineage.LineageVertexProvider;
 import org.apache.flink.table.data.RowData;
 import org.apache.flink.util.Preconditions;
 
@@ -51,11 +57,12 @@
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.util.Arrays;
+import java.util.Collections;
 
 /** InputFormat for {@link JdbcDynamicTableSource}. */
 @Internal
 public class JdbcRowDataInputFormat extends RichInputFormat<RowData, InputSplit>
-        implements ResultTypeQueryable<RowData> {
+        implements LineageVertexProvider, ResultTypeQueryable<RowData> {
 
     private static final long serialVersionUID = 2L;
     private static final Logger LOG = LoggerFactory.getLogger(JdbcRowDataInputFormat.class);
@@ -296,6 +303,18 @@ public static Builder builder() {
         return new Builder();
     }
 
+    @Override
+    public LineageVertex getLineageVertex() {
+        DefaultTypeDatasetFacet defaultTypeDatasetFacet =
+                new DefaultTypeDatasetFacet(getProducedType());
+        String name = LineageUtils.nameOf(connectionProvider);
+        String namespace = LineageUtils.namespaceOf(connectionProvider);
+        LineageDataset dataset =
+                LineageUtils.datasetOf(name, namespace, Arrays.asList(defaultTypeDatasetFacet));
+        return LineageUtils.sourceLineageVertexOf(
+                Boundedness.BOUNDED, Collections.singleton(dataset));
+    }
+
     /** Builder for {@link JdbcRowDataInputFormat}. */
     public static class Builder {
         private JdbcConnectionOptions.JdbcConnectionOptionsBuilder connOptionsBuilder;

diff --git a/...ain/java/org/apache/flink/connector/jdbc/core/table/source/JdbcRowDataLookupFunction.java b/...ain/java/org/apache/flink/connector/jdbc/core/table/source/JdbcRowDataLookupFunction.java
@@ -20,18 +20,25 @@
 
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.connector.source.Boundedness;
 import org.apache.flink.connector.jdbc.core.database.dialect.JdbcDialect;
 import org.apache.flink.connector.jdbc.core.database.dialect.JdbcDialectConverter;
 import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider;
 import org.apache.flink.connector.jdbc.datasource.connections.SimpleJdbcConnectionProvider;
 import org.apache.flink.connector.jdbc.internal.options.InternalJdbcConnectionOptions;
+import org.apache.flink.connector.jdbc.lineage.DefaultTypeDatasetFacet;
+import org.apache.flink.connector.jdbc.lineage.LineageUtils;
 import org.apache.flink.connector.jdbc.statement.FieldNamedPreparedStatement;
+import org.apache.flink.streaming.api.lineage.LineageDataset;
+import org.apache.flink.streaming.api.lineage.LineageVertex;
+import org.apache.flink.streaming.api.lineage.LineageVertexProvider;
 import org.apache.flink.table.data.RowData;
 import org.apache.flink.table.functions.FunctionContext;
 import org.apache.flink.table.functions.LookupFunction;
 import org.apache.flink.table.types.DataType;
 import org.apache.flink.table.types.logical.LogicalType;
 import org.apache.flink.table.types.logical.RowType;
+import org.apache.flink.table.types.utils.TypeConversions;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -53,7 +60,7 @@
 
 /** A lookup function for {@link JdbcDynamicTableSource}. */
 @Internal
-public class JdbcRowDataLookupFunction extends LookupFunction {
+public class JdbcRowDataLookupFunction extends LookupFunction implements LineageVertexProvider {
 
     private static final Logger LOG = LoggerFactory.getLogger(JdbcRowDataLookupFunction.class);
     private static final long serialVersionUID = 2L;
@@ -67,6 +74,7 @@ public class JdbcRowDataLookupFunction extends LookupFunction {
 
     private final List<String> resolvedPredicates;
     private final Serializable[] pushdownParams;
+    private final RowType producedType;
 
     private transient FieldNamedPreparedStatement statement;
 
@@ -106,12 +114,12 @@ public JdbcRowDataLookupFunction(
                         .getSelectFromStatement(options.getTableName(), fieldNames, keyNames);
         JdbcDialect jdbcDialect = options.getDialect();
         this.jdbcDialectConverter = jdbcDialect.getRowConverter(rowType);
-        this.lookupKeyRowConverter =
-                jdbcDialect.getRowConverter(
-                        RowType.of(
-                                Arrays.stream(keyTypes)
-                                        .map(DataType::getLogicalType)
-                                        .toArray(LogicalType[]::new)));
+        this.producedType =
+                RowType.of(
+                        Arrays.stream(keyTypes)
+                                .map(DataType::getLogicalType)
+                                .toArray(LogicalType[]::new));
+        this.lookupKeyRowConverter = jdbcDialect.getRowConverter(producedType);
         this.resolvedPredicates = resolvedPredicates;
         this.pushdownParams = pushdownParams;
     }
@@ -224,4 +232,18 @@ public void close() throws IOException {
     public Connection getDbConnection() {
         return connectionProvider.getConnection();
     }
+
+    @Override
+    public LineageVertex getLineageVertex() {
+        DefaultTypeDatasetFacet defaultTypeDatasetFacet =
+                new DefaultTypeDatasetFacet(
+                        TypeConversions.fromDataTypeToLegacyInfo(
+                                TypeConversions.fromLogicalToDataType(producedType)));
+        String name = LineageUtils.nameOf(connectionProvider);
+        String namespace = LineageUtils.namespaceOf(connectionProvider);
+        LineageDataset dataset =
+                LineageUtils.datasetOf(name, namespace, Arrays.asList(defaultTypeDatasetFacet));
+        return LineageUtils.sourceLineageVertexOf(
+                Boundedness.BOUNDED, Collections.singleton(dataset));
+    }
 }
diff --git a/.../org/apache/flink/connector/jdbc/datasource/connections/SimpleJdbcConnectionProvider.java b/.../org/apache/flink/connector/jdbc/datasource/connections/SimpleJdbcConnectionProvider.java
@@ -149,4 +149,8 @@ public Connection reestablishConnection() throws SQLException, ClassNotFoundExce
         closeConnection();
         return getOrEstablishConnection();
     }
+
+    public String getDbURL() {
+        return this.jdbcOptions.getDbURL();
+    }
 }
diff --git a/...-core/src/main/java/org/apache/flink/connector/jdbc/internal/GenericJdbcSinkFunction.java b/...-core/src/main/java/org/apache/flink/connector/jdbc/internal/GenericJdbcSinkFunction.java
@@ -23,20 +23,25 @@
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.typeutils.InputTypeConfigurable;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.connector.jdbc.lineage.LineageUtils;
 import org.apache.flink.runtime.state.FunctionInitializationContext;
 import org.apache.flink.runtime.state.FunctionSnapshotContext;
 import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
 import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
+import org.apache.flink.streaming.api.lineage.LineageDataset;
+import org.apache.flink.streaming.api.lineage.LineageVertex;
+import org.apache.flink.streaming.api.lineage.LineageVertexProvider;
 import org.apache.flink.util.Preconditions;
 
 import javax.annotation.Nonnull;
 
 import java.io.IOException;
+import java.util.Collections;
 
 /** A generic SinkFunction for JDBC. */
 @Internal
 public class GenericJdbcSinkFunction<T> extends RichSinkFunction<T>
-        implements CheckpointedFunction, InputTypeConfigurable {
+        implements LineageVertexProvider, CheckpointedFunction, InputTypeConfigurable {
     private final JdbcOutputFormat<T, ?, ?> outputFormat;
     private JdbcOutputSerializer<T> serializer;
 
@@ -78,4 +83,12 @@ public void setInputType(TypeInformation<?> type, ExecutionConfig executionConfi
                 JdbcOutputSerializer.of(
                         ((TypeInformation<T>) type).createSerializer(executionConfig));
     }
+
+    @Override
+    public LineageVertex getLineageVertex() {
+        String name = LineageUtils.nameOf(outputFormat.connectionProvider);
+        String namespace = LineageUtils.namespaceOf(outputFormat.connectionProvider);
+        LineageDataset dataset = LineageUtils.datasetOf(name, namespace, Collections.emptyList());
+        return LineageUtils.lineageVertexOf(Collections.singleton(dataset));
+    }
 }