Skip to content

Commit db03518

Browse files
kokosingkewang1024
kokosing
and
kewang1024
committed
Pass JdbcIdentity To DriverConnectionFactory#openConnection
Co-Authored-By: Ke Wang <[email protected]>
1 parent 6b3c211 commit db03518

File tree

17 files changed

+204
-111
lines changed

17 files changed

+204
-111
lines changed

‎presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/BaseJdbcClient.java

Lines changed: 37 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -118,9 +118,9 @@ public void destroy()
118118
}
119119

120120
@Override
121-
public Set<String> getSchemaNames()
121+
public Set<String> getSchemaNames(JdbcIdentity identity)
122122
{
123-
try (Connection connection = connectionFactory.openConnection();
123+
try (Connection connection = connectionFactory.openConnection(identity);
124124
ResultSet resultSet = connection.getMetaData().getSchemas()) {
125125
ImmutableSet.Builder<String> schemaNames = ImmutableSet.builder();
126126
while (resultSet.next()) {
@@ -138,9 +138,9 @@ public Set<String> getSchemaNames()
138138
}
139139

140140
@Override
141-
public List<SchemaTableName> getTableNames(@Nullable String schema)
141+
public List<SchemaTableName> getTableNames(JdbcIdentity identity, @Nullable String schema)
142142
{
143-
try (Connection connection = connectionFactory.openConnection()) {
143+
try (Connection connection = connectionFactory.openConnection(identity)) {
144144
DatabaseMetaData metadata = connection.getMetaData();
145145
if (metadata.storesUpperCaseIdentifiers() && (schema != null)) {
146146
schema = schema.toUpperCase(ENGLISH);
@@ -160,9 +160,9 @@ public List<SchemaTableName> getTableNames(@Nullable String schema)
160160

161161
@Nullable
162162
@Override
163-
public JdbcTableHandle getTableHandle(SchemaTableName schemaTableName)
163+
public JdbcTableHandle getTableHandle(JdbcIdentity identity, SchemaTableName schemaTableName)
164164
{
165-
try (Connection connection = connectionFactory.openConnection()) {
165+
try (Connection connection = connectionFactory.openConnection(identity)) {
166166
DatabaseMetaData metadata = connection.getMetaData();
167167
String jdbcSchemaName = schemaTableName.getSchemaName();
168168
String jdbcTableName = schemaTableName.getTableName();
@@ -197,7 +197,7 @@ public JdbcTableHandle getTableHandle(SchemaTableName schemaTableName)
197197
@Override
198198
public List<JdbcColumnHandle> getColumns(ConnectorSession session, JdbcTableHandle tableHandle)
199199
{
200-
try (Connection connection = connectionFactory.openConnection()) {
200+
try (Connection connection = connectionFactory.openConnection(JdbcIdentity.from(session))) {
201201
try (ResultSet resultSet = getColumns(tableHandle, connection.getMetaData())) {
202202
List<JdbcColumnHandle> columns = new ArrayList<>();
203203
while (resultSet.next()) {
@@ -232,7 +232,7 @@ public Optional<ReadMapping> toPrestoType(ConnectorSession session, JdbcTypeHand
232232
}
233233

234234
@Override
235-
public ConnectorSplitSource getSplits(JdbcTableLayoutHandle layoutHandle)
235+
public ConnectorSplitSource getSplits(JdbcIdentity identity, JdbcTableLayoutHandle layoutHandle)
236236
{
237237
JdbcTableHandle tableHandle = layoutHandle.getTable();
238238
JdbcSplit jdbcSplit = new JdbcSplit(
@@ -246,10 +246,10 @@ public ConnectorSplitSource getSplits(JdbcTableLayoutHandle layoutHandle)
246246
}
247247

248248
@Override
249-
public Connection getConnection(JdbcSplit split)
249+
public Connection getConnection(JdbcIdentity identity, JdbcSplit split)
250250
throws SQLException
251251
{
252-
Connection connection = connectionFactory.openConnection();
252+
Connection connection = connectionFactory.openConnection(identity);
253253
try {
254254
connection.setReadOnly(true);
255255
}
@@ -276,10 +276,10 @@ public PreparedStatement buildSql(Connection connection, JdbcSplit split, List<J
276276
}
277277

278278
@Override
279-
public void createTable(ConnectorTableMetadata tableMetadata)
279+
public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
280280
{
281281
try {
282-
createTable(tableMetadata, null, tableMetadata.getTable().getTableName());
282+
createTable(tableMetadata, session, tableMetadata.getTable().getTableName());
283283
}
284284
catch (SQLException e) {
285285
throw new PrestoException(JDBC_ERROR, e);
@@ -315,11 +315,12 @@ protected JdbcOutputTableHandle createTable(ConnectorTableMetadata tableMetadata
315315
String schema = schemaTableName.getSchemaName();
316316
String table = schemaTableName.getTableName();
317317

318-
if (!getSchemaNames().contains(schema)) {
318+
JdbcIdentity identity = JdbcIdentity.from(session);
319+
if (!getSchemaNames(identity).contains(schema)) {
319320
throw new PrestoException(NOT_FOUND, "Schema not found: " + schema);
320321
}
321322

322-
try (Connection connection = connectionFactory.openConnection()) {
323+
try (Connection connection = connectionFactory.openConnection(identity)) {
323324
boolean uppercase = connection.getMetaData().storesUpperCaseIdentifiers();
324325
if (uppercase) {
325326
schema = schema.toUpperCase(ENGLISH);
@@ -377,24 +378,26 @@ protected String generateTemporaryTableName()
377378
return "tmp_presto_" + UUID.randomUUID().toString().replace("-", "");
378379
}
379380

381+
//todo
380382
@Override
381-
public void commitCreateTable(JdbcOutputTableHandle handle)
383+
public void commitCreateTable(JdbcIdentity identity, JdbcOutputTableHandle handle)
382384
{
383385
renameTable(
386+
identity,
384387
handle.getCatalogName(),
385388
new SchemaTableName(handle.getSchemaName(), handle.getTemporaryTableName()),
386389
new SchemaTableName(handle.getSchemaName(), handle.getTableName()));
387390
}
388391

389392
@Override
390-
public void renameTable(JdbcTableHandle handle, SchemaTableName newTable)
393+
public void renameTable(JdbcIdentity identity, JdbcTableHandle handle, SchemaTableName newTable)
391394
{
392-
renameTable(handle.getCatalogName(), handle.getSchemaTableName(), newTable);
395+
renameTable(identity, handle.getCatalogName(), handle.getSchemaTableName(), newTable);
393396
}
394397

395-
protected void renameTable(String catalogName, SchemaTableName oldTable, SchemaTableName newTable)
398+
protected void renameTable(JdbcIdentity identity, String catalogName, SchemaTableName oldTable, SchemaTableName newTable)
396399
{
397-
try (Connection connection = connectionFactory.openConnection()) {
400+
try (Connection connection = connectionFactory.openConnection(identity)) {
398401
DatabaseMetaData metadata = connection.getMetaData();
399402
String schemaName = oldTable.getSchemaName();
400403
String tableName = oldTable.getTableName();
@@ -418,21 +421,21 @@ protected void renameTable(String catalogName, SchemaTableName oldTable, SchemaT
418421
}
419422

420423
@Override
421-
public void finishInsertTable(JdbcOutputTableHandle handle)
424+
public void finishInsertTable(JdbcIdentity identity, JdbcOutputTableHandle handle)
422425
{
423426
String temporaryTable = quoted(handle.getCatalogName(), handle.getSchemaName(), handle.getTemporaryTableName());
424427
String targetTable = quoted(handle.getCatalogName(), handle.getSchemaName(), handle.getTableName());
425428
String insertSql = format("INSERT INTO %s SELECT * FROM %s", targetTable, temporaryTable);
426429
String cleanupSql = "DROP TABLE " + temporaryTable;
427430

428-
try (Connection connection = getConnection(handle)) {
431+
try (Connection connection = getConnection(identity, handle)) {
429432
execute(connection, insertSql);
430433
}
431434
catch (SQLException e) {
432435
throw new PrestoException(JDBC_ERROR, e);
433436
}
434437

435-
try (Connection connection = getConnection(handle)) {
438+
try (Connection connection = getConnection(identity, handle)) {
436439
execute(connection, cleanupSql);
437440
}
438441
catch (SQLException e) {
@@ -441,9 +444,9 @@ public void finishInsertTable(JdbcOutputTableHandle handle)
441444
}
442445

443446
@Override
444-
public void addColumn(JdbcTableHandle handle, ColumnMetadata column)
447+
public void addColumn(JdbcIdentity identity, JdbcTableHandle handle, ColumnMetadata column)
445448
{
446-
try (Connection connection = connectionFactory.openConnection()) {
449+
try (Connection connection = connectionFactory.openConnection(identity)) {
447450
String schema = handle.getSchemaName();
448451
String table = handle.getTableName();
449452
String columnName = column.getName();
@@ -465,9 +468,9 @@ public void addColumn(JdbcTableHandle handle, ColumnMetadata column)
465468
}
466469

467470
@Override
468-
public void renameColumn(JdbcTableHandle handle, JdbcColumnHandle jdbcColumn, String newColumnName)
471+
public void renameColumn(JdbcIdentity identity, JdbcTableHandle handle, JdbcColumnHandle jdbcColumn, String newColumnName)
469472
{
470-
try (Connection connection = connectionFactory.openConnection()) {
473+
try (Connection connection = connectionFactory.openConnection(identity)) {
471474
DatabaseMetaData metadata = connection.getMetaData();
472475
if (metadata.storesUpperCaseIdentifiers()) {
473476
newColumnName = newColumnName.toUpperCase(ENGLISH);
@@ -485,9 +488,9 @@ public void renameColumn(JdbcTableHandle handle, JdbcColumnHandle jdbcColumn, St
485488
}
486489

487490
@Override
488-
public void dropColumn(JdbcTableHandle handle, JdbcColumnHandle column)
491+
public void dropColumn(JdbcIdentity identity, JdbcTableHandle handle, JdbcColumnHandle column)
489492
{
490-
try (Connection connection = connectionFactory.openConnection()) {
493+
try (Connection connection = connectionFactory.openConnection(identity)) {
491494
String sql = format(
492495
"ALTER TABLE %s DROP COLUMN %s",
493496
quoted(handle.getCatalogName(), handle.getSchemaName(), handle.getTableName()),
@@ -500,13 +503,13 @@ public void dropColumn(JdbcTableHandle handle, JdbcColumnHandle column)
500503
}
501504

502505
@Override
503-
public void dropTable(JdbcTableHandle handle)
506+
public void dropTable(JdbcIdentity identity, JdbcTableHandle handle)
504507
{
505508
StringBuilder sql = new StringBuilder()
506509
.append("DROP TABLE ")
507510
.append(quoted(handle.getCatalogName(), handle.getSchemaName(), handle.getTableName()));
508511

509-
try (Connection connection = connectionFactory.openConnection()) {
512+
try (Connection connection = connectionFactory.openConnection(identity)) {
510513
execute(connection, sql.toString());
511514
}
512515
catch (SQLException e) {
@@ -515,9 +518,9 @@ public void dropTable(JdbcTableHandle handle)
515518
}
516519

517520
@Override
518-
public void rollbackCreateTable(JdbcOutputTableHandle handle)
521+
public void rollbackCreateTable(JdbcIdentity identity, JdbcOutputTableHandle handle)
519522
{
520-
dropTable(new JdbcTableHandle(
523+
dropTable(identity, new JdbcTableHandle(
521524
handle.getConnectorId(),
522525
new SchemaTableName(handle.getSchemaName(), handle.getTemporaryTableName()),
523526
handle.getCatalogName(),
@@ -537,10 +540,10 @@ public String buildInsertSql(JdbcOutputTableHandle handle)
537540
}
538541

539542
@Override
540-
public Connection getConnection(JdbcOutputTableHandle handle)
543+
public Connection getConnection(JdbcIdentity identity, JdbcOutputTableHandle handle)
541544
throws SQLException
542545
{
543-
return connectionFactory.openConnection();
546+
return connectionFactory.openConnection(identity);
544547
}
545548

546549
@Override

‎presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/ConnectionFactory.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
public interface ConnectionFactory
2121
extends AutoCloseable
2222
{
23-
Connection openConnection()
23+
Connection openConnection(JdbcIdentity identity)
2424
throws SQLException;
2525

2626
@Override

‎presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/DriverConnectionFactory.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ public DriverConnectionFactory(Driver driver, String connectionUrl, Properties c
5454
}
5555

5656
@Override
57-
public Connection openConnection()
57+
public Connection openConnection(JdbcIdentity identity)
5858
throws SQLException
5959
{
6060
Connection connection = driver.connect(connectionUrl, connectionProperties);

‎presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcClient.java

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -33,25 +33,25 @@
3333

3434
public interface JdbcClient
3535
{
36-
default boolean schemaExists(String schema)
36+
default boolean schemaExists(JdbcIdentity identity, String schema)
3737
{
38-
return getSchemaNames().contains(schema);
38+
return getSchemaNames(identity).contains(schema);
3939
}
4040

41-
Set<String> getSchemaNames();
41+
Set<String> getSchemaNames(JdbcIdentity identity);
4242

43-
List<SchemaTableName> getTableNames(@Nullable String schema);
43+
List<SchemaTableName> getTableNames(JdbcIdentity identity, @Nullable String schema);
4444

4545
@Nullable
46-
JdbcTableHandle getTableHandle(SchemaTableName schemaTableName);
46+
JdbcTableHandle getTableHandle(JdbcIdentity identity, SchemaTableName schemaTableName);
4747

4848
List<JdbcColumnHandle> getColumns(ConnectorSession session, JdbcTableHandle tableHandle);
4949

5050
Optional<ReadMapping> toPrestoType(ConnectorSession session, JdbcTypeHandle typeHandle);
5151

52-
ConnectorSplitSource getSplits(JdbcTableLayoutHandle layoutHandle);
52+
ConnectorSplitSource getSplits(JdbcIdentity identity, JdbcTableLayoutHandle layoutHandle);
5353

54-
Connection getConnection(JdbcSplit split)
54+
Connection getConnection(JdbcIdentity identity, JdbcSplit split)
5555
throws SQLException;
5656

5757
default void abortReadConnection(Connection connection)
@@ -63,31 +63,31 @@ default void abortReadConnection(Connection connection)
6363
PreparedStatement buildSql(Connection connection, JdbcSplit split, List<JdbcColumnHandle> columnHandles)
6464
throws SQLException;
6565

66-
void addColumn(JdbcTableHandle handle, ColumnMetadata column);
66+
void addColumn(JdbcIdentity identity, JdbcTableHandle handle, ColumnMetadata column);
6767

68-
void dropColumn(JdbcTableHandle handle, JdbcColumnHandle column);
68+
void dropColumn(JdbcIdentity identity, JdbcTableHandle handle, JdbcColumnHandle column);
6969

70-
void renameColumn(JdbcTableHandle handle, JdbcColumnHandle jdbcColumn, String newColumnName);
70+
void renameColumn(JdbcIdentity identity, JdbcTableHandle handle, JdbcColumnHandle jdbcColumn, String newColumnName);
7171

72-
void renameTable(JdbcTableHandle handle, SchemaTableName newTableName);
72+
void renameTable(JdbcIdentity identity, JdbcTableHandle handle, SchemaTableName newTableName);
7373

74-
void createTable(ConnectorTableMetadata tableMetadata);
74+
void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata);
7575

7676
JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata);
7777

78-
void commitCreateTable(JdbcOutputTableHandle handle);
78+
void commitCreateTable(JdbcIdentity identity, JdbcOutputTableHandle handle);
7979

8080
JdbcOutputTableHandle beginInsertTable(ConnectorSession session, ConnectorTableMetadata tableMetadata);
8181

82-
void finishInsertTable(JdbcOutputTableHandle handle);
82+
void finishInsertTable(JdbcIdentity identity, JdbcOutputTableHandle handle);
8383

84-
void dropTable(JdbcTableHandle jdbcTableHandle);
84+
void dropTable(JdbcIdentity identity, JdbcTableHandle jdbcTableHandle);
8585

86-
void rollbackCreateTable(JdbcOutputTableHandle handle);
86+
void rollbackCreateTable(JdbcIdentity identity, JdbcOutputTableHandle handle);
8787

8888
String buildInsertSql(JdbcOutputTableHandle handle);
8989

90-
Connection getConnection(JdbcOutputTableHandle handle)
90+
Connection getConnection(JdbcIdentity identity, JdbcOutputTableHandle handle)
9191
throws SQLException;
9292

9393
PreparedStatement getPreparedStatement(Connection connection, String sql)

0 commit comments

Comments
 (0)