Examples with RecordWriter org.apache.hive.hcatalog.streaming.RecordWriter used on opensource projects

Example 1 with RecordWriter

use of org.apache.hive.hcatalog.streaming.RecordWriter in project nifi by apache.

the class HiveWriterTest method setup.

@Before
public void setup() throws Exception {
    hiveEndPoint = mock(HiveEndPoint.class);
    txnsPerBatch = 100;
    autoCreatePartitions = true;
    callTimeout = 0;
    executorService = mock(ExecutorService.class);
    streamingConnection = mock(StreamingConnection.class);
    transactionBatch = mock(TransactionBatch.class);
    userGroupInformation = mock(UserGroupInformation.class);
    hiveConf = mock(HiveConf.class);
    recordWriter = mock(RecordWriter.class);
    recordWriterCallable = mock(Callable.class);
    when(recordWriterCallable.call()).thenReturn(recordWriter);
    when(hiveEndPoint.newConnection(autoCreatePartitions, hiveConf, userGroupInformation)).thenReturn(streamingConnection);
    when(streamingConnection.fetchTransactionBatch(txnsPerBatch, recordWriter)).thenReturn(transactionBatch);
    when(executorService.submit(isA(Callable.class))).thenAnswer(invocation -> {
        Future future = mock(Future.class);
        Answer<Object> answer = i -> ((Callable) invocation.getArguments()[0]).call();
        when(future.get()).thenAnswer(answer);
        when(future.get(anyLong(), any(TimeUnit.class))).thenAnswer(answer);
        return future;
    });
    when(userGroupInformation.doAs(isA(PrivilegedExceptionAction.class))).thenAnswer(invocation -> {
        try {
            try {
                return ((PrivilegedExceptionAction) invocation.getArguments()[0]).run();
            } catch (UncheckedExecutionException e) {
                // Creation of strict json writer will fail due to external deps, this gives us chance to catch it
                for (StackTraceElement stackTraceElement : e.getStackTrace()) {
                    if (stackTraceElement.toString().startsWith("org.apache.hive.hcatalog.streaming.StrictJsonWriter.<init>(")) {
                        return recordWriterCallable.call();
                    }
                }
                throw e;
            }
        } catch (IOException | Error | RuntimeException | InterruptedException e) {
            throw e;
        } catch (Throwable e) {
            throw new UndeclaredThrowableException(e);
        }
    });
    initWriter();
}