Examples with ConnectorSession io.prestosql.spi.connector.ConnectorSession used on opensource projects

Example 6 with ConnectorSession

use of io.prestosql.spi.connector.ConnectorSession in project hetu-core by openlookeng.

the class TableCommentSystemTable method cursor.

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession connectorSession, TupleDomain<Integer> constraint) {
    Optional<String> catalogFilter = stringFilter(constraint, 0);
    Optional<String> schemaFilter = stringFilter(constraint, 1);
    Optional<String> tableFilter = stringFilter(constraint, 2);
    Session session = toSession(transactionHandle, connectorSession);
    Builder table = InMemoryRecordSet.builder(COMMENT_TABLE);
    for (String catalog : filter(listCatalogs(session, metadata, accessControl).keySet(), catalogFilter)) {
        QualifiedTablePrefix prefix = tablePrefix(catalog, schemaFilter, tableFilter);
        Set<SchemaTableName> names = ImmutableSet.of();
        try {
            names = listTables(session, metadata, accessControl, prefix);
        } catch (PrestoException e) {
            // listTables throws an exception if cannot connect the database
            LOG.debug(e, "Failed to get tables for catalog: %s", catalog);
        }
        for (SchemaTableName name : names) {
            QualifiedObjectName tableName = new QualifiedObjectName(prefix.getCatalogName(), name.getSchemaName(), name.getTableName());
            Optional<String> comment = Optional.empty();
            try {
                comment = metadata.getTableHandle(session, tableName).map(handle -> metadata.getTableMetadata(session, handle)).map(metadata -> metadata.getMetadata().getComment()).get();
            } catch (PrestoException e) {
                // getTableHandle may throw an exception (e.g. Cassandra connector doesn't allow case insensitive column names)
                LOG.debug(e, "Failed to get metadata for table: %s", name);
            }
            table.addRow(prefix.getCatalogName(), name.getSchemaName(), name.getTableName(), comment.orElse(null));
        }
    }
    return table.build().cursor();
}

Example 7 with ConnectorSession

use of io.prestosql.spi.connector.ConnectorSession in project hetu-core by openlookeng.

the class CatalogJdbcTable method cursor.

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession connectorSession, TupleDomain<Integer> constraint) {
    Session session = toSession(transactionHandle, connectorSession);
    Builder table = InMemoryRecordSet.builder(METADATA);
    for (String name : listCatalogs(session, metadata, accessControl).keySet()) {
        table.addRow(name);
    }
    return table.build().cursor();
}

Example 8 with ConnectorSession

use of io.prestosql.spi.connector.ConnectorSession in project hetu-core by openlookeng.

the class FormatFunction method valueConverter.

