Examples with UserSession org.apache.drill.exec.rpc.user.UserSession used on opensource projects

Example 1 with UserSession

use of org.apache.drill.exec.rpc.user.UserSession in project drill by axbaretto.

the class TestClassTransformation method beforeTestClassTransformation.

@BeforeClass
public static void beforeTestClassTransformation() throws Exception {
    // Tests here require the byte-code merge technique and are meaningless
    // if the plain-old Java technique is selected. Force the plain-Java
    // technique to be off if it happened to be set on in the default
    // configuration.
    System.setProperty(CodeCompiler.PREFER_POJ_CONFIG, "false");
    final UserSession userSession = UserSession.Builder.newBuilder().withOptionManager(getDrillbitContext().getOptionManager()).build();
    sessionOptions = (SessionOptionManager) userSession.getOptions();
}

Example 2 with UserSession

use of org.apache.drill.exec.rpc.user.UserSession in project drill by axbaretto.

the class TestPauseInjection method pauseOnSpecificBit.

@Test
public void pauseOnSpecificBit() {
    final RemoteServiceSet remoteServiceSet = RemoteServiceSet.getLocalServiceSet();
    final ZookeeperHelper zkHelper = new ZookeeperHelper();
    zkHelper.startZookeeper(1);
    try {
        // 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);
            }
        }
    } finally {
        zkHelper.stopZookeeper();
    }
}

Example 3 with UserSession

use of org.apache.drill.exec.rpc.user.UserSession in project drill by apache.

the class TestPauseInjection method timedPauseOnSpecificBit.

@Test
public void timedPauseOnSpecificBit() {
    final RemoteServiceSet remoteServiceSet = RemoteServiceSet.getLocalServiceSet();
    final ZookeeperHelper zkHelper = new ZookeeperHelper();
    zkHelper.startZookeeper(1);
    final long pauseDuration = 2000L;
    final long expectedDuration = pauseDuration;
    try {
        // 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().addTimedPauseOnBit(DummyClass.class, DummyClass.PAUSES, drillbitEndpoint1, 0, pauseDuration).build();
        ControlsInjectionUtil.setControls(session, controls);
        {
            final ExtendedLatch trigger = new ExtendedLatch(1);
            final Pointer<Exception> ex = new Pointer<>();
            final QueryContext queryContext = new QueryContext(session, drillbitContext1, QueryId.getDefaultInstance());
            // 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);
            assertNull("No exception should be thrown.", ex.value);
            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);
            }
        }
    } finally {
        zkHelper.stopZookeeper();
    }
}

Example 4 with UserSession

use of org.apache.drill.exec.rpc.user.UserSession in project drill by apache.

the class TestClassTransformation method beforeTestClassTransformation.

@BeforeClass
public static void beforeTestClassTransformation() {
    // Tests here require the byte-code merge technique and are meaningless
    // if the plain-old Java technique is selected. Force the plain-Java
    // technique to be off if it happened to be set on in the default
    // configuration.
    System.setProperty(CodeCompiler.PREFER_POJ_CONFIG, "false");
    final UserSession userSession = UserSession.Builder.newBuilder().withOptionManager(getDrillbitContext().getOptionManager()).build();
    sessionOptions = userSession.getOptions();
}

Example 5 with UserSession

use of org.apache.drill.exec.rpc.user.UserSession in project drill by axbaretto.

the class DropTableHandler method getPlan.

/**
 * Function resolves the schema and invokes the drop method
 * (while IF EXISTS statement is used function invokes the drop method only if table exists).
 * Raises an exception if the schema is immutable.
 *
 * @param sqlNode - SqlDropTable (SQL parse tree of drop table [if exists] query)
 * @return - Single row indicating drop succeeded or table is not found while IF EXISTS statement is used,
 * raise exception otherwise
 */
@Override
public PhysicalPlan getPlan(SqlNode sqlNode) throws ValidationException, RelConversionException, IOException {
    SqlDropTable dropTableNode = ((SqlDropTable) sqlNode);
    String originalTableName = dropTableNode.getName();
    SchemaPlus defaultSchema = config.getConverter().getDefaultSchema();
    List<String> tableSchema = dropTableNode.getSchema();
    DrillConfig drillConfig = context.getConfig();
    UserSession session = context.getSession();
    AbstractSchema temporarySchema = resolveToTemporarySchema(tableSchema, defaultSchema, drillConfig);
    boolean isTemporaryTable = session.isTemporaryTable(temporarySchema, drillConfig, originalTableName);
    if (isTemporaryTable) {
        session.removeTemporaryTable(temporarySchema, originalTableName, drillConfig);
    } else {
        AbstractSchema drillSchema = SchemaUtilites.resolveToMutableDrillSchema(defaultSchema, tableSchema);
        Table tableToDrop = SqlHandlerUtil.getTableFromSchema(drillSchema, originalTableName);
        if (tableToDrop == null || tableToDrop.getJdbcTableType() != Schema.TableType.TABLE) {
            if (dropTableNode.checkTableExistence()) {
                return DirectPlan.createDirectPlan(context, false, String.format("Table [%s] not found", originalTableName));
            } else {
                throw UserException.validationError().message("Table [%s] not found", originalTableName).build(logger);
            }
        }
        SqlHandlerUtil.dropTableFromSchema(drillSchema, originalTableName);
    }
    String message = String.format("%s [%s] dropped", isTemporaryTable ? "Temporary table" : "Table", originalTableName);
    logger.info(message);
    return DirectPlan.createDirectPlan(context, true, message);
}