Examples with JobContext org.apache.samza.context.JobContext used on opensource projects

Example 1 with JobContext

use of org.apache.samza.context.JobContext in project samza by apache.

the class TestLocalTableProvider method testInit.

@Test
public void testInit() {
    Context context = mock(Context.class);
    JobContext jobContext = mock(JobContext.class);
    when(context.getJobContext()).thenReturn(jobContext);
    when(jobContext.getConfig()).thenReturn(new MapConfig());
    ContainerContext containerContext = mock(ContainerContext.class);
    when(context.getContainerContext()).thenReturn(containerContext);
    when(containerContext.getContainerMetricsRegistry()).thenReturn(new NoOpMetricsRegistry());
    TaskContext taskContext = mock(TaskContext.class);
    when(context.getTaskContext()).thenReturn(taskContext);
    when(taskContext.getStore(any())).thenReturn(mock(KeyValueStore.class));
    TableProvider tableProvider = createTableProvider("t1");
    tableProvider.init(context);
    Assert.assertNotNull(tableProvider.getTable());
}

Example 2 with JobContext

use of org.apache.samza.context.JobContext in project samza by apache.

the class KafkaChangelogStateBackendFactory method getRestoreManager.

@Override
public TaskRestoreManager getRestoreManager(JobContext jobContext, ContainerContext containerContext, TaskModel taskModel, ExecutorService restoreExecutor, MetricsRegistry metricsRegistry, Set<String> storesToRestore, Config config, Clock clock, File loggedStoreBaseDir, File nonLoggedStoreBaseDir, KafkaChangelogRestoreParams kafkaChangelogRestoreParams) {
    Map<String, SystemStream> storeChangelogs = new StorageConfig(config).getStoreChangelogs();
    Set<SystemStreamPartition> changelogSSPs = storeChangelogs.values().stream().flatMap(ss -> containerContext.getContainerModel().getTasks().values().stream().map(tm -> new SystemStreamPartition(ss, tm.getChangelogPartition()))).collect(Collectors.toSet());
    // filter out standby store-ssp pairs
    Map<String, SystemStream> filteredStoreChangelogs = filterStandbySystemStreams(storeChangelogs, containerContext.getContainerModel());
    SystemAdmins systemAdmins = new SystemAdmins(kafkaChangelogRestoreParams.getSystemAdmins());
    if (new TaskConfig(config).getTransactionalStateRestoreEnabled()) {
        return new TransactionalStateTaskRestoreManager(storesToRestore, jobContext, containerContext, taskModel, restoreExecutor, filteredStoreChangelogs, kafkaChangelogRestoreParams.getInMemoryStores(), kafkaChangelogRestoreParams.getStorageEngineFactories(), kafkaChangelogRestoreParams.getSerdes(), systemAdmins, kafkaChangelogRestoreParams.getStoreConsumers(), metricsRegistry, kafkaChangelogRestoreParams.getCollector(), getSspCache(systemAdmins, clock, changelogSSPs), loggedStoreBaseDir, nonLoggedStoreBaseDir, config, clock);
    } else {
        return new NonTransactionalStateTaskRestoreManager(storesToRestore, jobContext, containerContext, taskModel, restoreExecutor, filteredStoreChangelogs, kafkaChangelogRestoreParams.getInMemoryStores(), kafkaChangelogRestoreParams.getStorageEngineFactories(), kafkaChangelogRestoreParams.getSerdes(), systemAdmins, getStreamCache(systemAdmins, clock), kafkaChangelogRestoreParams.getStoreConsumers(), metricsRegistry, kafkaChangelogRestoreParams.getCollector(), jobContext.getJobModel().getMaxChangeLogStreamPartitions(), loggedStoreBaseDir, nonLoggedStoreBaseDir, config, clock);
    }
}

Example 3 with JobContext

use of org.apache.samza.context.JobContext in project samza by apache.

the class TestRemoteTableDescriptor method createMockContext.

