Skip to content

Commit 73c08e3

Browse files
xborderxborder
committed
Keep statement orchestration internal
Dispatch Avatica operations through package-private hooks and guard the public statement API surface.
1 parent fee4892 commit 73c08e3

7 files changed

Lines changed: 128 additions & 104 deletions

File tree

‎flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet
7070
throws SQLException {
7171
super(statement, state, signature, resultSetMetaData, timeZone, firstFrame);
7272
this.connection = (ArrowFlightConnection) statement.connection;
73-
this.flightInfo = ((ArrowFlightMetaStatement) statement).executeFlightInfoQuery();
73+
this.flightInfo = ((ArrowFlightInfoStatement) statement).executeFlightInfoQuery();
7474
}
7575

7676
/** Private constructor for fromFlightInfo. */

‎flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java‎

Lines changed: 29 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,10 @@ public ArrowFlightMetaImpl(final AvaticaConnection connection) {
4747

4848
@Override
4949
public void closeStatement(final StatementHandle statementHandle) {
50-
getMetaStatement(statementHandle).closeStatement();
50+
final AvaticaStatement statement = getStatement(statementHandle);
51+
if (statement instanceof ArrowFlightPreparedStatement) {
52+
((ArrowFlightPreparedStatement) statement).closePreparedResources();
53+
}
5154
}
5255

5356
@Override
@@ -60,7 +63,8 @@ public ExecuteResult execute(
6063
final StatementHandle statementHandle,
6164
final List<TypedValue> typedValues,
6265
final long maxRowCount) {
63-
return getMetaStatement(statementHandle).execute(statementHandle, typedValues, maxRowCount);
66+
return getPreparedStatement(statementHandle)
67+
.executeWithTypedValues(statementHandle, typedValues, maxRowCount);
6468
}
6569

6670
@Override
@@ -75,7 +79,8 @@ public ExecuteResult execute(
7579
public ExecuteBatchResult executeBatch(
7680
final StatementHandle statementHandle, final List<List<TypedValue>> parameterValuesList)
7781
throws IllegalStateException {
78-
return getMetaStatement(statementHandle).executeBatch(statementHandle, parameterValuesList);
82+
return getPreparedStatement(statementHandle)
83+
.executeBatchWithTypedValues(statementHandle, parameterValuesList);
7984
}
8085

8186
@Override
@@ -125,8 +130,13 @@ public ExecuteResult prepareAndExecute(
125130
final PrepareCallback callback)
126131
throws NoSuchStatementException {
127132
try {
128-
return getMetaStatement(handle)
129-
.prepareAndExecute(query, maxRowCount, maxRowsInFirstFrame, callback);
133+
final AvaticaStatement statement = getStatement(handle);
134+
if (statement instanceof ArrowFlightStatement) {
135+
return ((ArrowFlightStatement) statement)
136+
.prepareAndExecuteInternal(query, maxRowCount, maxRowsInFirstFrame, callback);
137+
}
138+
return ((ArrowFlightPreparedStatement) statement)
139+
.prepareAndExecuteInternal(query, maxRowCount, maxRowsInFirstFrame, callback);
130140
} catch (SQLException e) {
131141
// AvaticaStatement.executeInternal handles RuntimeException and preserves it as the cause of
132142
// the public SQLException. NoSuchStatementException would incorrectly trigger a statement
@@ -181,14 +191,24 @@ void setDefaultConnectionProperties() {
181191
.setTransactionIsolation(Connection.TRANSACTION_NONE);
182192
}
183193

184-
private ArrowFlightMetaStatement getMetaStatement(StatementHandle statementHandle) {
185-
AvaticaStatement statement = connection.statementMap.get(statementHandle.id);
186-
if (statement instanceof ArrowFlightMetaStatement) {
187-
return (ArrowFlightMetaStatement) statement;
194+
private AvaticaStatement getStatement(final StatementHandle statementHandle) {
195+
final AvaticaStatement statement = connection.statementMap.get(statementHandle.id);
196+
if (statement instanceof ArrowFlightStatement
197+
|| statement instanceof ArrowFlightPreparedStatement) {
198+
return statement;
188199
}
189200
throw new IllegalStateException("Statement not found: " + statementHandle);
190201
}
191202

203+
private ArrowFlightPreparedStatement getPreparedStatement(final StatementHandle statementHandle) {
204+
final AvaticaStatement statement = getStatement(statementHandle);
205+
if (statement instanceof ArrowFlightPreparedStatement) {
206+
return (ArrowFlightPreparedStatement) statement;
207+
}
208+
throw new IllegalStateException(
209+
"Statement operation is not supported for handle: " + statementHandle);
210+
}
211+
192212
static Signature buildDefaultSignature() {
193213
return buildSignature(null, StatementType.SELECT);
194214
}

‎flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java‎

Lines changed: 0 additions & 61 deletions
This file was deleted.

‎flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java‎

Lines changed: 4 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838

3939
/** Arrow Flight JDBC's implementation {@link java.sql.PreparedStatement}. */
4040
public class ArrowFlightPreparedStatement extends AvaticaPreparedStatement
41-
implements ArrowFlightMetaStatement {
41+
implements ArrowFlightInfoStatement {
4242

4343
private ArrowFlightSqlClientHandler.PreparedStatement preparedStatement;
4444

@@ -66,7 +66,7 @@ public ArrowFlightConnection getConnection() throws SQLException {
6666
return (ArrowFlightConnection) super.getConnection();
6767
}
6868

69-
ExecuteResult prepareAndExecute(final PrepareCallback callback) throws SQLException {
69+
ExecuteResult prepareAndExecuteInternal(final PrepareCallback callback) throws SQLException {
7070
ensurePrepared();
7171
final StatementType statementType = preparedStatement.getType();
7272
final long updateCount =
@@ -81,8 +81,7 @@ ExecuteResult prepareAndExecute(final PrepareCallback callback) throws SQLExcept
8181
return new ExecuteResult(Collections.singletonList(metaResultSet));
8282
}
8383

84-
@Override
85-
public ExecuteResult prepareAndExecute(
84+
ExecuteResult prepareAndExecuteInternal(
8685
final String query,
8786
final long maxRowCount,
8887
final int maxRowsInFirstFrame,
@@ -93,7 +92,7 @@ public ExecuteResult prepareAndExecute(
9392
.withQuery(query)
9493
.withExistingStatement(this)
9594
.build()
96-
.prepareAndExecute(callback);
95+
.prepareAndExecuteInternal(callback);
9796
}
9897

9998
Schema getDataSetSchema() {
@@ -163,25 +162,6 @@ ExecuteBatchResult executeBatchWithTypedValues(
163162
return new ExecuteBatchResult(updatedCounts);
164163
}
165164

166-
@Override
167-
public ExecuteResult execute(
168-
final StatementHandle statementHandle,
169-
final List<TypedValue> typedValues,
170-
final long maxRowCount) {
171-
return executeWithTypedValues(statementHandle, typedValues, maxRowCount);
172-
}
173-
174-
@Override
175-
public ExecuteBatchResult executeBatch(
176-
final StatementHandle statementHandle, final List<List<TypedValue>> parameterValuesList) {
177-
return executeBatchWithTypedValues(statementHandle, parameterValuesList);
178-
}
179-
180-
@Override
181-
public void closeStatement() {
182-
closePreparedResources();
183-
}
184-
185165
@Override
186166
public FlightInfo executeFlightInfoQuery() throws SQLException {
187167
ensurePrepared();

‎flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
import org.apache.calcite.avatica.Meta.StatementType;
3434

3535
/** A SQL statement for querying data from an Arrow Flight server. */
36-
public class ArrowFlightStatement extends AvaticaStatement implements ArrowFlightMetaStatement {
36+
public class ArrowFlightStatement extends AvaticaStatement implements ArrowFlightInfoStatement {
3737

3838
ArrowFlightStatement(
3939
final ArrowFlightConnection connection,
@@ -59,22 +59,19 @@ public ArrowFlightConnection getConnection() throws SQLException {
5959
return (ArrowFlightConnection) super.getConnection();
6060
}
6161

62-
@Override
63-
public ExecuteResult prepareAndExecute(
62+
ExecuteResult prepareAndExecuteInternal(
6463
final String query,
6564
final long maxRowCount,
6665
final int maxRowsInFirstFrame,
6766
final PrepareCallback callback)
6867
throws SQLException {
6968
// Keep Avatica Statement.execute(String) behavior: Avatica calls Meta.prepareAndExecute,
7069
// which resolves to this statement hook.
71-
this.closeStatement();
72-
7370
return ArrowFlightPreparedStatement.builder(getConnection())
7471
.withQuery(query)
7572
.withExistingStatement(this)
7673
.build()
77-
.prepareAndExecute(callback);
74+
.prepareAndExecuteInternal(callback);
7875
}
7976

8077
@Override
@@ -192,7 +189,7 @@ private void switchToDirectStatementMode() throws SQLException {
192189
if (existingStatement instanceof ArrowFlightPreparedStatement) {
193190
// Release resources from previously attached statement implementation before switching back
194191
// to direct statement mode for executeQuery/executeUpdate.
195-
((ArrowFlightPreparedStatement) existingStatement).closeStatement();
192+
((ArrowFlightPreparedStatement) existingStatement).closePreparedResources();
196193
}
197194
connection.statementMap.put(handle.id, this);
198195
}

‎flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.arrow.driver.jdbc;
18+
19+
import static org.junit.jupiter.api.Assertions.assertThrows;
20+
21+
import java.util.List;
22+
import org.apache.calcite.avatica.Meta.PrepareCallback;
23+
import org.apache.calcite.avatica.Meta.StatementHandle;
24+
import org.junit.jupiter.api.Test;
25+
26+
public class ArrowFlightStatementApiTest {
27+
28+
@Test
29+
public void testMetaOrchestrationMethodsAreNotPublicStatementApis() {
30+
assertNotPublic(
31+
ArrowFlightStatement.class,
32+
"prepareAndExecute",
33+
String.class,
34+
long.class,
35+
int.class,
36+
PrepareCallback.class);
37+
assertNotPublic(
38+
ArrowFlightPreparedStatement.class,
39+
"prepareAndExecute",
40+
String.class,
41+
long.class,
42+
int.class,
43+
PrepareCallback.class);
44+
assertNotPublic(
45+
ArrowFlightPreparedStatement.class,
46+
"execute",
47+
StatementHandle.class,
48+
List.class,
49+
long.class);
50+
assertNotPublic(
51+
ArrowFlightPreparedStatement.class, "executeBatch", StatementHandle.class, List.class);
52+
assertNotPublic(ArrowFlightPreparedStatement.class, "closeStatement");
53+
assertNotPublic(
54+
ArrowFlightStatement.class,
55+
"prepareAndExecuteInternal",
56+
String.class,
57+
long.class,
58+
int.class,
59+
PrepareCallback.class);
60+
assertNotPublic(
61+
ArrowFlightPreparedStatement.class,
62+
"prepareAndExecuteInternal",
63+
String.class,
64+
long.class,
65+
int.class,
66+
PrepareCallback.class);
67+
assertNotPublic(
68+
ArrowFlightPreparedStatement.class,
69+
"executeWithTypedValues",
70+
StatementHandle.class,
71+
List.class,
72+
long.class);
73+
assertNotPublic(
74+
ArrowFlightPreparedStatement.class,
75+
"executeBatchWithTypedValues",
76+
StatementHandle.class,
77+
List.class);
78+
assertNotPublic(ArrowFlightPreparedStatement.class, "closePreparedResources");
79+
}
80+
81+
private static void assertNotPublic(
82+
final Class<?> statementClass, final String methodName, final Class<?>... parameterTypes) {
83+
assertThrows(
84+
NoSuchMethodException.class,
85+
() -> statementClass.getMethod(methodName, parameterTypes),
86+
methodName + " must remain an internal orchestration method");
87+
}
88+
}

‎flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ public void testExecuteQueryMapsPreparedStatementCloseFailureToSQLException() {
248248
final ArrowFlightPreparedStatement preparedStatement = mock(ArrowFlightPreparedStatement.class);
249249
final FlightRuntimeException closeFailure =
250250
CallStatus.INTERNAL.withDescription("Prepared statement close failed").toRuntimeException();
251-
doThrow(closeFailure).when(preparedStatement).closeStatement();
251+
doThrow(closeFailure).when(preparedStatement).closePreparedResources();
252252
arrowConnection.statementMap.put(arrowStatement.handle.id, preparedStatement);
253253

254254
try {
@@ -268,7 +268,7 @@ public void testExecuteLargeUpdateMapsPreparedStatementCloseFailureToSQLExceptio
268268
final ArrowFlightPreparedStatement preparedStatement = mock(ArrowFlightPreparedStatement.class);
269269
final FlightRuntimeException closeFailure =
270270
CallStatus.INTERNAL.withDescription("Prepared statement close failed").toRuntimeException();
271-
doThrow(closeFailure).when(preparedStatement).closeStatement();
271+
doThrow(closeFailure).when(preparedStatement).closePreparedResources();
272272
arrowConnection.statementMap.put(arrowStatement.handle.id, preparedStatement);
273273

274274
try {

0 commit comments

Comments
 (0)