Examples with WorkUnitState org.apache.gobblin.configuration.WorkUnitState used on opensource projects

Example 91 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class JdbcExtractorTest method testUnsignedInt.

/**
 * Test for the metadata query to see if the check for unsigned int is present
 */
@Test
public void testUnsignedInt() throws SchemaException {
    State state = new WorkUnitState();
    state.setId("id");
    MysqlExtractor mysqlExtractor = new MysqlExtractor((WorkUnitState) state);
    List<Command> commands = mysqlExtractor.getSchemaMetadata("db", "table");
    assertTrue(commands.get(0).getCommandType() == JdbcCommand.JdbcCommandType.QUERY);
    assertTrue(commands.get(0).getParams().get(0).contains("bigint"));
    assertTrue(commands.get(1).getCommandType() == JdbcCommand.JdbcCommandType.QUERYPARAMS);
    assertTrue(!commands.get(1).getParams().get(0).contains("unsigned"));
    // set option to promote unsigned int to bigint
    state.setProp(ConfigurationKeys.SOURCE_QUERYBASED_PROMOTE_UNSIGNED_INT_TO_BIGINT, "true");
    commands = mysqlExtractor.getSchemaMetadata("db", "table");
    assertTrue(commands.get(0).getCommandType() == JdbcCommand.JdbcCommandType.QUERY);
    assertTrue(commands.get(0).getParams().get(0).contains("bigint"));
    assertTrue(commands.get(1).getCommandType() == JdbcCommand.JdbcCommandType.QUERYPARAMS);
    assertTrue(commands.get(1).getParams().get(0).contains("unsigned"));
}

Example 92 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class PostgresqlExtractorTest method setup.

@BeforeClass
public void setup() {
    output = new JdbcCommandOutput();
    try {
        output.put(new JdbcCommand(), buildMockResultSet());
    } catch (Exception e) {
        // hack for test failure
        assertEquals("PostgresqlExtractorTest: error initializing mock result set", "false");
    }
    state = new WorkUnitState();
    state.setId("id");
    postgresqlExtractor = new PostgresqlExtractor((WorkUnitState) state);
}

Example 93 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class JsonIntermediateToParquetGroupConverterTest method setUp.

@BeforeClass
public static void setUp() {
    Type listType = new TypeToken<JsonObject>() {
    }.getType();
    Gson gson = new Gson();
    JsonObject testData = gson.fromJson(new InputStreamReader(JsonIntermediateToParquetGroupConverter.class.getResourceAsStream(RESOURCE_PATH)), listType);
    testCases = testData.getAsJsonObject();
    SourceState source = new SourceState();
    workUnit = new WorkUnitState(source.createWorkUnit(source.createExtract(Extract.TableType.SNAPSHOT_ONLY, "test_namespace", "test_table")));
}

Example 94 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class JdbcPublisher method getStagingTables.

private static Map<String, List<WorkUnitState>> getStagingTables(Collection<? extends WorkUnitState> states, int branches, int i) {
    Map<String, List<WorkUnitState>> stagingTables = Maps.newHashMap();
    for (WorkUnitState workUnitState : states) {
        String stagingTableKey = ForkOperatorUtils.getPropertyNameForBranch(ConfigurationKeys.WRITER_STAGING_TABLE, branches, i);
        String stagingTable = Preconditions.checkNotNull(workUnitState.getProp(stagingTableKey));
        List<WorkUnitState> existing = stagingTables.get(stagingTable);
        if (existing == null) {
            existing = Lists.newArrayList();
            stagingTables.put(stagingTable, existing);
        }
        existing.add(workUnitState);
    }
    return stagingTables;
}

Example 95 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class GobblinMultiTaskAttempt method runWorkUnits.

/**
 * Run a given list of {@link WorkUnit}s of a job.
 *
 * <p>
 *   This method assumes that the given list of {@link WorkUnit}s have already been flattened and
 *   each {@link WorkUnit} contains the task ID in the property {@link ConfigurationKeys#TASK_ID_KEY}.
 * </p>
 *
 * @param countDownLatch a {@link java.util.concurrent.CountDownLatch} waited on for job completion
 * @return a list of {@link Task}s from the {@link WorkUnit}s
 */
private List<Task> runWorkUnits(CountUpAndDownLatch countDownLatch) {
    List<Task> tasks = Lists.newArrayList();
    while (this.workUnits.hasNext()) {
        WorkUnit workUnit = this.workUnits.next();
        String taskId = workUnit.getProp(ConfigurationKeys.TASK_ID_KEY);
        // skip tasks that executed successsfully in a prior attempt
        if (taskSuccessfulInPriorAttempt(taskId)) {
            continue;
        }
        countDownLatch.countUp();
        SubscopedBrokerBuilder<GobblinScopeTypes, ?> taskBrokerBuilder = this.jobBroker.newSubscopedBuilder(new TaskScopeInstance(taskId));
        WorkUnitState workUnitState = new WorkUnitState(workUnit, this.jobState, taskBrokerBuilder);
        workUnitState.setId(taskId);
        workUnitState.setProp(ConfigurationKeys.JOB_ID_KEY, this.jobId);
        workUnitState.setProp(ConfigurationKeys.TASK_ID_KEY, taskId);
        if (this.containerIdOptional.isPresent()) {
            workUnitState.setProp(ConfigurationKeys.TASK_ATTEMPT_ID_KEY, this.containerIdOptional.get());
        }
        // Create a new task from the work unit and submit the task to run
        Task task = createTaskRunnable(workUnitState, countDownLatch);
        this.taskStateTracker.registerNewTask(task);
        task.setTaskFuture(this.taskExecutor.submit(task));
        tasks.add(task);
    }
    new EventSubmitter.Builder(JobMetrics.get(this.jobId).getMetricContext(), "gobblin.runtime").build().submit(JobEvent.TASKS_SUBMITTED, "tasksCount", Long.toString(countDownLatch.getRegisteredParties()));
    return tasks;
}