private static BiFunction<ConnectorSession, Block, Object> valueConverter(FunctionAndTypeManager functionAndTypeManager, Type type, int position) {
    if (type.equals(UNKNOWN)) {
        return (session, block) -> null;
    }
    if (type.equals(BOOLEAN)) {
        return (session, block) -> type.getBoolean(block, position);
    }
    if (type.equals(TINYINT) || type.equals(SMALLINT) || type.equals(INTEGER) || type.equals(BIGINT)) {
        return (session, block) -> type.getLong(block, position);
    }
    if (type.equals(REAL)) {
        return (session, block) -> intBitsToFloat(toIntExact(type.getLong(block, position)));
    }
    if (type.equals(DOUBLE)) {
        return (session, block) -> type.getDouble(block, position);
    }
    if (type.equals(DATE)) {
        return (session, block) -> LocalDate.ofEpochDay(type.getLong(block, position));
    }
    if (type.equals(TIMESTAMP_WITH_TIME_ZONE)) {
        return (session, block) -> toZonedDateTime(type.getLong(block, position));
    }
    if (type.equals(TIMESTAMP)) {
        return (session, block) -> toLocalDateTime(type.getLong(block, position));
    }
    if (type.equals(TIME)) {
        return (session, block) -> toLocalTime(session, type.getLong(block, position));
    }
    // TODO: support TIME WITH TIME ZONE by making SqlTimeWithTimeZone implement TemporalAccessor
    if (type.equals(JSON)) {
        FunctionHandle functionHandle = functionAndTypeManager.resolveFunction(Optional.empty(), QualifiedObjectName.valueOf(DEFAULT_NAMESPACE, "json_format"), fromTypes(JSON));
        MethodHandle handle = functionAndTypeManager.getBuiltInScalarFunctionImplementation(functionHandle).getMethodHandle();
        return (session, block) -> convertToString(handle, type.getSlice(block, position));
    }
    if (isShortDecimal(type)) {
        int scale = ((DecimalType) type).getScale();
        return (session, block) -> BigDecimal.valueOf(type.getLong(block, position), scale);
    }
    if (isLongDecimal(type)) {
        int scale = ((DecimalType) type).getScale();
        return (session, block) -> new BigDecimal(decodeUnscaledValue(type.getSlice(block, position)), scale);
    }
    if (isVarcharType(type) || isCharType(type)) {
        return (session, block) -> type.getSlice(block, position).toStringUtf8();
    }
    BiFunction<ConnectorSession, Block, Object> function;
    if (type.getJavaType() == long.class) {
        function = (session, block) -> type.getLong(block, position);
    } else if (type.getJavaType() == double.class) {
        function = (session, block) -> type.getDouble(block, position);
    } else if (type.getJavaType() == boolean.class) {
        function = (session, block) -> type.getBoolean(block, position);
    } else if (type.getJavaType() == Slice.class) {
        function = (session, block) -> type.getSlice(block, position);
    } else {
        function = (session, block) -> type.getObject(block, position);
    }
    MethodHandle handle = castToVarchar(functionAndTypeManager, type);
    if ((handle == null) || (handle.type().parameterCount() != 1)) {
        throw new PrestoException(NOT_SUPPORTED, "Type not supported for formatting: " + type.getDisplayName());
    }
    return (session, block) -> convertToString(handle, function.apply(session, block));
}

Example 9 with ConnectorSession

use of io.prestosql.spi.connector.ConnectorSession in project hetu-core by openlookeng.

the class ArrayJoin method specializeArrayJoin.

private static BuiltInScalarFunctionImplementation specializeArrayJoin(Map<String, Type> types, FunctionAndTypeManager functionAndTypeManager, List<Boolean> nullableArguments, MethodHandle methodHandle) {
    Type type = types.get("T");
    List<ArgumentProperty> argumentProperties = nullableArguments.stream().map(nullable -> nullable ? valueTypeArgumentProperty(USE_BOXED_TYPE) : valueTypeArgumentProperty(RETURN_NULL_ON_NULL)).collect(toImmutableList());
    if (type instanceof UnknownType) {
        return new BuiltInScalarFunctionImplementation(false, argumentProperties, methodHandle.bindTo(null), Optional.of(STATE_FACTORY));
    } else {
        try {
            BuiltInScalarFunctionImplementation castFunction = functionAndTypeManager.getBuiltInScalarFunctionImplementation(functionAndTypeManager.lookupCast(CastType.CAST, type.getTypeSignature(), VARCHAR_TYPE_SIGNATURE));
            MethodHandle getter;
            Class<?> elementType = type.getJavaType();
            if (elementType == boolean.class) {
                getter = GET_BOOLEAN;
            } else if (elementType == double.class) {
                getter = GET_DOUBLE;
            } else if (elementType == long.class) {
                getter = GET_LONG;
            } else if (elementType == Slice.class) {
                getter = GET_SLICE;
            } else {
                throw new UnsupportedOperationException("Unsupported type: " + elementType.getName());
            }
            MethodHandle cast = castFunction.getMethodHandle();
            // if the cast doesn't take a ConnectorSession, create an adapter that drops the provided session
            if (cast.type().parameterArray()[0] != ConnectorSession.class) {
                cast = MethodHandles.dropArguments(cast, 0, ConnectorSession.class);
            }
            // Adapt a target cast that takes (ConnectorSession, ?) to one that takes (Block, int, ConnectorSession), which will be invoked by the implementation
            // The first two arguments (Block, int) are filtered through the element type's getXXX method to produce the underlying value that needs to be passed to
            // the cast.
            cast = MethodHandles.permuteArguments(cast, MethodType.methodType(Slice.class, cast.type().parameterArray()[1], cast.type().parameterArray()[0]), 1, 0);
            cast = MethodHandles.dropArguments(cast, 1, int.class);
            cast = MethodHandles.dropArguments(cast, 1, Block.class);
            cast = MethodHandles.foldArguments(cast, getter.bindTo(type));
            MethodHandle target = MethodHandles.insertArguments(methodHandle, 0, cast);
            return new BuiltInScalarFunctionImplementation(false, argumentProperties, target, Optional.of(STATE_FACTORY));
        } catch (PrestoException e) {
            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, String.format("Input type %s not supported", type), e);
        }
    }
}

