diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcClient.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcClient.java index befeb9298fda..9b1c1d13e9d8 100644 --- a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcClient.java +++ b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcClient.java @@ -63,6 +63,7 @@ import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Iterables.getOnlyElement; +import static io.prestosql.plugin.jdbc.BaseJdbcSessionProperties.getUnsupportedTypeHandlingStrategy; import static io.prestosql.plugin.jdbc.JdbcErrorCode.JDBC_ERROR; import static io.prestosql.plugin.jdbc.StandardColumnMappings.bigintWriteFunction; import static io.prestosql.plugin.jdbc.StandardColumnMappings.booleanWriteFunction; @@ -248,12 +249,22 @@ public List getColumns(ConnectorSession session, JdbcTableHand resultSet.getInt("DECIMAL_DIGITS"), Optional.empty()); Optional columnMapping = toPrestoType(session, connection, typeHandle); - // skip unsupported column types + String columnName = resultSet.getString("COLUMN_NAME"); if (columnMapping.isPresent()) { - String columnName = resultSet.getString("COLUMN_NAME"); boolean nullable = (resultSet.getInt("NULLABLE") != columnNoNulls); columns.add(new JdbcColumnHandle(columnName, typeHandle, columnMapping.get().getType(), nullable)); } + else { + UnsupportedTypeHandlingStrategy unsupportedTypeHandlingStrategy = getUnsupportedTypeHandlingStrategy(session); + switch (unsupportedTypeHandlingStrategy) { + case IGNORE: + break; + case FAIL: + throw new PrestoException(JDBC_ERROR, "Unsupported data type for column: " + columnName); + default: + throw new IllegalStateException("Unknown unsupported type handling strategy: " + unsupportedTypeHandlingStrategy); + } + } } if (columns.isEmpty()) { // In rare cases (e.g. PostgreSQL) a table might have no columns. diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcConfig.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcConfig.java index 6b507098ccf9..7e68ec63a562 100644 --- a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcConfig.java +++ b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcConfig.java @@ -14,6 +14,7 @@ package io.prestosql.plugin.jdbc; import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; import io.airlift.configuration.ConfigSecuritySensitive; import io.airlift.units.Duration; import io.airlift.units.MinDuration; @@ -21,6 +22,7 @@ import javax.annotation.Nullable; import javax.validation.constraints.NotNull; +import static io.prestosql.plugin.jdbc.UnsupportedTypeHandlingStrategy.FAIL; import static java.util.concurrent.TimeUnit.MINUTES; public class BaseJdbcConfig @@ -32,6 +34,7 @@ public class BaseJdbcConfig private String passwordCredentialName; private boolean caseInsensitiveNameMatching; private Duration caseInsensitiveNameMatchingCacheTtl = new Duration(1, MINUTES); + private UnsupportedTypeHandlingStrategy unsupportedTypeHandlingStrategy = FAIL; @NotNull public String getConnectionUrl() @@ -124,4 +127,18 @@ public BaseJdbcConfig setCaseInsensitiveNameMatchingCacheTtl(Duration caseInsens this.caseInsensitiveNameMatchingCacheTtl = caseInsensitiveNameMatchingCacheTtl; return this; } + + @NotNull + public UnsupportedTypeHandlingStrategy getUnsupportedTypeHandlingStrategy() + { + return unsupportedTypeHandlingStrategy; + } + + @Config("unsupported-type.handling-strategy") + @ConfigDescription("Configures how unsupported column data types should be handled") + public BaseJdbcConfig setUnsupportedTypeHandlingStrategy(UnsupportedTypeHandlingStrategy unsupportedTypeHandlingStrategy) + { + this.unsupportedTypeHandlingStrategy = unsupportedTypeHandlingStrategy; + return this; + } } diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcSessionProperties.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcSessionProperties.java new file mode 100644 index 000000000000..2bdba27c354f --- /dev/null +++ b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcSessionProperties.java @@ -0,0 +1,55 @@ +/* + * Licensed 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 io.prestosql.plugin.jdbc; + +import com.google.common.collect.ImmutableList; +import io.prestosql.spi.connector.ConnectorSession; +import io.prestosql.spi.session.PropertyMetadata; + +import javax.inject.Inject; + +import java.util.List; + +import static io.prestosql.spi.session.PropertyMetadata.enumProperty; + +public final class BaseJdbcSessionProperties + implements SessionPropertiesProvider +{ + public static final String UNSUPPORTED_TYPE_HANDLING_STRATEGY = "unsupported_type_handling_strategy"; + + private final List> sessionProperties; + + @Inject + public BaseJdbcSessionProperties(BaseJdbcConfig config) + { + sessionProperties = ImmutableList.of( + enumProperty( + UNSUPPORTED_TYPE_HANDLING_STRATEGY, + "Configures how unsupported column data types should be handled", + UnsupportedTypeHandlingStrategy.class, + config.getUnsupportedTypeHandlingStrategy(), + false)); + } + + @Override + public List> getSessionProperties() + { + return sessionProperties; + } + + public static UnsupportedTypeHandlingStrategy getUnsupportedTypeHandlingStrategy(ConnectorSession session) + { + return session.getProperty(UNSUPPORTED_TYPE_HANDLING_STRATEGY, UnsupportedTypeHandlingStrategy.class); + } +} diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcConnector.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcConnector.java index 54067c6bf880..b636dd33e04c 100644 --- a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcConnector.java +++ b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcConnector.java @@ -13,6 +13,7 @@ */ package io.prestosql.plugin.jdbc; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import io.airlift.bootstrap.LifeCycleManager; import io.airlift.log.Logger; @@ -25,10 +26,12 @@ import io.prestosql.spi.connector.ConnectorSplitManager; import io.prestosql.spi.connector.ConnectorTransactionHandle; import io.prestosql.spi.procedure.Procedure; +import io.prestosql.spi.session.PropertyMetadata; import io.prestosql.spi.transaction.IsolationLevel; import javax.inject.Inject; +import java.util.List; import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -53,6 +56,7 @@ public class JdbcConnector private final JdbcPageSinkProvider jdbcPageSinkProvider; private final Optional accessControl; private final Set procedures; + private final Set sessionProperties; private final ConcurrentMap transactions = new ConcurrentHashMap<>(); @@ -64,7 +68,8 @@ public JdbcConnector( JdbcRecordSetProvider jdbcRecordSetProvider, JdbcPageSinkProvider jdbcPageSinkProvider, Optional accessControl, - Set procedures) + Set procedures, + Set sessionProperties) { this.lifeCycleManager = requireNonNull(lifeCycleManager, "lifeCycleManager is null"); this.jdbcMetadataFactory = requireNonNull(jdbcMetadataFactory, "jdbcMetadataFactory is null"); @@ -73,6 +78,7 @@ public JdbcConnector( this.jdbcPageSinkProvider = requireNonNull(jdbcPageSinkProvider, "jdbcPageSinkProvider is null"); this.accessControl = requireNonNull(accessControl, "accessControl is null"); this.procedures = ImmutableSet.copyOf(requireNonNull(procedures, "procedures is null")); + this.sessionProperties = requireNonNull(sessionProperties, "sessionProperties is null"); } @Override @@ -142,6 +148,14 @@ public Set getProcedures() return procedures; } + @Override + public List> getSessionProperties() + { + ImmutableList.Builder builder = ImmutableList.>builder(); + this.sessionProperties.forEach(sessionPropertiesProvider -> builder.addAll(sessionPropertiesProvider.getSessionProperties())); + return builder.build(); + } + @Override public final void shutdown() { diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcModule.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcModule.java index ddeecefa30e5..fd3176815dd8 100644 --- a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcModule.java +++ b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcModule.java @@ -21,6 +21,7 @@ import com.google.inject.Singleton; import io.prestosql.plugin.jdbc.jmx.StatisticsAwareConnectionFactory; import io.prestosql.plugin.jdbc.jmx.StatisticsAwareJdbcClient; +import com.google.inject.multibindings.Multibinder; import io.prestosql.spi.connector.ConnectorAccessControl; import io.prestosql.spi.procedure.Procedure; @@ -45,6 +46,8 @@ public void configure(Binder binder) { newOptionalBinder(binder, ConnectorAccessControl.class); newSetBinder(binder, Procedure.class); + Multibinder sessionProperties = newSetBinder(binder, SessionPropertiesProvider.class); + sessionProperties.addBinding().to(BaseJdbcSessionProperties.class).in(Scopes.SINGLETON); binder.bind(JdbcMetadataFactory.class).in(Scopes.SINGLETON); binder.bind(JdbcSplitManager.class).in(Scopes.SINGLETON); binder.bind(JdbcRecordSetProvider.class).in(Scopes.SINGLETON); diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/SessionPropertiesProvider.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/SessionPropertiesProvider.java new file mode 100644 index 000000000000..63a4d699ca9c --- /dev/null +++ b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/SessionPropertiesProvider.java @@ -0,0 +1,23 @@ +/* + * Licensed 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 io.prestosql.plugin.jdbc; + +import io.prestosql.spi.session.PropertyMetadata; + +import java.util.List; + +public interface SessionPropertiesProvider +{ + List> getSessionProperties(); +} diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/UnsupportedTypeHandlingStrategy.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/UnsupportedTypeHandlingStrategy.java new file mode 100644 index 000000000000..3dc95cb0b2c8 --- /dev/null +++ b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/UnsupportedTypeHandlingStrategy.java @@ -0,0 +1,21 @@ +/* + * Licensed 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 io.prestosql.plugin.jdbc; + +public enum UnsupportedTypeHandlingStrategy +{ + IGNORE, + FAIL, + /**/; +} diff --git a/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/JdbcQueryRunner.java b/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/JdbcQueryRunner.java index ab0c4f4a2855..18d32ea6151a 100644 --- a/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/JdbcQueryRunner.java +++ b/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/JdbcQueryRunner.java @@ -46,6 +46,13 @@ public static DistributedQueryRunner createJdbcQueryRunner(TpchTable... table public static DistributedQueryRunner createJdbcQueryRunner(Iterable> tables) throws Exception + + { + return createJdbcQueryRunner(tables, TestingH2JdbcModule.createProperties()); + } + + public static DistributedQueryRunner createJdbcQueryRunner(Iterable> tables, Map properties) + throws Exception { DistributedQueryRunner queryRunner = null; try { @@ -54,7 +61,6 @@ public static DistributedQueryRunner createJdbcQueryRunner(Iterable queryRunner.installPlugin(new TpchPlugin()); queryRunner.createCatalog("tpch", "tpch"); - Map properties = TestingH2JdbcModule.createProperties(); createSchema(properties, "tpch"); queryRunner.installPlugin(new JdbcPlugin("base-jdbc", new TestingH2JdbcModule())); diff --git a/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/TestBaseJdbcConfig.java b/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/TestBaseJdbcConfig.java index 082b3f65a74d..880c03e4e3b8 100644 --- a/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/TestBaseJdbcConfig.java +++ b/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/TestBaseJdbcConfig.java @@ -20,6 +20,8 @@ import java.util.Map; +import static io.prestosql.plugin.jdbc.UnsupportedTypeHandlingStrategy.FAIL; +import static io.prestosql.plugin.jdbc.UnsupportedTypeHandlingStrategy.IGNORE; import static java.util.concurrent.TimeUnit.MINUTES; import static java.util.concurrent.TimeUnit.SECONDS; @@ -35,7 +37,8 @@ public void testDefaults() .setUserCredentialName(null) .setPasswordCredentialName(null) .setCaseInsensitiveNameMatching(false) - .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, MINUTES))); + .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, MINUTES)) + .setUnsupportedTypeHandlingStrategy(FAIL)); } @Test @@ -49,6 +52,7 @@ public void testExplicitPropertyMappings() .put("password-credential-name", "bar") .put("case-insensitive-name-matching", "true") .put("case-insensitive-name-matching.cache-ttl", "1s") + .put("unsupported-type.handling-strategy", "IGNORE") .build(); BaseJdbcConfig expected = new BaseJdbcConfig() @@ -58,7 +62,8 @@ public void testExplicitPropertyMappings() .setUserCredentialName("foo") .setPasswordCredentialName("bar") .setCaseInsensitiveNameMatching(true) - .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, SECONDS)); + .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, SECONDS)) + .setUnsupportedTypeHandlingStrategy(IGNORE); ConfigAssertions.assertFullMapping(properties, expected); } diff --git a/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/TestJdbcDistributedQueries.java b/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/TestJdbcDistributedQueries.java index b3363262efc0..b05356c50fa9 100644 --- a/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/TestJdbcDistributedQueries.java +++ b/presto-base-jdbc/src/test/java/io/prestosql/plugin/jdbc/TestJdbcDistributedQueries.java @@ -13,21 +13,67 @@ */ package io.prestosql.plugin.jdbc; +import com.google.common.collect.ImmutableList; import io.airlift.tpch.TpchTable; +import io.prestosql.Session; import io.prestosql.tests.AbstractTestQueries; +import io.prestosql.tests.sql.JdbcSqlExecutor; +import io.prestosql.tests.sql.TestTable; +import org.testng.annotations.Test; +import java.util.Map; +import java.util.Optional; +import java.util.Properties; + +import static io.prestosql.plugin.jdbc.BaseJdbcSessionProperties.UNSUPPORTED_TYPE_HANDLING_STRATEGY; import static io.prestosql.plugin.jdbc.JdbcQueryRunner.createJdbcQueryRunner; +import static io.prestosql.plugin.jdbc.UnsupportedTypeHandlingStrategy.IGNORE; +import static io.prestosql.tests.sql.TestTable.TABLE_NAME_PLACEHOLDER; +import static java.lang.String.format; public class TestJdbcDistributedQueries extends AbstractTestQueries { + private final Map properties; + public TestJdbcDistributedQueries() { - super(() -> createJdbcQueryRunner(TpchTable.getTables())); + this(TestingH2JdbcModule.createProperties()); + } + + public TestJdbcDistributedQueries(Map properties) + { + super(() -> createJdbcQueryRunner(ImmutableList.copyOf(TpchTable.getTables()), properties)); + this.properties = properties; } @Override public void testLargeIn() { } + + @Test + public void testFailureOnUnknown() + { + try (TestTable table = new TestTable( + getSqlExecutor(), + "tpch.test_failure_on_unknown_type", + format("CREATE TABLE %s (i int, x GEOMETRY)", TABLE_NAME_PLACEHOLDER), + Optional.of("(1, 'POINT(7 52)')"))) { + assertQueryFails("SELECT * FROM " + table.getName(), "Unsupported data type for column: X"); + assertQuery(onUnsupportedType(IGNORE), "SELECT * FROM " + table.getName(), "VALUES 1"); + } + } + + private Session onUnsupportedType(UnsupportedTypeHandlingStrategy unsupportedTypeHandlingStrategy) + { + return Session.builder(getSession()) + .setCatalogSessionProperty("jdbc", UNSUPPORTED_TYPE_HANDLING_STRATEGY, unsupportedTypeHandlingStrategy.name()) + .build(); + } + + private JdbcSqlExecutor getSqlExecutor() + { + return new JdbcSqlExecutor(properties.get("connection-url"), new Properties()); + } } diff --git a/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlIntegrationSmokeTest.java b/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlIntegrationSmokeTest.java index c7a7f02a3f5e..642f4b10f9a8 100644 --- a/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlIntegrationSmokeTest.java +++ b/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlIntegrationSmokeTest.java @@ -13,6 +13,8 @@ */ package io.prestosql.plugin.mysql; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import io.airlift.testing.mysql.TestingMySqlServer; import io.prestosql.Session; import io.prestosql.testing.MaterializedResult; @@ -51,7 +53,12 @@ public TestMySqlIntegrationSmokeTest() public TestMySqlIntegrationSmokeTest(TestingMySqlServer mysqlServer) { - super(() -> createMySqlQueryRunner(mysqlServer, ORDERS)); + super(() -> createMySqlQueryRunner( + mysqlServer, + ImmutableMap.builder() + .put("unsupported-type.handling-strategy", "IGNORE") + .build(), + ImmutableList.of(ORDERS))); this.mysqlServer = mysqlServer; } diff --git a/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlTypeMapping.java b/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlTypeMapping.java index 819a655d8f50..2ec10c9e163d 100644 --- a/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlTypeMapping.java +++ b/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlTypeMapping.java @@ -75,7 +75,12 @@ public TestMySqlTypeMapping() private TestMySqlTypeMapping(TestingMySqlServer mysqlServer) { - super(() -> createMySqlQueryRunner(mysqlServer, ImmutableMap.of(), ImmutableList.of())); + super(() -> createMySqlQueryRunner( + mysqlServer, + ImmutableMap.builder() + .put("unsupported-type.handling-strategy", "IGNORE") + .build(), + ImmutableList.of())); this.mysqlServer = mysqlServer; } diff --git a/presto-phoenix/src/test/java/io/prestosql/plugin/phoenix/TestPhoenixConfig.java b/presto-phoenix/src/test/java/io/prestosql/plugin/phoenix/TestPhoenixConfig.java index 15b45bbdc92e..b50cdc71b66f 100644 --- a/presto-phoenix/src/test/java/io/prestosql/plugin/phoenix/TestPhoenixConfig.java +++ b/presto-phoenix/src/test/java/io/prestosql/plugin/phoenix/TestPhoenixConfig.java @@ -32,7 +32,7 @@ public void testDefaults() .setConnectionUrl(null) .setResourceConfigFiles("") .setCaseInsensitiveNameMatching(false) - .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, MINUTES))); + .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, MINUTES)); } @Test diff --git a/presto-sqlserver/src/test/java/io/prestosql/plugin/sqlserver/TestSqlServerIntegrationSmokeTest.java b/presto-sqlserver/src/test/java/io/prestosql/plugin/sqlserver/TestSqlServerIntegrationSmokeTest.java index 788c165d1b0f..a92e518507de 100644 --- a/presto-sqlserver/src/test/java/io/prestosql/plugin/sqlserver/TestSqlServerIntegrationSmokeTest.java +++ b/presto-sqlserver/src/test/java/io/prestosql/plugin/sqlserver/TestSqlServerIntegrationSmokeTest.java @@ -13,6 +13,8 @@ */ package io.prestosql.plugin.sqlserver; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import io.prestosql.tests.AbstractTestIntegrationSmokeTest; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; @@ -35,7 +37,12 @@ public TestSqlServerIntegrationSmokeTest() public TestSqlServerIntegrationSmokeTest(TestingSqlServer testingSqlServer) { - super(() -> createSqlServerQueryRunner(testingSqlServer, ORDERS)); + super(() -> createSqlServerQueryRunner( + testingSqlServer, + ImmutableMap.builder() + .put("unsupported-type.handling-strategy", "IGNORE") + .build(), + ImmutableList.of(ORDERS))); this.sqlServer = testingSqlServer; } diff --git a/presto-tests/src/main/java/io/prestosql/tests/sql/TestTable.java b/presto-tests/src/main/java/io/prestosql/tests/sql/TestTable.java index 98e0efa0131f..5c9d032a6cf7 100644 --- a/presto-tests/src/main/java/io/prestosql/tests/sql/TestTable.java +++ b/presto-tests/src/main/java/io/prestosql/tests/sql/TestTable.java @@ -13,21 +13,48 @@ */ package io.prestosql.tests.sql; +import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; +import static java.lang.String.format; + public class TestTable implements AutoCloseable { + public static final String TABLE_NAME_PLACEHOLDER = "{TABLE_NAME}"; + private final SqlExecutor sqlExecutor; private final String name; private static final AtomicInteger instanceCounter = new AtomicInteger(); public TestTable(SqlExecutor sqlExecutor, String namePrefix, String createDdlTemplate) + { + this(sqlExecutor, namePrefix, createDdlTemplate, Optional.empty()); + } + + public TestTable(SqlExecutor sqlExecutor, String namePrefix, String createDdlTemplate, Optional data) { this.sqlExecutor = sqlExecutor; this.name = namePrefix + "_" + instanceCounter.incrementAndGet(); - sqlExecutor.execute(createDdlTemplate.replace("{TABLE_NAME}", this.name)); + sqlExecutor.execute(createDdlTemplate.replace(TABLE_NAME_PLACEHOLDER, this.name)); + if (!data.isPresent()) { + return; + } + try { + sqlExecutor.execute(format("INSERT INTO %s VALUES %s", name, data.get())); + } + catch (Exception e) { + try { + close(); + } + catch (Exception innerException) { + if (e != innerException) { + e.addSuppressed(innerException); + } + } + throw e; + } } public String getName()