Examples with DrillbitContext org.apache.drill.exec.server.DrillbitContext used on opensource projects

Example 11 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class TestWriteToDisk method test.

@Test
@SuppressWarnings("static-method")
public void test() throws Exception {
    final List<ValueVector> vectorList = Lists.newArrayList();
    final DrillConfig config = DrillConfig.create();
    try (final RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
        final Drillbit bit = new Drillbit(config, serviceSet)) {
        bit.run();
        final DrillbitContext context = bit.getContext();
        final MaterializedField intField = MaterializedField.create("int", Types.required(TypeProtos.MinorType.INT));
        final MaterializedField binField = MaterializedField.create("binary", Types.required(TypeProtos.MinorType.VARBINARY));
        try (final IntVector intVector = (IntVector) TypeHelper.getNewVector(intField, context.getAllocator());
            final VarBinaryVector binVector = (VarBinaryVector) TypeHelper.getNewVector(binField, context.getAllocator())) {
            AllocationHelper.allocate(intVector, 4, 4);
            AllocationHelper.allocate(binVector, 4, 5);
            vectorList.add(intVector);
            vectorList.add(binVector);
            intVector.getMutator().setSafe(0, 0);
            binVector.getMutator().setSafe(0, "ZERO".getBytes());
            intVector.getMutator().setSafe(1, 1);
            binVector.getMutator().setSafe(1, "ONE".getBytes());
            intVector.getMutator().setSafe(2, 2);
            binVector.getMutator().setSafe(2, "TWO".getBytes());
            intVector.getMutator().setSafe(3, 3);
            binVector.getMutator().setSafe(3, "THREE".getBytes());
            intVector.getMutator().setValueCount(4);
            binVector.getMutator().setValueCount(4);
            VectorContainer container = new VectorContainer();
            container.addCollection(vectorList);
            container.setRecordCount(4);
            WritableBatch batch = WritableBatch.getBatchNoHVWrap(container.getRecordCount(), container, false);
            VectorAccessibleSerializable wrap = new VectorAccessibleSerializable(batch, context.getAllocator());
            Configuration conf = new Configuration();
            conf.set(FileSystem.FS_DEFAULT_NAME_KEY, FileSystem.DEFAULT_FS);
            final VectorAccessibleSerializable newWrap = new VectorAccessibleSerializable(context.getAllocator());
            try (final FileSystem fs = FileSystem.get(conf)) {
                final File tempDir = Files.createTempDir();
                tempDir.deleteOnExit();
                final Path path = new Path(tempDir.getAbsolutePath(), "drillSerializable");
                try (final FSDataOutputStream out = fs.create(path)) {
                    wrap.writeToStream(out);
                    out.close();
                }
                try (final FSDataInputStream in = fs.open(path)) {
                    newWrap.readFromStream(in);
                }
            }
            final VectorAccessible newContainer = newWrap.get();
            for (VectorWrapper<?> w : newContainer) {
                try (ValueVector vv = w.getValueVector()) {
                    int values = vv.getAccessor().getValueCount();
                    for (int i = 0; i < values; i++) {
                        final Object o = vv.getAccessor().getObject(i);
                        if (o instanceof byte[]) {
                            System.out.println(new String((byte[]) o));
                        } else {
                            System.out.println(o);
                        }
                    }
                }
            }
        }
    }
}

Example 12 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class TestCustomTunnel method ensureRoundTripJackson.

@Test
public void ensureRoundTripJackson() throws Exception {
    final DrillbitContext context = getDrillbitContext();
    final MesgA mesgA = new MesgA();
    mesgA.fieldA = "123";
    mesgA.fieldB = "okra";
    final TestCustomMessageHandlerJackson handler = new TestCustomMessageHandlerJackson(mesgA);
    context.getController().registerCustomHandler(1003, handler, new ControlTunnel.JacksonSerDe<MesgA>(MesgA.class), new ControlTunnel.JacksonSerDe<MesgB>(MesgB.class));
    final ControlTunnel loopbackTunnel = context.getController().getTunnel(context.getEndpoint());
    final CustomTunnel<MesgA, MesgB> tunnel = loopbackTunnel.getCustomTunnel(1003, new ControlTunnel.JacksonSerDe<MesgA>(MesgA.class), new ControlTunnel.JacksonSerDe<MesgB>(MesgB.class));
    CustomFuture<MesgB> future = tunnel.send(mesgA);
    assertEquals(expectedB, future.get());
}

Example 13 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class TestCustomTunnel method ensureRoundTrip.

@Test
public void ensureRoundTrip() throws Exception {
    final DrillbitContext context = getDrillbitContext();
    final TestCustomMessageHandler handler = new TestCustomMessageHandler(context.getEndpoint(), false);
    context.getController().registerCustomHandler(1001, handler, DrillbitEndpoint.PARSER);
    final ControlTunnel loopbackTunnel = context.getController().getTunnel(context.getEndpoint());
    final CustomTunnel<DrillbitEndpoint, QueryId> tunnel = loopbackTunnel.getCustomTunnel(1001, DrillbitEndpoint.class, QueryId.PARSER);
    CustomFuture<QueryId> future = tunnel.send(context.getEndpoint());
    assertEquals(expectedId, future.get());
}

Example 14 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class TestExceptionInjection method injectionOnSpecificBit.

