Publish Lineage Information from Custom Applications to OCI Data Catalog

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>org.example</groupId>
    <artifactId>OpenLineageExtension</artifactId>
    <version>1.0-SNAPSHOT</version>

    <properties>
        <maven.compiler.source>11</maven.compiler.source>
        <maven.compiler.target>11</maven.compiler.target>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
        <openlineage.version>1.8.0</openlineage.version>
        <oci.sdk.version>3.41.2</oci.sdk.version>
    </properties>

    <dependencies>
        <!-- https://mvnrepository.com/artifact/io.openlineage/openlineage-spark -->
        <dependency>
            <groupId>io.openlineage</groupId>
            <artifactId>openlineage-spark</artifactId>
            <version>${openlineage.version}</version>
        </dependency>

        <!-- https://mvnrepository.com/artifact/org.projectlombok/lombok -->
        <dependency>
            <groupId>org.projectlombok</groupId>
            <artifactId>lombok</artifactId>
            <version>1.18.30</version>
        </dependency>

        <!-- https://mvnrepository.com/artifact/org.apache.spark/spark-sql -->
        <dependency>
            <groupId>org.apache.spark</groupId>
            <artifactId>spark-sql_2.13</artifactId>
            <version>3.5.0</version>
        </dependency>

        <!-- https://mvnrepository.com/artifact/org.scala-lang/scala-library -->
        <dependency>
            <groupId>org.scala-lang</groupId>
            <artifactId>scala-library</artifactId>
            <version>2.13.12</version>
        </dependency>

        <dependency>
            <groupId>com.oracle.oci.sdk</groupId>
            <artifactId>oci-java-sdk-datacatalog</artifactId>
            <version>${oci.sdk.version}</version>
        </dependency>

        <dependency>
            <groupId>com.oracle.oci.sdk</groupId>
            <artifactId>oci-java-sdk-common</artifactId>
            <version>${oci.sdk.version}</version>
        </dependency>

        <dependency>
            <groupId>com.oracle.oci.sdk</groupId>
            <artifactId>oci-java-sdk-common-httpclient-jersey</artifactId>
            <version>${oci.sdk.version}</version>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-compiler-plugin</artifactId>
                <version>2.0.2</version>
                <configuration>
                    <source>11</source>
                    <target>11</target>
                </configuration>
            </plugin>
        </plugins>
    </build>
</project>

package io.openlineage.client.transports;

import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.Setter;
import lombok.ToString;
import lombok.With;

@NoArgsConstructor
@AllArgsConstructor
@ToString
@With
public final class OciConfig implements TransportConfig {

    @Getter
    @Setter
    private String catalogId;
    @Getter
    @Setter
    private String dataAssetKey;
    @Getter
    @Setter
    private String authType;
    @Getter
    @Setter
    private String authProfile;
    @Getter
    @Setter
    private String endpoint;
}

package io.openlineage.client.transports;

