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

Example 26 with DrillbitContext

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

the class TestRepeatedFunction method testRepeated.

@Test
public void testRepeated() throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final PhysicalPlanReader reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    final PhysicalPlan plan = reader.readPhysicalPlan(Files.toString(DrillFileUtils.getResourceAsFile("/physical_repeated_1.json"), Charsets.UTF_8));
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    boolean oneIsOne = false;
    int size = 0;
    final int[] sizes = { 1, 2, 0, 6 };
    while (exec.next()) {
        final IntVector c1 = exec.getValueVectorById(new SchemaPath("cnt", ExpressionPosition.UNKNOWN), IntVector.class);
        final BitVector c2 = exec.getValueVectorById(new SchemaPath("has_min", ExpressionPosition.UNKNOWN), BitVector.class);
        for (int i = 0; i < exec.getRecordCount(); i++) {
            final int curSize = sizes[size % sizes.length];
            assertEquals(curSize, c1.getAccessor().get(i));
            switch(curSize) {
                case 1:
                    assertEquals(oneIsOne, 1 == c2.getAccessor().get(i));
                    oneIsOne = !oneIsOne;
                    break;
                case 2:
                    assertEquals(1, c2.getAccessor().get(i));
                    break;
                case 0:
                    assertEquals(0, c2.getAccessor().get(i));
                    break;
                case 6:
                    assertEquals(1, c2.getAccessor().get(i));
                    break;
            }
            size++;
        }
    }
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 27 with DrillbitContext

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

the class TestTraceMultiRecordBatch method testFilter.

@Test
public void testFilter() throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final PhysicalPlanReader reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    final PhysicalPlan plan = reader.readPhysicalPlan(Files.toString(DrillFileUtils.getResourceAsFile("/trace/multi_record_batch_trace.json"), Charsets.UTF_8));
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    while (exec.next()) {
        for (final ValueVector vv : exec) {
            vv.clear();
        }
    }
    exec.close();
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 28 with DrillbitContext

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

the class TestSimpleUnion method testUnion.

@Test
public void testUnion() throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final PhysicalPlanReader reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    final PhysicalPlan plan = reader.readPhysicalPlan(Files.toString(DrillFileUtils.getResourceAsFile("/union/test1.json"), Charsets.UTF_8));
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    // first batch : 0-row schema-only batch.
    final int[] counts = new int[] { 0, 100, 50 };
    int i = 0;
    while (exec.next()) {
        System.out.println("iteration count:" + exec.getRecordCount());
        assertEquals(counts[i++], exec.getRecordCount());
    }
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 29 with DrillbitContext

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

the class TestDrillSpnegoAuthenticator method setupTest.

@BeforeClass
public static void setupTest() throws Exception {
    spnegoHelper = new KerberosHelper(TestSpnegoAuthentication.class.getSimpleName(), primaryName);
    spnegoHelper.setupKdc(dirTestWatcher.getTmpDir());
    sun.security.krb5.Config.refresh();
    // (2) Reset the default realm.
    final Field defaultRealm = KerberosName.class.getDeclaredField("defaultRealm");
    defaultRealm.setAccessible(true);
    defaultRealm.set(null, KerberosUtil.getDefaultRealm());
    // Create a DrillbitContext with service principal and keytab for DrillSpnegoLoginService
    final DrillConfig newConfig = new DrillConfig(DrillConfig.create().withValue(ExecConstants.HTTP_AUTHENTICATION_MECHANISMS, ConfigValueFactory.fromIterable(Lists.newArrayList("spnego"))).withValue(ExecConstants.HTTP_SPNEGO_PRINCIPAL, ConfigValueFactory.fromAnyRef(spnegoHelper.SERVER_PRINCIPAL)).withValue(ExecConstants.HTTP_SPNEGO_KEYTAB, ConfigValueFactory.fromAnyRef(spnegoHelper.serverKeytab.toString())));
    // Create mock objects for optionManager and AuthConfiguration
    final SystemOptionManager optionManager = Mockito.mock(SystemOptionManager.class);
    Mockito.when(optionManager.getOption(ExecConstants.ADMIN_USERS_VALIDATOR)).thenReturn(ExecConstants.ADMIN_USERS_VALIDATOR.DEFAULT_ADMIN_USERS);
    Mockito.when(optionManager.getOption(ExecConstants.ADMIN_USER_GROUPS_VALIDATOR)).thenReturn(ExecConstants.ADMIN_USER_GROUPS_VALIDATOR.DEFAULT_ADMIN_USER_GROUPS);
    final DrillbitContext drillbitContext = Mockito.mock(DrillbitContext.class);
    Mockito.when(drillbitContext.getConfig()).thenReturn(newConfig);
    Mockito.when(drillbitContext.getOptionManager()).thenReturn(optionManager);
    Authenticator.AuthConfiguration authConfiguration = Mockito.mock(Authenticator.AuthConfiguration.class);
    spnegoAuthenticator = new DrillSpnegoAuthenticator("SPNEGO");
    DrillSpnegoLoginService spnegoLoginService = new DrillSpnegoLoginService(drillbitContext);
    Mockito.when(authConfiguration.getLoginService()).thenReturn(spnegoLoginService);
    Mockito.when(authConfiguration.getIdentityService()).thenReturn(new DefaultIdentityService());
    Mockito.when(authConfiguration.isSessionRenewedOnAuthentication()).thenReturn(true);
    // Set the login service and identity service inside SpnegoAuthenticator
    spnegoAuthenticator.setConfiguration(authConfiguration);
}

Example 30 with DrillbitContext

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

the class TestSpnegoAuthentication method testFORMEnabledWithPlainDisabled.

/**
 * Validate failure in creating FORM security handler when PAM authenticator is absent. PAM authenticator is provided
 * via {@link PlainFactory#getAuthenticator()}
 * @throws Exception
 */
@Test
public void testFORMEnabledWithPlainDisabled() throws Exception {
    try {
        final DrillConfig newConfig = new DrillConfig(DrillConfig.create().withValue(ExecConstants.USER_AUTHENTICATION_ENABLED, ConfigValueFactory.fromAnyRef(true)).withValue(ExecConstants.HTTP_AUTHENTICATION_MECHANISMS, ConfigValueFactory.fromIterable(Lists.newArrayList("form"))).withValue(ExecConstants.HTTP_SPNEGO_PRINCIPAL, ConfigValueFactory.fromAnyRef(spnegoHelper.SERVER_PRINCIPAL)).withValue(ExecConstants.HTTP_SPNEGO_KEYTAB, ConfigValueFactory.fromAnyRef(spnegoHelper.serverKeytab.toString())));
        final ScanResult scanResult = ClassPathScanner.fromPrescan(newConfig);
        final AuthenticatorProviderImpl authenticatorProvider = Mockito.mock(AuthenticatorProviderImpl.class);
        Mockito.when(authenticatorProvider.containsFactory(PlainFactory.SIMPLE_NAME)).thenReturn(false);
        final DrillbitContext context = Mockito.mock(DrillbitContext.class);
        Mockito.when(context.getClasspathScan()).thenReturn(scanResult);
        Mockito.when(context.getConfig()).thenReturn(newConfig);
        Mockito.when(context.getAuthProvider()).thenReturn(authenticatorProvider);
        final DrillHttpSecurityHandlerProvider securityProvider = new DrillHttpSecurityHandlerProvider(newConfig, context);
        fail();
    } catch (Exception ex) {
        assertTrue(ex instanceof DrillbitStartupException);
    }
}