Support OpenLineage in spark-3.x-bigquery connectors

pom.xml

@@ -44,6 +44,20 @@
  https://github.com/GoogleCloudDataproc/spark-bigquery-connector/issues
  </url>
  </issueManagement>
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-catalyst_2.13</artifactId>
+ <version>3.5.0</version>
+ <scope>compile</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_2.13</artifactId>
+ <version>3.5.0</version>
+ <scope>compile</scope>
+ </dependency>
+ </dependencies>
 
  <modules>
  <module>spark-bigquery-parent</module>

spark-bigquery-dsv2/spark-3.1-bigquery-lib/pom.xml

@@ -16,6 +16,18 @@
  <spark.version>3.1.0</spark.version>
  <shade.skip>true</shade.skip>
  </properties>
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-compiler-plugin</artifactId>
+ <configuration>
+ <source>8</source>
+ <target>8</target>
+ </configuration>
+ </plugin>
+ </plugins>
+ </build>
  <licenses>
  <license>
  <name>Apache License, Version 2.0</name>

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

@@ -15,13 +15,15 @@
  */
 package com.google.cloud.spark.bigquery.v2;
 
+import com.google.cloud.bigquery.connector.common.BigQueryUtil;
 import com.google.cloud.spark.bigquery.DataSourceVersion;
 import com.google.cloud.spark.bigquery.SparkBigQueryConfig;
 import com.google.cloud.spark.bigquery.v2.context.BigQueryDataSourceReaderContext;
 import com.google.cloud.spark.bigquery.v2.context.BigQueryDataSourceReaderModule;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.inject.Injector;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Supplier;
@@ -97,4 +99,14 @@ public WriteBuilder newWriteBuilder(LogicalWriteInfo info) {
  // context
  return new BigQueryWriteBuilder(injector, info, SaveMode.Append);
  }
+
+ @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();
+ }
 }

spark-bigquery-dsv2/spark-bigquery-dsv2-common/pom.xml

@@ -88,6 +88,14 @@
  </execution>
  </executions>
  </plugin>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-compiler-plugin</artifactId>
+ <configuration>
+ <source>8</source>
+ <target>8</target>
+ </configuration>
+ </plugin>
  </plugins>
  </build>
 </project>

spark-bigquery-dsv2/spark-bigquery-dsv2-common/src/main/java/com/google/cloud/spark/bigquery/BigQueryRelationProvider.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed 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
+ *
+ * https://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 com.google.cloud.spark.bigquery;
+
+import static scala.collection.JavaConverters.mapAsJavaMap;
+
+import com.google.common.collect.ImmutableMap;
+import java.util.HashMap;
+import java.util.Optional;
+import org.apache.spark.sql.SQLContext;
+import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.types.StructType;
+import scala.Option;
+import scala.collection.immutable.Map;
+
+/** For OpenLineage */
+public abstract class BigQueryRelationProvider {
+ public SparkBigQueryConfig createSparkBigQueryConfig(
+ SQLContext sqlContext, Map<String, String> options, Option<StructType> schema) {
+ java.util.Map<String, String> optionsMap = new HashMap<>(mapAsJavaMap(options));
+ DataSourceVersion.V1.updateOptionsMap(optionsMap);
+ SparkSession spark = sqlContext.sparkSession();
+ ImmutableMap<String, String> globalOptions =
+ ImmutableMap.copyOf(mapAsJavaMap(spark.conf().getAll()));
+
+ return SparkBigQueryConfig.from(
+ ImmutableMap.copyOf(optionsMap),
+ globalOptions,
+ spark.sparkContext().hadoopConfiguration(),
+ ImmutableMap.of(),
+ 1,
+ spark.sqlContext().conf(),
+ spark.version(),
+ Optional.ofNullable(schema.getOrElse(null)),
+ true);
+ }
+}

spark-bigquery-dsv2/spark-bigquery-dsv2-common/src/main/java/com/google/cloud/spark/bigquery/v2/BaseBigQuerySource.java

@@ -15,9 +15,11 @@
  */
 package com.google.cloud.spark.bigquery.v2;
 
+import com.google.cloud.spark.bigquery.BigQueryRelationProvider;
 import org.apache.spark.sql.sources.DataSourceRegister;
 
-public abstract class BaseBigQuerySource implements DataSourceRegister {
+public abstract class BaseBigQuerySource extends BigQueryRelationProvider
+ implements DataSourceRegister {
 
  @Override
  public String shortName() {