Get lineage out of query

Signed-off-by: Dominik Dębowczyk <[email protected]>

Author: ddebowczyk92

spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/integration/OpenLineageIntegrationTestBase.java

@@ -18,6 +18,8 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Scanner;
@@ -26,6 +28,7 @@
 import org.apache.spark.sql.SaveMode;
 import org.apache.spark.sql.SparkSession;
 import org.json.JSONObject;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.ClassRule;
 import org.junit.Test;
@@ -50,6 +53,12 @@ public void createTestTable() {
     testTable = "test_" + System.nanoTime();
   }
 
+  @After
+  public void clearLineageFile() throws FileNotFoundException {
+    PrintWriter pw = new PrintWriter(lineageFile);
+    pw.close();
+  }
+
   protected static class CustomSessionFactory extends ExternalResource {
     SparkSession spark;
     File lineageFile;
@@ -73,6 +82,9 @@ protected void before() throws Throwable {
   }
 
   private List<JSONObject> parseEventLogs(File file) throws Exception {
+    // Adding a 5-second cushion, as it takes some time for all OpenLineage events to be emitted and
+    // saved to the file
+    Thread.sleep(5 * 1000);
     List<JSONObject> eventList;
     try (Scanner scanner = new Scanner(file)) {
       eventList = new ArrayList<>();
@@ -117,4 +129,37 @@ public void testLineageEvent() throws Exception {
           assertThat(getFieldName(event, "outputs")).matches(fullTableName);
         });
   }
+
+  @Test
+  public void testLineageEventWithQueryInput() throws Exception {
+    String fullTableName = testDataset.toString() + "." + testTable;
+    Dataset<Row> readDF =
+        spark
+            .read()
+            .format("bigquery")
+            .option("viewsEnabled", true)
+            .option("materializationDataset", testDataset.toString())
+            .option("query", "SELECT * FROM `bigquery-public-data.samples.shakespeare`")
+            .load();
+
+    readDF.createOrReplaceTempView("words");
+    Dataset<Row> writeDF =
+        spark.sql("SELECT word, SUM(word_count) AS word_count FROM words GROUP BY word");
+    writeDF
+        .write()
+        .format("bigquery")
+        .mode(SaveMode.Append)
+        .option("table", fullTableName)
+        .option("temporaryGcsBucket", TestConstants.TEMPORARY_GCS_BUCKET)
+        .option("writeMethod", "direct")
+        .save();
+    List<JSONObject> eventList = parseEventLogs(lineageFile);
+    assertThat(eventList)
+        .isNotEmpty(); // check if there is at least one event with both input and output
+    eventList.forEach(
+        (event) -> { // check if each of these events have the correct input and output
+          assertThat(getFieldName(event, "inputs")).matches(TestConstants.SHAKESPEARE_TABLE);
+          assertThat(getFieldName(event, "outputs")).matches(fullTableName);
+        });
+  }
 }

spark-bigquery-dsv2/spark-3.1-bigquery-lib/src/main/java/com/google/cloud/spark/bigquery/v2/Spark31BigQueryTable.java

@@ -109,10 +109,16 @@ public WriteBuilder newWriteBuilder(LogicalWriteInfo info) {
   @Override
   public Map<String, String> properties() {
     SparkBigQueryConfig config = injector.getInstance(SparkBigQueryConfig.class);
-    return ImmutableMap.<String, String>builder()
-        .put("openlineage.dataset.name", BigQueryUtil.friendlyTableName(config.getTableId()))
-        .put("openlineage.dataset.namespace", "bigquery")
-        .put("openlineage.dataset.storageDatasetFacet.storageLayer", "bigquery")
-        .build();
+    ImmutableMap.Builder<String, String> propertiesBuilder =
+        ImmutableMap.<String, String>builder()
+            .put("openlineage.dataset.namespace", "bigquery")
+            .put("openlineage.dataset.storageDatasetFacet.storageLayer", "bigquery");
+    if (config.getQuery().isPresent()) {
+      propertiesBuilder.put("openlineage.dataset.query", config.getQuery().get());
+    } else {
+      propertiesBuilder.put(
+          "openlineage.dataset.name", BigQueryUtil.friendlyTableName(config.getTableId()));
+    }
+    return propertiesBuilder.build();
   }
 }

spark-bigquery-parent/pom.xml

@@ -66,7 +66,7 @@
         <guava.version>33.4.0-jre</guava.version>
         <jackson.version>2.18.2</jackson.version>
         <netty.version>4.1.119.Final</netty.version>
-        <openlineage-spark.version>1.27.0</openlineage-spark.version>
+        <openlineage-spark.version>1.31.0</openlineage-spark.version>
         <paranamer.version>2.8</paranamer.version>
         <!-- Don't upgrade protobuf until bigquerystorage is compatible with 4.x version of proto.
         Ref : https://github.com/protocolbuffers/protobuf/issues/16452-->