Examples with OpenLineageContext io.openlineage.spark.api.OpenLineageContext used on opensource projects

Example 1 with OpenLineageContext

use of io.openlineage.spark.api.OpenLineageContext in project OpenLineage by OpenLineage.

the class InternalEventHandlerFactory method createInputDatasetBuilder.

@Override
public Collection<PartialFunction<Object, List<InputDataset>>> createInputDatasetBuilder(OpenLineageContext context) {
    ImmutableList builders = ImmutableList.<PartialFunction<Object, List<InputDataset>>>builder().addAll(generate(eventHandlerFactories, factory -> factory.createInputDatasetBuilder(context))).addAll(DatasetBuilderFactoryProvider.getInstance().getInputBuilders(context)).build();
    context.getInputDatasetBuilders().addAll(builders);
    return builders;
}

Example 2 with OpenLineageContext

use of io.openlineage.spark.api.OpenLineageContext in project OpenLineage by OpenLineage.

the class DataSourceV2RelationDatasetBuilderTest method provideBuildersWithSparkListeners.

private static Stream<Arguments> provideBuildersWithSparkListeners() {
    OpenLineageContext context = mock(OpenLineageContext.class);
    DatasetFactory factory = mock(DatasetFactory.class);
    return Stream.of(Arguments.of(new DataSourceV2RelationInputDatasetBuilder(context, factory), mock(SparkListenerJobStart.class), true), Arguments.of(new DataSourceV2RelationInputDatasetBuilder(context, factory), mock(SparkListenerSQLExecutionStart.class), true), Arguments.of(new DataSourceV2RelationInputDatasetBuilder(context, factory), mock(SparkListenerJobEnd.class), false), Arguments.of(new DataSourceV2RelationInputDatasetBuilder(context, factory), mock(SparkListenerSQLExecutionEnd.class), false), Arguments.of(new DataSourceV2RelationOutputDatasetBuilder(context, factory), mock(SparkListenerJobStart.class), false), Arguments.of(new DataSourceV2RelationOutputDatasetBuilder(context, factory), mock(SparkListenerSQLExecutionStart.class), false), Arguments.of(new DataSourceV2RelationOutputDatasetBuilder(context, factory), mock(SparkListenerJobEnd.class), true), Arguments.of(new DataSourceV2RelationOutputDatasetBuilder(context, factory), mock(SparkListenerSQLExecutionEnd.class), true));
}

Example 3 with OpenLineageContext

use of io.openlineage.spark.api.OpenLineageContext in project OpenLineage by OpenLineage.

the class DataSourceV2RelationDatasetBuilderTest method provideBuilders.

private static Stream<Arguments> provideBuilders() {
    OpenLineageContext context = mock(OpenLineageContext.class);
    DatasetFactory factory = mock(DatasetFactory.class);
    OpenLineage openLineage = mock(OpenLineage.class);
    return Stream.of(Arguments.of(new DataSourceV2RelationInputDatasetBuilder(context, factory), mock(DataSourceV2Relation.class), context, factory, openLineage), Arguments.of(new DataSourceV2RelationOutputDatasetBuilder(context, factory), mock(DataSourceV2Relation.class), context, factory, openLineage));
}

Example 4 with OpenLineageContext

use of io.openlineage.spark.api.OpenLineageContext in project OpenLineage by OpenLineage.

the class OpenLineageSparkListenerTest method testSqlEventWithJobEventEmitsOnce.

@Test
public void testSqlEventWithJobEventEmitsOnce() {
    SparkSession sparkSession = mock(SparkSession.class);
    SparkContext sparkContext = mock(SparkContext.class);
    EventEmitter emitter = mock(EventEmitter.class);
    QueryExecution qe = mock(QueryExecution.class);
    LogicalPlan query = UnresolvedRelation$.MODULE$.apply(TableIdentifier.apply("tableName"));
    SparkPlan plan = mock(SparkPlan.class);
    when(sparkSession.sparkContext()).thenReturn(sparkContext);
    when(sparkContext.appName()).thenReturn("appName");
    when(qe.optimizedPlan()).thenReturn(new InsertIntoHadoopFsRelationCommand(new Path("file:///tmp/dir"), null, false, Seq$.MODULE$.empty(), Option.empty(), null, Map$.MODULE$.empty(), query, SaveMode.Overwrite, Option.empty(), Option.empty(), Seq$.MODULE$.<String>empty()));
    when(qe.executedPlan()).thenReturn(plan);
    when(plan.sparkContext()).thenReturn(sparkContext);
    when(plan.nodeName()).thenReturn("execute");
    OpenLineageContext olContext = OpenLineageContext.builder().sparkSession(Optional.of(sparkSession)).sparkContext(sparkSession.sparkContext()).openLineage(new OpenLineage(OpenLineageClient.OPEN_LINEAGE_CLIENT_URI)).queryExecution(qe).build();
    olContext.getOutputDatasetQueryPlanVisitors().add(new InsertIntoHadoopFsRelationVisitor(olContext));
    ExecutionContext executionContext = new StaticExecutionContextFactory(emitter).createSparkSQLExecutionContext(1L, emitter, qe, olContext);
    executionContext.start(new SparkListenerSQLExecutionStart(1L, "", "", "", new SparkPlanInfo("name", "string", Seq$.MODULE$.empty(), Map$.MODULE$.empty(), Seq$.MODULE$.empty()), 1L));
    executionContext.start(new SparkListenerJobStart(0, 2L, Seq$.MODULE$.<StageInfo>empty(), new Properties()));
    ArgumentCaptor<OpenLineage.RunEvent> lineageEvent = ArgumentCaptor.forClass(OpenLineage.RunEvent.class);
    verify(emitter, times(2)).emit(lineageEvent.capture());
}

Example 5 with OpenLineageContext

use of io.openlineage.spark.api.OpenLineageContext in project OpenLineage by OpenLineage.

the class InternalEventHandlerFactory method createRunFacetBuilders.

@Override
public Collection<CustomFacetBuilder<?, ? extends RunFacet>> createRunFacetBuilders(OpenLineageContext context) {
    Builder<CustomFacetBuilder<?, ? extends RunFacet>> listBuilder;
    listBuilder = ImmutableList.<CustomFacetBuilder<?, ? extends RunFacet>>builder().addAll(generate(eventHandlerFactories, factory -> factory.createRunFacetBuilders((context)))).add(new ErrorFacetBuilder(), new LogicalPlanRunFacetBuilder(context), new SparkVersionFacetBuilder(context));
    if (DatabricksEnvironmentFacetBuilder.isDatabricksRuntime()) {
        listBuilder.add(new DatabricksEnvironmentFacetBuilder(context));
    }
    return listBuilder.build();
}