@SuppressWarnings("static-method")
@Test
public void injectionOnSpecificBit() {
    final RemoteServiceSet remoteServiceSet = RemoteServiceSet.getLocalServiceSet();
    final ZookeeperHelper zkHelper = new ZookeeperHelper();
    zkHelper.startZookeeper(1);
    // Creating two drillbits
    final Drillbit drillbit1, drillbit2;
    final DrillConfig drillConfig = zkHelper.getConfig();
    try {
        drillbit1 = Drillbit.start(drillConfig, remoteServiceSet);
        drillbit2 = Drillbit.start(drillConfig, remoteServiceSet);
    } catch (DrillbitStartupException e) {
        throw new RuntimeException("Failed to start drillbits.", e);
    }
    final DrillbitContext drillbitContext1 = drillbit1.getContext();
    final DrillbitContext drillbitContext2 = drillbit2.getContext();
    final UserSession session = UserSession.Builder.newBuilder().withCredentials(UserBitShared.UserCredentials.newBuilder().setUserName("foo").build()).withUserProperties(UserProperties.getDefaultInstance()).withOptionManager(drillbitContext1.getOptionManager()).build();
    final String passthroughDesc = "<<injected from descPassthrough>>";
    final int nSkip = 7;
    final int nFire = 3;
    final Class<? extends Throwable> exceptionClass = RuntimeException.class;
    // only drillbit1's (address, port)
    final String controls = Controls.newBuilder().addExceptionOnBit(DummyClass.class, passthroughDesc, exceptionClass, drillbitContext1.getEndpoint(), nSkip, nFire).build();
    ControlsInjectionUtil.setControls(session, controls);
    {
        final QueryContext queryContext1 = new QueryContext(session, drillbitContext1, QueryId.getDefaultInstance());
        final DummyClass class1 = new DummyClass(queryContext1);
        // these shouldn't throw
        for (int i = 0; i < nSkip; ++i) {
            class1.descPassthroughMethod(passthroughDesc);
        }
        // these should throw
        for (int i = 0; i < nFire; ++i) {
            assertPassthroughThrows(class1, exceptionClass.getName(), passthroughDesc);
        }
        // this shouldn't throw
        class1.descPassthroughMethod(passthroughDesc);
        try {
            queryContext1.close();
        } catch (Exception e) {
            fail();
        }
    }
    {
        final QueryContext queryContext2 = new QueryContext(session, drillbitContext2, QueryId.getDefaultInstance());
        final DummyClass class2 = new DummyClass(queryContext2);
        // these shouldn't throw
        for (int i = 0; i < nSkip; ++i) {
            class2.descPassthroughMethod(passthroughDesc);
        }
        // these shouldn't throw
        for (int i = 0; i < nFire; ++i) {
            class2.descPassthroughMethod(passthroughDesc);
        }
        // this shouldn't throw
        class2.descPassthroughMethod(passthroughDesc);
        try {
            queryContext2.close();
        } catch (Exception e) {
            fail();
        }
    }
}

Example 15 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class TestPauseInjection method pauseOnSpecificBit.

@Test
public void pauseOnSpecificBit() {
    final RemoteServiceSet remoteServiceSet = RemoteServiceSet.getLocalServiceSet();
    final ZookeeperHelper zkHelper = new ZookeeperHelper();
    zkHelper.startZookeeper(1);
    // Creating two drillbits
    final Drillbit drillbit1, drillbit2;
    final DrillConfig drillConfig = zkHelper.getConfig();
    try {
        drillbit1 = Drillbit.start(drillConfig, remoteServiceSet);
        drillbit2 = Drillbit.start(drillConfig, remoteServiceSet);
    } catch (final DrillbitStartupException e) {
        throw new RuntimeException("Failed to start two drillbits.", e);
    }
    final DrillbitContext drillbitContext1 = drillbit1.getContext();
    final DrillbitContext drillbitContext2 = drillbit2.getContext();
    final UserSession session = UserSession.Builder.newBuilder().withCredentials(UserCredentials.newBuilder().setUserName("foo").build()).withUserProperties(UserProperties.getDefaultInstance()).withOptionManager(drillbitContext1.getOptionManager()).build();
    final DrillbitEndpoint drillbitEndpoint1 = drillbitContext1.getEndpoint();
    final String controls = Controls.newBuilder().addPauseOnBit(DummyClass.class, DummyClass.PAUSES, drillbitEndpoint1).build();
    ControlsInjectionUtil.setControls(session, controls);
    {
        final long expectedDuration = 1000L;
        final ExtendedLatch trigger = new ExtendedLatch(1);
        final Pointer<Exception> ex = new Pointer<>();
        final QueryContext queryContext = new QueryContext(session, drillbitContext1, QueryId.getDefaultInstance());
        (new ResumingThread(queryContext, trigger, ex, expectedDuration)).start();
        // test that the pause happens
        final DummyClass dummyClass = new DummyClass(queryContext, trigger);
        final long actualDuration = dummyClass.pauses();
        assertTrue(String.format("Test should stop for at least %d milliseconds.", expectedDuration), expectedDuration <= actualDuration);
        assertTrue("No exception should be thrown.", ex.value == null);
        try {
            queryContext.close();
        } catch (final Exception e) {
            fail("Failed to close query context: " + e);
        }
    }
    {
        final ExtendedLatch trigger = new ExtendedLatch(1);
        final QueryContext queryContext = new QueryContext(session, drillbitContext2, QueryId.getDefaultInstance());
        // if the resume did not happen, the test would hang
        final DummyClass dummyClass = new DummyClass(queryContext, trigger);
        dummyClass.pauses();
        try {
            queryContext.close();
        } catch (final Exception e) {
            fail("Failed to close query context: " + e);
        }
    }
}