private Context createMockContext(TableDescriptor tableDescriptor) {
    Context context = mock(Context.class);
    ContainerContext containerContext = mock(ContainerContext.class);
    when(context.getContainerContext()).thenReturn(containerContext);
    MetricsRegistry metricsRegistry = mock(MetricsRegistry.class);
    when(metricsRegistry.newTimer(anyString(), anyString())).thenReturn(mock(Timer.class));
    when(metricsRegistry.newCounter(anyString(), anyString())).thenReturn(mock(Counter.class));
    when(containerContext.getContainerMetricsRegistry()).thenReturn(metricsRegistry);
    TaskContextImpl taskContext = mock(TaskContextImpl.class);
    when(context.getTaskContext()).thenReturn(taskContext);
    TaskName taskName = new TaskName("MyTask");
    TaskModel taskModel = mock(TaskModel.class);
    when(taskModel.getTaskName()).thenReturn(taskName);
    when(context.getTaskContext().getTaskModel()).thenReturn(taskModel);
    ContainerModel containerModel = mock(ContainerModel.class);
    when(containerModel.getTasks()).thenReturn(ImmutableMap.of(taskName, taskModel));
    when(containerContext.getContainerModel()).thenReturn(containerModel);
    String containerId = "container-1";
    JobModel jobModel = mock(JobModel.class);
    when(taskContext.getJobModel()).thenReturn(jobModel);
    when(jobModel.getContainers()).thenReturn(ImmutableMap.of(containerId, containerModel));
    JobContext jobContext = mock(JobContext.class);
    Config jobConfig = new MapConfig(tableDescriptor.toConfig(new MapConfig()));
    when(jobContext.getConfig()).thenReturn(jobConfig);
    when(context.getJobContext()).thenReturn(jobContext);
    return context;
}

Example 4 with JobContext

use of org.apache.samza.context.JobContext in project samza by apache.

the class TestLocalTableRead method createTable.

private LocalTable createTable(boolean isTimerDisabled) {
    Map<String, String> config = new HashMap<>();
    if (isTimerDisabled) {
        config.put(MetricsConfig.METRICS_TIMER_ENABLED, "false");
    }
    Context context = mock(Context.class);
    JobContext jobContext = mock(JobContext.class);
    when(context.getJobContext()).thenReturn(jobContext);
    when(jobContext.getConfig()).thenReturn(new MapConfig(config));
    ContainerContext containerContext = mock(ContainerContext.class);
    when(context.getContainerContext()).thenReturn(containerContext);
    when(containerContext.getContainerMetricsRegistry()).thenReturn(metricsRegistry);
    LocalTable table = new LocalTable("t1", kvStore);
    table.init(context);
    return table;
}

Example 5 with JobContext

use of org.apache.samza.context.JobContext in project samza by apache.

the class QueryTranslator method translate.

/**
 * For unit testing only
 */
@VisibleForTesting
void translate(SamzaSqlQueryParser.QueryInfo queryInfo, StreamApplicationDescriptor appDesc, int queryId) {
    QueryPlanner planner = new QueryPlanner(sqlConfig.getRelSchemaProviders(), sqlConfig.getInputSystemStreamConfigBySource(), sqlConfig.getUdfMetadata(), sqlConfig.isQueryPlanOptimizerEnabled());
    final RelRoot relRoot = planner.plan(queryInfo.getSelectQuery());
    SamzaSqlExecutionContext executionContext = new SamzaSqlExecutionContext(sqlConfig);
    TranslatorContext translatorContext = new TranslatorContext(appDesc, relRoot, executionContext);
    translate(relRoot, sqlConfig.getOutputSystemStreams().get(queryId), translatorContext, queryId);
    Map<Integer, TranslatorContext> translatorContexts = new HashMap<>();
    translatorContexts.put(queryId, translatorContext.clone());
    appDesc.withApplicationTaskContextFactory(new ApplicationTaskContextFactory<SamzaSqlApplicationContext>() {

        @Override
        public SamzaSqlApplicationContext create(ExternalContext externalContext, JobContext jobContext, ContainerContext containerContext, TaskContext taskContext, ApplicationContainerContext applicationContainerContext) {
            return new SamzaSqlApplicationContext(translatorContexts);
        }
    });
}