org.apache.arrow.adbc.driver.jdbc.InfoMetadataBuilder Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of adbc-driver-jdbc Show documentation
An ADBC driver wrapping the JDBC API.
The newest version!
/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.arrow.adbc.driver.jdbc;

import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.apache.arrow.adbc.core.AdbcDriver;
import org.apache.arrow.adbc.core.AdbcInfoCode;
import org.apache.arrow.adbc.core.StandardSchemas;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.UInt4Vector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.complex.DenseUnionVector;
import org.checkerframework.checker.nullness.qual.Nullable;

/** Helper class to track state needed to build up the info structure. */
final class InfoMetadataBuilder implements AutoCloseable {
  private static final byte STRING_VALUE_TYPE_ID = (byte) 0;
  private static final byte BIGINT_VALUE_TYPE_ID = (byte) 2;
  private static final Map SUPPORTED_CODES = new HashMap<>();
  private final BufferAllocator allocator;
  private final Collection requestedCodes;
  private final DatabaseMetaData dbmd;
  private VectorSchemaRoot root;

  final UInt4Vector infoCodes;
  final DenseUnionVector infoValues;
  final VarCharVector stringValues;
  final BigIntVector bigIntValues;

  @FunctionalInterface
  interface AddInfo {
    void accept(InfoMetadataBuilder builder, int rowIndex) throws SQLException;
  }

  static {
    SUPPORTED_CODES.put(
        AdbcInfoCode.VENDOR_NAME.getValue(),
        (b, idx) -> {
          b.setStringValue(idx, b.dbmd.getDatabaseProductName());
        });
    SUPPORTED_CODES.put(
        AdbcInfoCode.VENDOR_VERSION.getValue(),
        (b, idx) -> {
          b.setStringValue(idx, b.dbmd.getDatabaseProductVersion());
        });
    SUPPORTED_CODES.put(
        AdbcInfoCode.DRIVER_NAME.getValue(),
        (b, idx) -> {
          final String driverName = "ADBC JDBC Driver (" + b.dbmd.getDriverName() + ")";
          b.setStringValue(idx, driverName);
        });
    SUPPORTED_CODES.put(
        AdbcInfoCode.DRIVER_VERSION.getValue(),
        (b, idx) -> {
          final String driverVersion = b.dbmd.getDriverVersion() + " (ADBC Driver Version 0.0.1)";
          b.setStringValue(idx, driverVersion);
        });
    SUPPORTED_CODES.put(
        AdbcInfoCode.DRIVER_ADBC_VERSION.getValue(),
        (b, idx) -> {
          b.setBigIntValue(idx, AdbcDriver.ADBC_VERSION_1_1_0);
        });
  }

  InfoMetadataBuilder(BufferAllocator allocator, Connection connection, int @Nullable [] infoCodes)
      throws SQLException {
    this.allocator = allocator;
    if (infoCodes == null) {
      this.requestedCodes = new ArrayList<>(SUPPORTED_CODES.keySet());
    } else {
      this.requestedCodes = IntStream.of(infoCodes).boxed().collect(Collectors.toList());
    }
    this.root = VectorSchemaRoot.create(StandardSchemas.GET_INFO_SCHEMA, allocator);
    this.dbmd = connection.getMetaData();
    this.infoCodes = (UInt4Vector) root.getVector(0);
    this.infoValues = (DenseUnionVector) root.getVector(1);
    this.stringValues = this.infoValues.getVarCharVector(STRING_VALUE_TYPE_ID);
    this.bigIntValues = this.infoValues.getBigIntVector(BIGINT_VALUE_TYPE_ID);
  }

  void setBigIntValue(int index, long value) {
    infoValues.setValueCount(index + 1);
    infoValues.setTypeId(index, BIGINT_VALUE_TYPE_ID);
    bigIntValues.setSafe(index, value);
    infoValues
        .getOffsetBuffer()
        .setInt((long) index * DenseUnionVector.OFFSET_WIDTH, bigIntValues.getValueCount());
    bigIntValues.setValueCount(bigIntValues.getValueCount() + 1);
  }

  void setStringValue(int index, final String value) {
    infoValues.setValueCount(index + 1);
    infoValues.setTypeId(index, STRING_VALUE_TYPE_ID);
    stringValues.setSafe(index, value.getBytes(StandardCharsets.UTF_8));
    infoValues
        .getOffsetBuffer()
        .setInt((long) index * DenseUnionVector.OFFSET_WIDTH, stringValues.getLastSet());
  }

  VectorSchemaRoot build() throws SQLException {
    int rowIndex = 0;
    for (final Integer code : requestedCodes) {
      final AddInfo metadata = SUPPORTED_CODES.get(code);
      if (metadata == null) {
        continue;
      }
      infoCodes.setSafe(rowIndex, code);
      metadata.accept(this, rowIndex++);
    }
    root.setRowCount(rowIndex);
    VectorSchemaRoot result = root;
    try {
      root = VectorSchemaRoot.create(StandardSchemas.GET_INFO_SCHEMA, allocator);
    } catch (RuntimeException e) {
      result.close();
      throw e;
    }
    return result;
  }

  @Override
  public void close() throws Exception {
    AutoCloseables.close(root);
  }
}