Example 10 with ConnectorSession

use of io.prestosql.spi.connector.ConnectorSession in project hetu-core by openlookeng.

the class TestHivePageSink method prepareHivePageSink.

// Used to test snapshot. Input pages has 1 row and 1 column. Partition is based on this column.
private HivePageSink prepareHivePageSink() throws IOException {
    // Mock all relevant dependencies
    HiveWriterFactory writerFactory = mock(HiveWriterFactory.class);
    HiveColumnHandle hiveColumnHandle = mock(HiveColumnHandle.class);
    HdfsEnvironment hdfsEnvironment = mock(HdfsEnvironment.class);
    PageIndexerFactory pageIndexerFactory = mock(PageIndexerFactory.class);
    PageIndexer pageIndexer = mock(PageIndexer.class);
    JsonCodec jsonCodec = mock(JsonCodec.class);
    ConnectorSession connectorSession = mock(ConnectorSession.class);
    // Mocked necessary but uninteresting methods
    when(connectorSession.isSnapshotEnabled()).thenReturn(true);
    when(connectorSession.getTaskId()).thenReturn(OptionalInt.of(1));
    when(pageIndexerFactory.createPageIndexer(anyObject())).thenReturn(pageIndexer);
    when(jsonCodec.toJsonBytes(anyObject())).thenReturn(new byte[0]);
    when(writerFactory.isTxnTable()).thenReturn(false);
    HiveWriter hiveWriter = mock(HiveWriter.class);
    when(hiveWriter.getVerificationTask()).thenReturn(Optional.empty());
    when(writerFactory.createWriter(anyObject(), anyObject(), anyObject())).thenReturn(hiveWriter);
    when(writerFactory.createWriterForSnapshotMerge(anyObject(), anyObject(), anyObject())).thenReturn(hiveWriter);
    when(writerFactory.getPartitionName(anyObject(), anyInt())).thenReturn(Optional.empty());
    when(hiveColumnHandle.isPartitionKey()).thenReturn(true);
    // When hdfsEnvironment.doAs() is called, simply invoke the passed in action
    when(hdfsEnvironment.doAs(anyObject(), (GenericExceptionAction) anyObject())).thenAnswer(invocation -> ((GenericExceptionAction) invocation.getArguments()[1]).run());
    doAnswer(invocation -> {
        ((Runnable) invocation.getArguments()[1]).run();
        return null;
    }).when(hdfsEnvironment).doAs(anyObject(), (Runnable) anyObject());
    // The only entry in the page is a integer. We use it to determine partition index.
    // That is, page1 with value 0 is in partition 0; page2 with value 1 is in partition 1.
    // Some functions' return values depend on the number of partitions.
    // Store that as an array entry below, so that other mocked methods can use it.
    int[] maxIndex = new int[1];
    when(pageIndexer.indexPage(anyObject())).thenAnswer(invocation -> {
        maxIndex[0] = (int) ((Page) invocation.getArguments()[0]).getBlock(0).get(0);
        return new int[] { maxIndex[0] };
    });
    when(pageIndexer.getMaxIndex()).thenAnswer(invocation -> maxIndex[0]);
    doAnswer(invocation -> {
        assertEquals(((List) invocation.getArguments()[0]).size(), maxIndex[0] + 1);
        return null;
    }).when(writerFactory).mergeSubFiles(anyObject());
    return new HivePageSink(writerFactory, Collections.singletonList(hiveColumnHandle), Optional.empty(), pageIndexerFactory, mock(TypeManager.class), hdfsEnvironment, 10, mock(ListeningExecutorService.class), jsonCodec, connectorSession, HiveACIDWriteType.INSERT, mock(HiveWritableTableHandle.class));
}