import com.oracle.bmc.auth.AuthenticationDetailsProvider;
import com.oracle.bmc.auth.ConfigFileAuthenticationDetailsProvider;
import com.oracle.bmc.auth.SessionTokenAuthenticationDetailsProvider;
import com.oracle.bmc.datacatalog.DataCatalogClient;
import com.oracle.bmc.datacatalog.model.ImportLineageDetails;
import com.oracle.bmc.datacatalog.requests.ImportLineageRequest;
import com.oracle.bmc.datacatalog.responses.ImportLineageResponse;
import io.openlineage.client.OpenLineage;
import io.openlineage.client.OpenLineageClientUtils;
import lombok.SneakyThrows;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class OciTransport extends Transport {

    private static final Logger log = LoggerFactory.getLogger(OciTransport.class);
    private final OciConfig config;
    private final String SESSION_TOKEN_AUTH = "security_token";
    private final boolean isValidConfig;

    public OciTransport(final OciConfig ociConfig) {
        config = ociConfig;
        isValidConfig = isValidConfig();
    }

    /***
     * Selectively emit only complete events to OCI Data Catalog
     * @param runEvent
     */
    @Override
    public void emit(OpenLineage.RunEvent runEvent) {
        if (runEvent.getEventType().equals(OpenLineage.RunEvent.EventType.COMPLETE)) {
            log.info("Found event type - Complete. Emitting event using OciTransport");
            final String eventAsJson = OpenLineageClientUtils.toJson(runEvent);
            emit(eventAsJson);
        } else {
            log.info("Found event type - {}. Skipping emission",runEvent.getEventType());
        }
    }

    @SneakyThrows
    @Override
    public void emit(String eventAsJson) {
        if (!isValidConfig) {
            log.info("OCI config is not valid. Please update OCI config. Skipping lineage emission.");
            return;
        }
        AuthenticationDetailsProvider provider;
        if (config.getAuthType()!=null && config.getAuthType().equals(SESSION_TOKEN_AUTH)) {
            provider = new SessionTokenAuthenticationDetailsProvider(config.getAuthProfile());
        } else {
            provider = new ConfigFileAuthenticationDetailsProvider(config.getAuthProfile());
        }

        try (DataCatalogClient client = DataCatalogClient.builder().build(provider)) {
            if (config.getEndpoint() != null && !config.getEndpoint().isBlank()) {
                client.setEndpoint(config.getEndpoint());
            }
            ImportLineageRequest importLineageRequest = ImportLineageRequest
                .builder()
                .importLineageDetails(ImportLineageDetails
                    .builder()
                    .lineagePayload(eventAsJson.getBytes())
                    .build())
                .dataAssetKey(config.getDataAssetKey())
                .catalogId(config.getCatalogId())
                .build();
            ImportLineageResponse response = client.importLineage(importLineageRequest);
            log.info("Pushed lineage event to catalog - {}, opc-request-id = {}", config.getCatalogId(), response.getOpcRequestId());
        } catch (Exception ex) {
            log.warn("Failed to push lineage {}", ex.getMessage(), ex);
        }
    }

    private boolean isValidConfig() {
        log.info("OCI Config - {}", config.toString());
        if (config.getAuthProfile() == null || config.getAuthProfile().isBlank()) {
            log.info("OCI config is missing Authentication profile. Lineage will not be emitted to OCI Data Catalog. " +
                "Check documentation on how to create Authentication profile.");
            return false;
        }
        if (config.getCatalogId() == null || config.getCatalogId().isBlank()) {
            log.info("OCI config is missing CatalogId. Lineage will not be emitted to OCI Data Catalog. " +
                "Check documentation on how to get CatalogId.");
            return false;
        }
        if (config.getDataAssetKey() == null || config.getDataAssetKey().isBlank()) {
            log.info("OCI config is missing DataAssetResourceId. Lineage will not be emitted to OCI Data Catalog. " +
                "Check documentation on how to create DataAsset and get its DataAssetResourceId");
            return false;
        }
        return true;
    }
}

package io.openlineage.client.transports;

public class OciTransportBuilder implements TransportBuilder {

    @Override
    public String getType() {
        return "oci";
    }

    @Override
    public TransportConfig getConfig() {
        return new OciConfig();
    }

    @Override
    public Transport build(TransportConfig config) {
        return new OciTransport((OciConfig) config);
    }
}

--packages "io.openlineage:openlineage-spark:1.8.0,com.oracle.oci.sdk:oci-java-sdk-common:{oci-sdk-version},com.oracle.oci.sdk:oci-java-sdk-common-httpclient-jersey:{oci-sdk-version},com.oracle.pic.dcat:datacatalog-java-client:{oci-sdk-version},io.openlineage:openlineage-oci-extension:{generated-from-above-mentioned-code-snippet}"
--conf "spark.extraListeners=io.openlineage.spark.agent.OpenLineageSparkListener"
--conf "spark.openlineage.debugFacet=enabled"
--conf "spark.openlineage.transport.type=oci"
--conf "spark.openlineage.transport.catalogId={catalog instance OCID}"
--conf "spark.openlineage.transport.dataAssetKey={UUID of the dataAsset registered in Data catalog}"
--conf "spark.openlineage.transport.authProfile={user authentication profile}"
--conf "spark.openlineage.transport.authType={optionally specify security_token as auth type}"
--conf "spark.openlineage.application.name={application name to display on catalog UI}"