From bd9a646c37270799766096be55112d1dfd600b28 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Diego=20Fern=C3=A1ndez=20Giraldo?= Date: Sat, 4 Nov 2023 11:01:41 -0600 Subject: [PATCH] GH-33475: [Java] Add parameter binding for Prepared Statements in JDBC driver (#38404) This PR is a combination of #33961 and #14627. The goal is to support parametrized queries through the Arrow Flight SQL JDBC driver. An Arrow Flight SQL server returns a Schema for the `PreparedStatement` parameters. The driver then converts the `Field` list associated with the Schema into a list of `AvaticaParameter`. When the user sets values for the parameters, Avatica generates a list of `TypedValue`, which we then bind to each parameter vector. This conversion between Arrow and Avatica is handled by implementations of a `AvaticaParameterConverter` interface for each Arrow type. This interface which provides 2 methods: - createParameter: Create an `AvaticaParameter` from the given Arrow `Field`. - bindParameter: Cast the given `TypedValue` and bind it to the `FieldVector` at the specified index. This PR purposely leaves out a few features: - We currently naively cast the `TypedValue` values assuming users set the type correctly. If this cast fails, we raise an exception letting the user know that the cast is not supported. This could be improved in subsequent PRs to do smarter conversions from other types. - We currently don't provide conversions for complex types such as List, Map, Struct, Union, Interval, and Duration. The stubs are there so they can be implemented as needed. - Tests for specific types have not been implemented. I'm not very familiar with a lot of these JDBC types so it's hard to implement rigorous tets. * Closes: #33475 * Closes: #35536 Authored-by: Diego Fernandez Signed-off-by: David Li --- .../driver/jdbc/ArrowFlightJdbcFactory.java | 9 - .../ArrowFlightJdbcFlightStreamResultSet.java | 8 +- ...owFlightJdbcVectorSchemaRootResultSet.java | 19 +- .../driver/jdbc/ArrowFlightMetaImpl.java | 90 ++++--- .../jdbc/ArrowFlightPreparedStatement.java | 33 --- .../client/ArrowFlightSqlClientHandler.java | 20 ++ .../converter/AvaticaParameterConverter.java | 42 ++++ .../impl/BaseAvaticaParameterConverter.java | 42 ++++ .../impl/BinaryAvaticaParameterConverter.java | 49 ++++ .../impl/BoolAvaticaParameterConverter.java | 49 ++++ .../impl/DateAvaticaParameterConverter.java | 53 ++++ .../DecimalAvaticaParameterConverter.java | 51 ++++ .../DurationAvaticaParameterConverter.java | 43 ++++ ...edSizeBinaryAvaticaParameterConverter.java | 49 ++++ ...ixedSizeListAvaticaParameterConverter.java | 43 ++++ ...loatingPointAvaticaParameterConverter.java | 53 ++++ .../impl/IntAvaticaParameterConverter.java | 79 ++++++ .../IntervalAvaticaParameterConverter.java | 49 ++++ .../LargeBinaryAvaticaParameterConverter.java | 49 ++++ .../LargeListAvaticaParameterConverter.java | 43 ++++ .../LargeUtf8AvaticaParameterConverter.java | 50 ++++ .../impl/ListAvaticaParameterConverter.java | 43 ++++ .../impl/MapAvaticaParameterConverter.java | 43 ++++ .../impl/NullAvaticaParameterConverter.java | 49 ++++ .../impl/StructAvaticaParameterConverter.java | 43 ++++ .../impl/TimeAvaticaParameterConverter.java | 61 +++++ .../TimestampAvaticaParameterConverter.java | 78 ++++++ .../impl/UnionAvaticaParameterConverter.java | 43 ++++ .../impl/Utf8AvaticaParameterConverter.java | 50 ++++ .../jdbc/utils/AvaticaParameterBinder.java | 233 ++++++++++++++++++ .../arrow/driver/jdbc/utils/ConvertUtils.java | 154 +++++++++++- .../ArrowFlightPreparedStatementTest.java | 80 ++++++ .../jdbc/utils/MockFlightSqlProducer.java | 81 +++++- .../arrow/flight/sql/FlightSqlClient.java | 42 ++-- 34 files changed, 1809 insertions(+), 114 deletions(-) create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/AvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BaseAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BoolAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DateAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DecimalAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DurationAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FixedSizeBinaryAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FixedSizeListAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FloatingPointAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/IntAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/IntervalAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeBinaryAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeListAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeUtf8AvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/ListAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/MapAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/NullAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/StructAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimeAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimestampAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UnionAvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8AvaticaParameterConverter.java create mode 100644 java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java index 216e4cd002bc3..16bdede02d039 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java @@ -17,8 +17,6 @@ package org.apache.arrow.driver.jdbc; -import static org.apache.arrow.driver.jdbc.utils.ConvertUtils.convertArrowFieldsToColumnMetaDataList; - import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.Properties; @@ -26,7 +24,6 @@ import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; import org.apache.arrow.memory.RootAllocator; -import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaFactory; import org.apache.calcite.avatica.AvaticaResultSetMetaData; @@ -89,12 +86,6 @@ public ArrowFlightPreparedStatement newPreparedStatement( ArrowFlightSqlClientHandler.PreparedStatement preparedStatement = flightConnection.getMeta().getPreparedStatement(statementHandle); - if (preparedStatement == null) { - preparedStatement = flightConnection.getClientHandler().prepare(signature.sql); - } - final Schema resultSetSchema = preparedStatement.getDataSetSchema(); - signature.columns.addAll(convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields())); - return ArrowFlightPreparedStatement.newPreparedStatement( flightConnection, preparedStatement, statementHandle, signature, resultType, resultSetConcurrency, resultSetHoldability); diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java index 2e42cf0166b06..e23267ebe9ebf 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java @@ -93,7 +93,7 @@ static ArrowFlightJdbcFlightStreamResultSet fromFlightInfo( final TimeZone timeZone = TimeZone.getDefault(); final QueryState state = new QueryState(); - final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null); + final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null); final AvaticaResultSetMetaData resultSetMetaData = new AvaticaResultSetMetaData(null, null, signature); @@ -154,11 +154,7 @@ private void populateDataForCurrentFlightStream() throws SQLException { currentVectorSchemaRoot = originalRoot; } - if (schema != null) { - populateData(currentVectorSchemaRoot, schema); - } else { - populateData(currentVectorSchemaRoot); - } + populateData(currentVectorSchemaRoot, schema); } @Override diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java index 20a2af6a84aa4..626ae95bc5bbe 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java @@ -17,20 +17,18 @@ package org.apache.arrow.driver.jdbc; -import static java.util.Objects.isNull; - import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.HashSet; import java.util.List; +import java.util.Objects; import java.util.Set; import java.util.TimeZone; import org.apache.arrow.driver.jdbc.utils.ConvertUtils; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaResultSet; import org.apache.calcite.avatica.AvaticaResultSetMetaData; @@ -74,7 +72,7 @@ public static ArrowFlightJdbcVectorSchemaRootResultSet fromVectorSchemaRoot( final TimeZone timeZone = TimeZone.getDefault(); final QueryState state = new QueryState(); - final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null); + final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null); final AvaticaResultSetMetaData resultSetMetaData = new AvaticaResultSetMetaData(null, null, signature); @@ -93,17 +91,12 @@ protected AvaticaResultSet execute() throws SQLException { } void populateData(final VectorSchemaRoot vectorSchemaRoot) { - final List fields = vectorSchemaRoot.getSchema().getFields(); - final List columns = ConvertUtils.convertArrowFieldsToColumnMetaDataList(fields); - signature.columns.clear(); - signature.columns.addAll(columns); - - this.vectorSchemaRoot = vectorSchemaRoot; - execute2(new ArrowFlightJdbcCursor(vectorSchemaRoot), this.signature.columns); + populateData(vectorSchemaRoot, null); } void populateData(final VectorSchemaRoot vectorSchemaRoot, final Schema schema) { - final List columns = ConvertUtils.convertArrowFieldsToColumnMetaDataList(schema.getFields()); + Schema currentSchema = schema == null ? vectorSchemaRoot.getSchema() : schema; + final List columns = ConvertUtils.convertArrowFieldsToColumnMetaDataList(currentSchema.getFields()); signature.columns.clear(); signature.columns.addAll(columns); @@ -137,7 +130,7 @@ public void close() { } catch (final Exception e) { exceptions.add(e); } - if (!isNull(statement)) { + if (!Objects.isNull(statement)) { try { super.close(); } catch (final Exception e) { diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java index f825e7d13cef5..382750914992f 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java @@ -17,8 +17,6 @@ package org.apache.arrow.driver.jdbc; -import static java.lang.String.format; - import java.sql.Connection; import java.sql.SQLException; import java.sql.SQLTimeoutException; @@ -29,7 +27,10 @@ import java.util.concurrent.ConcurrentHashMap; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement; +import org.apache.arrow.driver.jdbc.utils.AvaticaParameterBinder; +import org.apache.arrow.driver.jdbc.utils.ConvertUtils; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.ColumnMetaData; @@ -54,12 +55,20 @@ public ArrowFlightMetaImpl(final AvaticaConnection connection) { setDefaultConnectionProperties(); } - static Signature newSignature(final String sql) { + /** + * Construct a signature. + */ + static Signature newSignature(final String sql, Schema resultSetSchema, Schema parameterSchema) { + List columnMetaData = resultSetSchema == null ? + new ArrayList<>() : ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields()); + List parameters = parameterSchema == null ? + new ArrayList<>() : ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields()); + return new Signature( - new ArrayList(), + columnMetaData, sql, - Collections.emptyList(), - Collections.emptyMap(), + parameters, + Collections.emptyMap(), null, // unnecessary, as SQL requests use ArrowFlightJdbcCursor StatementType.SELECT ); @@ -84,23 +93,28 @@ public void commit(final ConnectionHandle connectionHandle) { public ExecuteResult execute(final StatementHandle statementHandle, final List typedValues, final long maxRowCount) { Preconditions.checkArgument(connection.id.equals(statementHandle.connectionId), - "Connection IDs are not consistent"); + "Connection IDs are not consistent"); + PreparedStatement preparedStatement = getPreparedStatement(statementHandle); + + if (preparedStatement == null) { + throw new IllegalStateException("Prepared statement not found: " + statementHandle); + } + + + new AvaticaParameterBinder(preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator()) + .bind(typedValues); + if (statementHandle.signature == null) { // Update query - final StatementHandleKey key = new StatementHandleKey(statementHandle); - PreparedStatement preparedStatement = statementHandlePreparedStatementMap.get(key); - if (preparedStatement == null) { - throw new IllegalStateException("Prepared statement not found: " + statementHandle); - } long updatedCount = preparedStatement.executeUpdate(); return new ExecuteResult(Collections.singletonList(MetaResultSet.count(statementHandle.connectionId, - statementHandle.id, updatedCount))); + statementHandle.id, updatedCount))); } else { // TODO Why is maxRowCount ignored? return new ExecuteResult( - Collections.singletonList(MetaResultSet.create( - statementHandle.connectionId, statementHandle.id, - true, statementHandle.signature, null))); + Collections.singletonList(MetaResultSet.create( + statementHandle.connectionId, statementHandle.id, + true, statementHandle.signature, null))); } } @@ -114,7 +128,23 @@ public ExecuteResult execute(final StatementHandle statementHandle, public ExecuteBatchResult executeBatch(final StatementHandle statementHandle, final List> parameterValuesList) throws IllegalStateException { - throw new IllegalStateException("executeBatch not implemented."); + Preconditions.checkArgument(connection.id.equals(statementHandle.connectionId), + "Connection IDs are not consistent"); + PreparedStatement preparedStatement = getPreparedStatement(statementHandle); + + if (preparedStatement == null) { + throw new IllegalStateException("Prepared statement not found: " + statementHandle); + } + + final AvaticaParameterBinder binder = new AvaticaParameterBinder(preparedStatement, + ((ArrowFlightConnection) connection).getBufferAllocator()); + for (int i = 0; i < parameterValuesList.size(); i++) { + binder.bind(parameterValuesList.get(i), i); + } + + // Update query + long[] updatedCounts = {preparedStatement.executeUpdate()}; + return new ExecuteBatchResult(updatedCounts); } @Override @@ -126,18 +156,24 @@ public Frame fetch(final StatementHandle statementHandle, final long offset, * the results. */ throw AvaticaConnection.HELPER.wrap( - format("%s does not use frames.", this), + String.format("%s does not use frames.", this), AvaticaConnection.HELPER.unsupported()); } + private PreparedStatement prepareForHandle(final String query, StatementHandle handle) { + final PreparedStatement preparedStatement = + ((ArrowFlightConnection) connection).getClientHandler().prepare(query); + handle.signature = newSignature(query, preparedStatement.getDataSetSchema(), + preparedStatement.getParameterSchema()); + statementHandlePreparedStatementMap.put(new StatementHandleKey(handle), preparedStatement); + return preparedStatement; + } + @Override public StatementHandle prepare(final ConnectionHandle connectionHandle, final String query, final long maxRowCount) { final StatementHandle handle = super.createStatement(connectionHandle); - handle.signature = newSignature(query); - final PreparedStatement preparedStatement = - ((ArrowFlightConnection) connection).getClientHandler().prepare(query); - statementHandlePreparedStatementMap.put(new StatementHandleKey(handle), preparedStatement); + prepareForHandle(query, handle); return handle; } @@ -157,20 +193,18 @@ public ExecuteResult prepareAndExecute(final StatementHandle handle, final PrepareCallback callback) throws NoSuchStatementException { try { - final PreparedStatement preparedStatement = - ((ArrowFlightConnection) connection).getClientHandler().prepare(query); + PreparedStatement preparedStatement = prepareForHandle(query, handle); final StatementType statementType = preparedStatement.getType(); - statementHandlePreparedStatementMap.put(new StatementHandleKey(handle), preparedStatement); - final Signature signature = newSignature(query); + final long updateCount = statementType.equals(StatementType.UPDATE) ? preparedStatement.executeUpdate() : -1; synchronized (callback.getMonitor()) { callback.clear(); - callback.assign(signature, null, updateCount); + callback.assign(handle.signature, null, updateCount); } callback.execute(); final MetaResultSet metaResultSet = MetaResultSet.create(handle.connectionId, handle.id, - false, signature, null); + false, handle.signature, null); return new ExecuteResult(Collections.singletonList(metaResultSet)); } catch (SQLTimeoutException e) { // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and Runtime Exceptions. diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index 8784e39840b6a..7203f02daa9a1 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -17,15 +17,12 @@ package org.apache.arrow.driver.jdbc; -import java.sql.Connection; import java.sql.PreparedStatement; import java.sql.SQLException; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; -import org.apache.arrow.driver.jdbc.utils.ConvertUtils; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.util.Preconditions; -import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaPreparedStatement; import org.apache.calcite.avatica.Meta.Signature; import org.apache.calcite.avatica.Meta.StatementHandle; @@ -50,36 +47,6 @@ private ArrowFlightPreparedStatement(final ArrowFlightConnection connection, this.preparedStatement = Preconditions.checkNotNull(preparedStatement); } - /** - * Creates a new {@link ArrowFlightPreparedStatement} from the provided information. - * - * @param connection the {@link Connection} to use. - * @param statementHandle the {@link StatementHandle} to use. - * @param signature the {@link Signature} to use. - * @param resultSetType the ResultSet type. - * @param resultSetConcurrency the ResultSet concurrency. - * @param resultSetHoldability the ResultSet holdability. - * @return a new {@link PreparedStatement}. - * @throws SQLException on error. - */ - static ArrowFlightPreparedStatement createNewPreparedStatement( - final ArrowFlightConnection connection, - final StatementHandle statementHandle, - final Signature signature, - final int resultSetType, - final int resultSetConcurrency, - final int resultSetHoldability) throws SQLException { - - final ArrowFlightSqlClientHandler.PreparedStatement prepare = connection.getClientHandler().prepare(signature.sql); - final Schema resultSetSchema = prepare.getDataSetSchema(); - - signature.columns.addAll(ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields())); - - return new ArrowFlightPreparedStatement( - connection, prepare, statementHandle, - signature, resultSetType, resultSetConcurrency, resultSetHoldability); - } - static ArrowFlightPreparedStatement newPreparedStatement(final ArrowFlightConnection connection, final ArrowFlightSqlClientHandler.PreparedStatement preparedStmt, final StatementHandle statementHandle, diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java index 38e5a9bb362d2..66372092b8e99 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java @@ -49,6 +49,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.Meta.StatementType; import org.slf4j.Logger; @@ -206,6 +207,15 @@ public interface PreparedStatement extends AutoCloseable { */ Schema getDataSetSchema(); + /** + * Gets the {@link Schema} of the parameters for this {@link PreparedStatement}. + * + * @return {@link Schema}. + */ + Schema getParameterSchema(); + + void setParameters(VectorSchemaRoot parameters); + @Override void close(); } @@ -241,6 +251,16 @@ public Schema getDataSetSchema() { return preparedStatement.getResultSetSchema(); } + @Override + public Schema getParameterSchema() { + return preparedStatement.getParameterSchema(); + } + + @Override + public void setParameters(VectorSchemaRoot parameters) { + preparedStatement.setParameters(parameters); + } + @Override public void close() { try { diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/AvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/AvaticaParameterConverter.java new file mode 100644 index 0000000000000..c01e688f37396 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/AvaticaParameterConverter.java @@ -0,0 +1,42 @@ +/* + * 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.driver.jdbc.converter; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * Interface for a class in charge of converting between AvaticaParameters and TypedValues and + * Arrow. + */ +public interface AvaticaParameterConverter { + + /** + * Bind a TypedValue to a FieldVector at the given index. + * + * @param vector FieldVector that the parameter should be bound to. + * @param typedValue TypedValue to bind as a parameter. + * @param index Vector index that the TypedValue should be bound to. + * @return Whether the value was set successfully. + */ + boolean bindParameter(FieldVector vector, TypedValue typedValue, int index); + + AvaticaParameter createParameter(Field field); +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BaseAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BaseAvaticaParameterConverter.java new file mode 100644 index 0000000000000..f5cf8358b7a14 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BaseAvaticaParameterConverter.java @@ -0,0 +1,42 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.driver.jdbc.converter.AvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.utils.SqlTypes; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.SqlType; + +/** + * Base AvaticaParameterConverter with a generic createParameter method that can be used by most + * Arrow types. + */ +abstract class BaseAvaticaParameterConverter implements AvaticaParameterConverter { + protected AvaticaParameter createParameter(Field field, boolean signed) { + final String name = field.getName(); + final ArrowType arrowType = field.getType(); + final String typeName = arrowType.toString(); + final int precision = 0; // Would have to know about the actual number + final int scale = 0; // According to https://www.postgresql.org/docs/current/datatype-numeric.html + final int jdbcType = SqlTypes.getSqlTypeIdFromArrowType(arrowType); + final String className = SqlType.valueOf(jdbcType).clazz.getCanonicalName(); + return new AvaticaParameter(signed, precision, scale, jdbcType, typeName, className, name); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryAvaticaParameterConverter.java new file mode 100644 index 0000000000000..d244848955e52 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryAvaticaParameterConverter.java @@ -0,0 +1,49 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Binary Arrow types. + */ +public class BinaryAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public BinaryAvaticaParameterConverter(ArrowType.Binary type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + byte[] value = (byte[]) typedValue.toJdbc(null); + if (vector instanceof VarBinaryVector) { + ((VarBinaryVector) vector).setSafe(index, value); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BoolAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BoolAvaticaParameterConverter.java new file mode 100644 index 0000000000000..6725154d03c25 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BoolAvaticaParameterConverter.java @@ -0,0 +1,49 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Bool Arrow types. + */ +public class BoolAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public BoolAvaticaParameterConverter(ArrowType.Bool type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + boolean value = (boolean) typedValue.toLocal(); + if (vector instanceof BitVector) { + ((BitVector) vector).setSafe(index, value ? 1 : 0); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DateAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DateAvaticaParameterConverter.java new file mode 100644 index 0000000000000..0da1dabe43721 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DateAvaticaParameterConverter.java @@ -0,0 +1,53 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.DateDayVector; +import org.apache.arrow.vector.DateMilliVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Date Arrow types. + */ +public class DateAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public DateAvaticaParameterConverter(ArrowType.Date type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + int value = (int) typedValue.toLocal(); + if (vector instanceof DateMilliVector) { + ((DateMilliVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof DateDayVector) { + ((DateDayVector) vector).setSafe(index, value); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DecimalAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DecimalAvaticaParameterConverter.java new file mode 100644 index 0000000000000..fad43e2e06a76 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DecimalAvaticaParameterConverter.java @@ -0,0 +1,51 @@ +/* + * 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.driver.jdbc.converter.impl; + +import java.math.BigDecimal; + +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Decimal Arrow types. + */ +public class DecimalAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public DecimalAvaticaParameterConverter(ArrowType.Decimal type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + BigDecimal value = (BigDecimal) typedValue.toLocal(); + if (vector instanceof DecimalVector) { + ((DecimalVector) vector).setSafe(index, value); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, true); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DurationAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DurationAvaticaParameterConverter.java new file mode 100644 index 0000000000000..89f2fc1d5c12f --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/DurationAvaticaParameterConverter.java @@ -0,0 +1,43 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Duration Arrow types. + */ +public class DurationAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public DurationAvaticaParameterConverter(ArrowType.Duration type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FixedSizeBinaryAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FixedSizeBinaryAvaticaParameterConverter.java new file mode 100644 index 0000000000000..a90434f695ac3 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FixedSizeBinaryAvaticaParameterConverter.java @@ -0,0 +1,49 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.FixedSizeBinaryVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for FixedSizeBinary Arrow types. + */ +public class FixedSizeBinaryAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public FixedSizeBinaryAvaticaParameterConverter(ArrowType.FixedSizeBinary type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + byte[] value = (byte[]) typedValue.toJdbc(null); + if (vector instanceof FixedSizeBinaryVector) { + ((FixedSizeBinaryVector) vector).setSafe(index, value); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FixedSizeListAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FixedSizeListAvaticaParameterConverter.java new file mode 100644 index 0000000000000..60231a2460286 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FixedSizeListAvaticaParameterConverter.java @@ -0,0 +1,43 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for FixedSizeList Arrow types. + */ +public class FixedSizeListAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public FixedSizeListAvaticaParameterConverter(ArrowType.FixedSizeList type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FloatingPointAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FloatingPointAvaticaParameterConverter.java new file mode 100644 index 0000000000000..9f305a2b6f20d --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/FloatingPointAvaticaParameterConverter.java @@ -0,0 +1,53 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for FloatingPoint Arrow types. + */ +public class FloatingPointAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public FloatingPointAvaticaParameterConverter(ArrowType.FloatingPoint type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + Number value = (Number) typedValue.value; + if (vector instanceof Float4Vector) { + ((Float4Vector) vector).setSafe(index, value.floatValue()); + return true; + } else if (vector instanceof Float8Vector) { + ((Float8Vector) vector).setSafe(index, value.doubleValue()); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, true); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/IntAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/IntAvaticaParameterConverter.java new file mode 100644 index 0000000000000..6684e8d32c9a9 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/IntAvaticaParameterConverter.java @@ -0,0 +1,79 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.TinyIntVector; +import org.apache.arrow.vector.UInt1Vector; +import org.apache.arrow.vector.UInt2Vector; +import org.apache.arrow.vector.UInt4Vector; +import org.apache.arrow.vector.UInt8Vector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Int Arrow types. + */ +public class IntAvaticaParameterConverter extends BaseAvaticaParameterConverter { + private final ArrowType.Int type; + + public IntAvaticaParameterConverter(ArrowType.Int type) { + this.type = type; + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + Number value = (Number) typedValue.value; + if (vector instanceof TinyIntVector) { + ((TinyIntVector) vector).setSafe(index, value.intValue()); + return true; + } else if (vector instanceof SmallIntVector) { + ((SmallIntVector) vector).setSafe(index, value.intValue()); + return true; + } else if (vector instanceof IntVector) { + ((IntVector) vector).setSafe(index, value.intValue()); + return true; + } else if (vector instanceof BigIntVector) { + ((BigIntVector) vector).setSafe(index, value.longValue()); + return true; + } else if (vector instanceof UInt1Vector) { + ((UInt1Vector) vector).setSafe(index, value.intValue()); + return true; + } else if (vector instanceof UInt2Vector) { + ((UInt2Vector) vector).setSafe(index, value.intValue()); + return true; + } else if (vector instanceof UInt4Vector) { + ((UInt4Vector) vector).setSafe(index, value.intValue()); + return true; + } else if (vector instanceof UInt8Vector) { + ((UInt8Vector) vector).setSafe(index, value.longValue()); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, type.getIsSigned()); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/IntervalAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/IntervalAvaticaParameterConverter.java new file mode 100644 index 0000000000000..724275d51091e --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/IntervalAvaticaParameterConverter.java @@ -0,0 +1,49 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Interval Arrow types. + */ +public class IntervalAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public IntervalAvaticaParameterConverter(ArrowType.Interval type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + // Object value = typedValue.toLocal(); + // if (vector instanceof IntervalDayVector) { + // ((IntervalDayVector) vector).setSafe(index, () value); + // } else if (vector instanceof IntervalYearVector) { + // ((IntervalYearVector) vector).setSafe(index, () value); + // } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeBinaryAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeBinaryAvaticaParameterConverter.java new file mode 100644 index 0000000000000..133ec2072d583 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeBinaryAvaticaParameterConverter.java @@ -0,0 +1,49 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.LargeVarBinaryVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for LargeBinary Arrow types. + */ +public class LargeBinaryAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public LargeBinaryAvaticaParameterConverter(ArrowType.LargeBinary type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + byte[] value = (byte[]) typedValue.toJdbc(null); + if (vector instanceof LargeVarBinaryVector) { + ((LargeVarBinaryVector) vector).setSafe(index, value); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeListAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeListAvaticaParameterConverter.java new file mode 100644 index 0000000000000..6ef6920474860 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeListAvaticaParameterConverter.java @@ -0,0 +1,43 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for LargeList Arrow types. + */ +public class LargeListAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public LargeListAvaticaParameterConverter(ArrowType.LargeList type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeUtf8AvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeUtf8AvaticaParameterConverter.java new file mode 100644 index 0000000000000..d412ab007ac67 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/LargeUtf8AvaticaParameterConverter.java @@ -0,0 +1,50 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.LargeVarCharVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.util.Text; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for LargeUtf8 Arrow types. + */ +public class LargeUtf8AvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public LargeUtf8AvaticaParameterConverter(ArrowType.LargeUtf8 type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + String value = (String) typedValue.toLocal(); + if (vector instanceof LargeVarCharVector) { + ((LargeVarCharVector) vector).setSafe(index, new Text(value)); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/ListAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/ListAvaticaParameterConverter.java new file mode 100644 index 0000000000000..aec59cb4d428e --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/ListAvaticaParameterConverter.java @@ -0,0 +1,43 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for List Arrow types. + */ +public class ListAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public ListAvaticaParameterConverter(ArrowType.List type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/MapAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/MapAvaticaParameterConverter.java new file mode 100644 index 0000000000000..feac3794d222b --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/MapAvaticaParameterConverter.java @@ -0,0 +1,43 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Map Arrow types. + */ +public class MapAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public MapAvaticaParameterConverter(ArrowType.Map type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/NullAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/NullAvaticaParameterConverter.java new file mode 100644 index 0000000000000..e2c184fb11a09 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/NullAvaticaParameterConverter.java @@ -0,0 +1,49 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.NullVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Null Arrow types. + */ +public class NullAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public NullAvaticaParameterConverter(ArrowType.Null type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + Object value = typedValue.toLocal(); + if (vector instanceof NullVector) { + if (value != null) { throw new RuntimeException("Can't set non-null value on NullVector"); } + vector.setNull(index); + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/StructAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/StructAvaticaParameterConverter.java new file mode 100644 index 0000000000000..5dfe923cb516e --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/StructAvaticaParameterConverter.java @@ -0,0 +1,43 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Struct Arrow types. + */ +public class StructAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public StructAvaticaParameterConverter(ArrowType.Struct type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimeAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimeAvaticaParameterConverter.java new file mode 100644 index 0000000000000..c6b79537fd435 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimeAvaticaParameterConverter.java @@ -0,0 +1,61 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.TimeMicroVector; +import org.apache.arrow.vector.TimeMilliVector; +import org.apache.arrow.vector.TimeNanoVector; +import org.apache.arrow.vector.TimeSecVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Time Arrow types. + */ +public class TimeAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public TimeAvaticaParameterConverter(ArrowType.Time type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + int value = (int) typedValue.toLocal(); + if (vector instanceof TimeMicroVector) { + ((TimeMicroVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeMilliVector) { + ((TimeMilliVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeNanoVector) { + ((TimeNanoVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeSecVector) { + ((TimeSecVector) vector).setSafe(index, value); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimestampAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimestampAvaticaParameterConverter.java new file mode 100644 index 0000000000000..3c1940b75cfa7 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimestampAvaticaParameterConverter.java @@ -0,0 +1,78 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.TimeStampMicroTZVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TimeStampMilliTZVector; +import org.apache.arrow.vector.TimeStampMilliVector; +import org.apache.arrow.vector.TimeStampNanoTZVector; +import org.apache.arrow.vector.TimeStampNanoVector; +import org.apache.arrow.vector.TimeStampSecTZVector; +import org.apache.arrow.vector.TimeStampSecVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Timestamp Arrow types. + */ +public class TimestampAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public TimestampAvaticaParameterConverter(ArrowType.Timestamp type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + // FIXME: how should we handle TZ? Do we need to convert the value to the TZ on the vector? + long value = (long) typedValue.toLocal(); + if (vector instanceof TimeStampSecVector) { + ((TimeStampSecVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeStampMicroVector) { + ((TimeStampMicroVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeStampMilliVector) { + ((TimeStampMilliVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeStampNanoVector) { + ((TimeStampNanoVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeStampSecTZVector) { + ((TimeStampSecTZVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeStampMicroTZVector) { + ((TimeStampMicroTZVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeStampMilliTZVector) { + ((TimeStampMilliTZVector) vector).setSafe(index, value); + return true; + } else if (vector instanceof TimeStampNanoTZVector) { + ((TimeStampNanoTZVector) vector).setSafe(index, value); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UnionAvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UnionAvaticaParameterConverter.java new file mode 100644 index 0000000000000..6b171e685579a --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UnionAvaticaParameterConverter.java @@ -0,0 +1,43 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Union Arrow types. + */ +public class UnionAvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public UnionAvaticaParameterConverter(ArrowType.Union type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8AvaticaParameterConverter.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8AvaticaParameterConverter.java new file mode 100644 index 0000000000000..9223e5361d2d5 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8AvaticaParameterConverter.java @@ -0,0 +1,50 @@ +/* + * 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.driver.jdbc.converter.impl; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.util.Text; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * AvaticaParameterConverter for Utf8 Arrow types. + */ +public class Utf8AvaticaParameterConverter extends BaseAvaticaParameterConverter { + + public Utf8AvaticaParameterConverter(ArrowType.Utf8 type) { + } + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + String value = (String) typedValue.toLocal(); + if (vector instanceof VarCharVector) { + ((VarCharVector) vector).setSafe(index, new Text(value)); + return true; + } + return false; + } + + @Override + public AvaticaParameter createParameter(Field field) { + return createParameter(field, false); + } +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java new file mode 100644 index 0000000000000..c5be4697db7c5 --- /dev/null +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java @@ -0,0 +1,233 @@ +/* + * 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.driver.jdbc.utils; + +import java.util.List; + +import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement; +import org.apache.arrow.driver.jdbc.converter.impl.BinaryAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.BoolAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.DateAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.DecimalAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.DurationAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.FixedSizeBinaryAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.FixedSizeListAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.FloatingPointAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.IntAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.IntervalAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.LargeBinaryAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.LargeListAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.LargeUtf8AvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.ListAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.MapAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.NullAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.StructAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.TimeAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.TimestampAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.calcite.avatica.remote.TypedValue; + +/** + * Convert Avatica PreparedStatement parameters from a list of TypedValue to Arrow and bind them to the + * VectorSchemaRoot representing the PreparedStatement parameters. + *

+ * NOTE: Make sure to close the parameters VectorSchemaRoot once we're done with them. + */ +public class AvaticaParameterBinder { + private final PreparedStatement preparedStatement; + private final VectorSchemaRoot parameters; + + public AvaticaParameterBinder(PreparedStatement preparedStatement, BufferAllocator bufferAllocator) { + this.parameters = VectorSchemaRoot.create(preparedStatement.getParameterSchema(), bufferAllocator); + this.preparedStatement = preparedStatement; + } + + /** + * Bind the given Avatica values to the prepared statement. + * @param typedValues The parameter values. + */ + public void bind(List typedValues) { + bind(typedValues, 0); + } + + /** + * Bind the given Avatica values to the prepared statement at the given index. + * @param typedValues The parameter values. + * @param index index for parameter. + */ + public void bind(List typedValues, int index) { + if (preparedStatement.getParameterSchema().getFields().size() != typedValues.size()) { + throw new IllegalStateException( + String.format("Prepared statement has %s parameters, but only received %s", + preparedStatement.getParameterSchema().getFields().size(), + typedValues.size())); + } + + for (int i = 0; i < typedValues.size(); i++) { + bind(parameters.getVector(i), typedValues.get(i), index); + } + + if (!typedValues.isEmpty()) { + parameters.setRowCount(index + 1); + preparedStatement.setParameters(parameters); + } + } + + /** + * Bind a TypedValue to the given index on the FieldVctor. + * + * @param vector FieldVector to bind to. + * @param typedValue TypedValue to bind to the vector. + * @param index Vector index to bind the value at. + */ + private void bind(FieldVector vector, TypedValue typedValue, int index) { + try { + if (!vector.getField().getType().accept(new BinderVisitor(vector, typedValue, index))) { + throw new RuntimeException( + String.format("Binding to vector type %s is not yet supported", vector.getClass())); + } + } catch (ClassCastException e) { + throw new RuntimeException( + String.format("Binding value of type %s is not yet supported for expected Arrow type %s", + typedValue.type, vector.getField().getType())); + } + } + + private static class BinderVisitor implements ArrowType.ArrowTypeVisitor { + private final FieldVector vector; + private final TypedValue typedValue; + private final int index; + + private BinderVisitor(FieldVector vector, TypedValue value, int index) { + this.vector = vector; + this.typedValue = value; + this.index = index; + } + + @Override + public Boolean visit(ArrowType.Null type) { + return new NullAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Struct type) { + return new StructAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.List type) { + return new ListAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.LargeList type) { + return new LargeListAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.FixedSizeList type) { + return new FixedSizeListAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Union type) { + return new UnionAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Map type) { + return new MapAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Int type) { + return new IntAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.FloatingPoint type) { + return new FloatingPointAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Utf8 type) { + return new Utf8AvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.LargeUtf8 type) { + return new LargeUtf8AvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Binary type) { + return new BinaryAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.LargeBinary type) { + return new LargeBinaryAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.FixedSizeBinary type) { + return new FixedSizeBinaryAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Bool type) { + return new BoolAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Decimal type) { + return new DecimalAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Date type) { + return new DateAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Time type) { + return new TimeAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Timestamp type) { + return new TimestampAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Interval type) { + return new IntervalAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + + @Override + public Boolean visit(ArrowType.Duration type) { + return new DurationAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); + } + } + +} diff --git a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java index 324f991ef09e9..b21b03340e9f9 100644 --- a/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java +++ b/java/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java @@ -22,15 +22,37 @@ import java.util.stream.Collectors; import java.util.stream.Stream; +import org.apache.arrow.driver.jdbc.converter.impl.BinaryAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.BoolAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.DateAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.DecimalAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.DurationAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.FixedSizeBinaryAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.FixedSizeListAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.FloatingPointAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.IntAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.IntervalAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.LargeBinaryAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.LargeListAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.LargeUtf8AvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.ListAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.MapAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.NullAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.StructAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.TimeAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.TimestampAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter; import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.ColumnMetaData; import org.apache.calcite.avatica.proto.Common; import org.apache.calcite.avatica.proto.Common.ColumnMetaData.Builder; /** - * Convert Fields To Column MetaData List functions. + * Convert objects between Arrow and Avatica. */ public final class ConvertUtils { @@ -113,4 +135,134 @@ public static void setOnColumnMetaDataBuilder(final Builder builder, builder.setSearchable(searchable); } } + + /** + * Convert Fields To Avatica Parameters. + * + * @param fields list of {@link Field}. + * @return list of {@link AvaticaParameter}. + */ + public static List convertArrowFieldsToAvaticaParameters(final List fields) { + return fields.stream() + .map(field -> field.getType().accept(new ConverterVisitor(field))) + .collect(Collectors.toList()); + } + + private static class ConverterVisitor implements ArrowType.ArrowTypeVisitor { + private final Field field; + + private ConverterVisitor(Field field) { + this.field = field; + } + + @Override + public AvaticaParameter visit(ArrowType.Null type) { + return new NullAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Struct type) { + return new StructAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.List type) { + return new ListAvaticaParameterConverter(type).createParameter(field); + + } + + @Override + public AvaticaParameter visit(ArrowType.LargeList type) { + return new LargeListAvaticaParameterConverter(type).createParameter(field); + + } + + @Override + public AvaticaParameter visit(ArrowType.FixedSizeList type) { + return new FixedSizeListAvaticaParameterConverter(type).createParameter(field); + + } + + @Override + public AvaticaParameter visit(ArrowType.Union type) { + return new UnionAvaticaParameterConverter(type).createParameter(field); + + } + + @Override + public AvaticaParameter visit(ArrowType.Map type) { + return new MapAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Int type) { + return new IntAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.FloatingPoint type) { + return new FloatingPointAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Utf8 type) { + return new Utf8AvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.LargeUtf8 type) { + return new LargeUtf8AvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Binary type) { + return new BinaryAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.LargeBinary type) { + return new LargeBinaryAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.FixedSizeBinary type) { + return new FixedSizeBinaryAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Bool type) { + return new BoolAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Decimal type) { + return new DecimalAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Date type) { + return new DateAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Time type) { + return new TimeAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Timestamp type) { + return new TimestampAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Interval type) { + return new IntervalAvaticaParameterConverter(type).createParameter(field); + } + + @Override + public AvaticaParameter visit(ArrowType.Duration type) { + return new DurationAvaticaParameterConverter(type).createParameter(field); + } + } + } diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java index df2577e955881..b19f049544ada 100644 --- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java @@ -20,14 +20,26 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.junit.jupiter.api.Assertions.assertEquals; +import java.nio.charset.StandardCharsets; import java.sql.Connection; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; +import java.util.Arrays; +import java.util.Collections; import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.flight.sql.FlightSqlUtils; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.Types; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.Text; import org.junit.AfterClass; import org.junit.Before; import org.junit.BeforeClass; @@ -73,6 +85,39 @@ public void testSimpleQueryNoParameterBinding() throws SQLException { } } + @Test + public void testQueryWithParameterBinding() throws SQLException { + final String query = "Fake query with parameters"; + final Schema schema = new Schema(Collections.singletonList(Field.nullable("", Types.MinorType.INT.getType()))); + PRODUCER.addSelectQuery(query, schema, + Collections.singletonList(listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(schema, + allocator)) { + ((IntVector) root.getVector(0)).setSafe(0, 10); + root.setRowCount(1); + listener.start(root); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + })); + + PRODUCER.addExpectedParameters(query, + new Schema(Collections.singletonList(Field.nullable("", ArrowType.Utf8.INSTANCE))), + Collections.singletonList(Collections.singletonList(new Text("foo".getBytes(StandardCharsets.UTF_8))))); + try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) { + preparedStatement.setString(1, "foo"); + try (final ResultSet resultSet = preparedStatement.executeQuery()) { + resultSet.next(); + assert true; + } + } + } + + @Test @Ignore("https://github.com/apache/arrow/issues/34741: flaky test") public void testPreparedStatementExecutionOnce() throws SQLException { @@ -107,4 +152,39 @@ public void testUpdateQuery() throws SQLException { assertEquals(42, updated); } } + + @Test + public void testUpdateQueryWithParameters() throws SQLException { + String query = "Fake update with parameters"; + PRODUCER.addUpdateQuery(query, /*updatedRows*/42); + PRODUCER.addExpectedParameters(query, + new Schema(Collections.singletonList(Field.nullable("", ArrowType.Utf8.INSTANCE))), + Collections.singletonList(Collections.singletonList(new Text("foo".getBytes(StandardCharsets.UTF_8))))); + try (final PreparedStatement stmt = connection.prepareStatement(query)) { + // TODO: make sure this is validated on the server too + stmt.setString(1, "foo"); + int updated = stmt.executeUpdate(); + assertEquals(42, updated); + } + } + + @Test + public void testUpdateQueryWithBatchedParameters() throws SQLException { + String query = "Fake update with batched parameters"; + PRODUCER.addUpdateQuery(query, /*updatedRows*/42); + PRODUCER.addExpectedParameters(query, + new Schema(Collections.singletonList(Field.nullable("", ArrowType.Utf8.INSTANCE))), + Arrays.asList( + Collections.singletonList(new Text("foo".getBytes(StandardCharsets.UTF_8))), + Collections.singletonList(new Text("bar".getBytes(StandardCharsets.UTF_8))))); + try (final PreparedStatement stmt = connection.prepareStatement(query)) { + // TODO: make sure this is validated on the server too + stmt.setString(1, "foo"); + stmt.addBatch(); + stmt.setString(1, "bar"); + stmt.addBatch(); + int[] updated = stmt.executeBatch(); + assertEquals(42, updated[0]); + } + } } diff --git a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java index 0299eeb46d93b..75a7396931c8e 100644 --- a/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java +++ b/java/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java @@ -34,6 +34,7 @@ import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.Objects; import java.util.UUID; import java.util.function.BiConsumer; import java.util.function.Consumer; @@ -75,6 +76,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Schema; @@ -97,7 +99,9 @@ public final class MockFlightSqlProducer implements FlightSqlProducer { private final Map>> updateResultProviders = new HashMap<>(); - private SqlInfoBuilder sqlInfoBuilder = new SqlInfoBuilder(); + private final SqlInfoBuilder sqlInfoBuilder = new SqlInfoBuilder(); + private final Map parameterSchemas = new HashMap<>(); + private final Map>> expectedParameterValues = new HashMap<>(); private final Map actionTypeCounter = new HashMap<>(); @@ -192,6 +196,12 @@ void addUpdateQuery(final String sqlCommand, format("Attempted to overwrite pre-existing query: <%s>.", sqlCommand)); } + /** Registers parameters expected to be provided with a prepared statement. */ + public void addExpectedParameters(String query, Schema parameterSchema, List> expectedValues) { + parameterSchemas.put(query, parameterSchema); + expectedParameterValues.put(query, expectedValues); + } + @Override public void createPreparedStatement(final ActionCreatePreparedStatementRequest request, final CallContext callContext, @@ -223,6 +233,13 @@ public void createPreparedStatement(final ActionCreatePreparedStatementRequest r return; } + final Schema parameterSchema = parameterSchemas.get(query); + if (parameterSchema != null) { + final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + MessageSerializer.serialize(new WriteChannel(Channels.newChannel(outputStream)), parameterSchema); + resultBuilder.setParameterSchema(ByteString.copyFrom(outputStream.toByteArray())); + } + listener.onNext(new Result(pack(resultBuilder.build()).toByteArray())); } catch (final Throwable t) { listener.onError(t); @@ -330,6 +347,51 @@ public Runnable acceptPutStatement(final CommandStatementUpdate commandStatement }; } + private boolean validateParameters(String query, + FlightStream flightStream, + StreamListener streamListener) { + final List> expectedValues = expectedParameterValues.get(query); + if (expectedValues != null) { + int index = 0; + while (flightStream.next()) { + final VectorSchemaRoot root = flightStream.getRoot(); + for (int i = 0; i < root.getRowCount(); i++) { + if (index >= expectedValues.size()) { + streamListener.onError(CallStatus.INVALID_ARGUMENT + .withDescription("More parameter rows provided than expected") + .toRuntimeException()); + return true; + } + List expectedRow = expectedValues.get(index++); + if (root.getFieldVectors().size() != expectedRow.size()) { + streamListener.onError(CallStatus.INVALID_ARGUMENT + .withDescription("Parameter count mismatch") + .toRuntimeException()); + return true; + } + + for (int paramIndex = 0; paramIndex < expectedRow.size(); paramIndex++) { + Object expected = expectedRow.get(paramIndex); + Object actual = root.getVector(paramIndex).getObject(i); + if (!Objects.equals(expected, actual)) { + streamListener.onError(CallStatus.INVALID_ARGUMENT + .withDescription(String.format("Parameter mismatch. Expected: %s Actual: %s", expected, actual)) + .toRuntimeException()); + return true; + } + } + } + } + if (index < expectedValues.size()) { + streamListener.onError(CallStatus.INVALID_ARGUMENT + .withDescription("Fewer parameter rows provided than expected") + .toRuntimeException()); + return true; + } + } + return false; + } + @Override public Runnable acceptPutPreparedStatementUpdate( final CommandPreparedStatementUpdate commandPreparedStatementUpdate, @@ -339,6 +401,11 @@ public Runnable acceptPutPreparedStatementUpdate( final String query = Preconditions.checkNotNull( preparedStatements.get(handle), format("No query registered under handle: <%s>.", handle)); + + if (validateParameters(query, flightStream, streamListener)) { + return () -> { }; + } + return acceptPutStatement( CommandStatementUpdate.newBuilder().setQuery(query).build(), callContext, flightStream, streamListener); @@ -349,8 +416,16 @@ public Runnable acceptPutPreparedStatementQuery( final CommandPreparedStatementQuery commandPreparedStatementQuery, final CallContext callContext, final FlightStream flightStream, final StreamListener streamListener) { - // TODO Implement this method. - throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + final ByteString handle = commandPreparedStatementQuery.getPreparedStatementHandle(); + final String query = Preconditions.checkNotNull( + preparedStatements.get(handle), + format("No query registered under handle: <%s>.", handle)); + + if (validateParameters(query, flightStream, streamListener)) { + return () -> { }; + } + + return streamListener::onCompleted; } @Override diff --git a/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java b/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java index e72354513013e..93d933f00f38f 100644 --- a/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java +++ b/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java @@ -951,12 +951,11 @@ public static class PreparedStatement implements AutoCloseable { * {@code PreparedStatement} setters. */ public void setParameters(final VectorSchemaRoot parameterBindingRoot) { - if (this.parameterBindingRoot != null) { - if (this.parameterBindingRoot.equals(parameterBindingRoot)) { - return; - } - this.parameterBindingRoot.close(); + if (parameterBindingRoot == this.parameterBindingRoot) { + // Nothing to do if we're attempting to set the same parameters again. + return; } + clearParameters(); this.parameterBindingRoot = parameterBindingRoot; } @@ -1038,19 +1037,25 @@ public FlightInfo execute(final CallOption... options) { .toByteArray()); if (parameterBindingRoot != null && parameterBindingRoot.getRowCount() > 0) { - final SyncPutListener putListener = new SyncPutListener(); - - FlightClient.ClientStreamListener listener = - client.startPut(descriptor, parameterBindingRoot, putListener, options); - - listener.putNext(); - listener.completed(); - listener.getResult(); + putParameters(descriptor, options); } return client.getInfo(descriptor, options); } + private SyncPutListener putParameters(FlightDescriptor descriptor, CallOption... options) { + final SyncPutListener putListener = new SyncPutListener(); + + FlightClient.ClientStreamListener listener = + client.startPut(descriptor, parameterBindingRoot, putListener, options); + + listener.putNext(); + listener.completed(); + listener.getResult(); + + return putListener; + } + /** * Checks whether this client is open. * @@ -1074,11 +1079,8 @@ public long executeUpdate(final CallOption... options) { .build()) .toByteArray()); setParameters(parameterBindingRoot == null ? VectorSchemaRoot.of() : parameterBindingRoot); - final SyncPutListener putListener = new SyncPutListener(); - final FlightClient.ClientStreamListener listener = - client.startPut(descriptor, parameterBindingRoot, putListener, options); - listener.putNext(); - listener.completed(); + SyncPutListener putListener = putParameters(descriptor, options); + try { final PutResult read = putListener.read(); try (final ArrowBuf metadata = read.getApplicationMetadata()) { @@ -1112,9 +1114,7 @@ public void close(final CallOption... options) { final Iterator closePreparedStatementResults = client.doAction(action, options); closePreparedStatementResults.forEachRemaining(result -> { }); - if (parameterBindingRoot != null) { - parameterBindingRoot.close(); - } + clearParameters(); } @Override