From a34eb43e77987c123dae10243f3de4b948a41ff4 Mon Sep 17 00:00:00 2001 From: taherk77 Date: Sat, 25 Mar 2023 16:46:54 +0530 Subject: [PATCH 1/8] Added Spanner connector --- plugin/trino-spanner/pom.xml | 149 ++++++++++ .../trino/plugin/spanner/SpannerClient.java | 274 ++++++++++++++++++ .../trino/plugin/spanner/SpannerConfig.java | 19 ++ .../trino/plugin/spanner/SpannerModule.java | 98 +++++++ .../trino/plugin/spanner/SpannerPlugin.java | 28 ++ .../spanner/SpannerTableProperties.java | 65 +++++ .../src/test/java/SpannerSqlQueryRunner.java | 158 ++++++++++ pom.xml | 1 + 8 files changed, 792 insertions(+) create mode 100644 plugin/trino-spanner/pom.xml create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerPlugin.java create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java create mode 100644 plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java diff --git a/plugin/trino-spanner/pom.xml b/plugin/trino-spanner/pom.xml new file mode 100644 index 000000000000..af2be234f422 --- /dev/null +++ b/plugin/trino-spanner/pom.xml @@ -0,0 +1,149 @@ + + + 4.0.0 + + trino-root + io.trino + 411-SNAPSHOT + ../../pom.xml + + trino-spanner + + + 17 + 17 + + + + + io.trino + trino-base-jdbc + + + io.trino + trino-plugin-toolkit + + + io.airlift + configuration + + + io.airlift + log + + + com.google.cloud + google-cloud-spanner-jdbc + 2.9.9 + + + com.google.guava + guava + + + com.google.inject + guice + + + javax.validation + validation-api + + + io.airlift + concurrent + runtime + + + + io.airlift + log-manager + runtime + + + + io.airlift + units + runtime + + + + + io.trino + trino-spi + provided + + + + io.airlift + slice + provided + + + + com.fasterxml.jackson.core + jackson-annotations + provided + + + + org.openjdk.jol + jol-core + provided + + + + + io.trino + trino-base-jdbc + test-jar + test + + + + io.trino + trino-jmx + test + + + io.trino + trino-main + test + + + io.trino + trino-main + test-jar + test + + + io.trino + trino-testing + test + + + + io.trino + trino-tpch + test + + + + io.trino.tpch + tpch + test + + + io.airlift + testing + test + + + org.testng + testng + test + + + + diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java new file mode 100644 index 000000000000..bbf640613f68 --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java @@ -0,0 +1,274 @@ +package io.trino.plugin.spanner; + +import com.google.common.base.Preconditions; +import io.trino.plugin.jdbc.BaseJdbcClient; +import io.trino.plugin.jdbc.BaseJdbcConfig; +import io.trino.plugin.jdbc.ColumnMapping; +import io.trino.plugin.jdbc.ConnectionFactory; +import io.trino.plugin.jdbc.JdbcStatisticsConfig; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.plugin.jdbc.JdbcTypeHandle; +import io.trino.plugin.jdbc.QueryBuilder; +import io.trino.plugin.jdbc.RemoteTableName; +import io.trino.plugin.jdbc.StandardColumnMappings; +import io.trino.plugin.jdbc.WriteMapping; +import io.trino.plugin.jdbc.logging.RemoteQueryModifier; +import io.trino.plugin.jdbc.mapping.IdentifierMapping; +import io.trino.spi.ErrorCode; +import io.trino.spi.ErrorCodeSupplier; +import io.trino.spi.ErrorType; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.type.Type; +import io.trino.spi.type.TypeManager; +import io.trino.spi.type.VarcharType; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static io.trino.plugin.jdbc.StandardColumnMappings.bigintWriteFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.varcharWriteFunction; +import static io.trino.spi.ErrorType.INTERNAL_ERROR; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.SmallintType.SMALLINT; +import static io.trino.spi.type.TinyintType.TINYINT; +import static java.lang.String.format; +import static java.lang.String.join; + +public class SpannerClient + extends BaseJdbcClient +{ + private final String DEFAULT_SCHEMA = "public"; + private final SpannerConfig config; + + public SpannerClient(BaseJdbcConfig config, SpannerConfig spannerConfig, JdbcStatisticsConfig statisticsConfig, ConnectionFactory connectionFactory, QueryBuilder queryBuilder, TypeManager typeManager, IdentifierMapping identifierMapping, RemoteQueryModifier queryModifier) + { + super("`", connectionFactory, queryBuilder, config.getJdbcTypesMappedToVarchar(), identifierMapping, queryModifier, true); + this.config = spannerConfig; + } + + @Override + public Optional toColumnMapping(ConnectorSession session, Connection connection, JdbcTypeHandle typeHandle) + { + return Optional.of(StandardColumnMappings.varcharColumnMapping(VarcharType.VARCHAR, true)); + } + + @Override + public WriteMapping toWriteMapping(ConnectorSession session, Type type) + { + if (type == TINYINT || type == SMALLINT || type == INTEGER || type == BIGINT) { + return WriteMapping.longMapping("INT64", bigintWriteFunction()); + } + return WriteMapping.sliceMapping("STRING(MAX)", varcharWriteFunction()); + } + + @Override + protected Optional> getTableTypes() + { + return Optional.of(Arrays.asList("BASE TABLE", "VIEW")); + } + + @Override + public Collection listSchemas(Connection connection) + { + Set schemas = new HashSet<>(Collections.singleton(DEFAULT_SCHEMA)); + + try { + ResultSet resultSet = connection.getMetaData().getSchemas(null, null); + while (resultSet.next()) { + schemas.add(resultSet.getString(1)); + } + return schemas; + } + catch (SQLException e) { + return schemas; + } + } + + @Override + public boolean schemaExists(ConnectorSession session, String schema) + { + System.out.println("Called schemaExists " + schema); + if (schema.equalsIgnoreCase(DEFAULT_SCHEMA)) { + return true; + } + else { + try { + Connection connection = connectionFactory.openConnection(session); + ResultSet schemas = connection.getMetaData().getSchemas(null, null); + boolean found = false; + while (schemas.next()) { + if (schemas.getString(1).equalsIgnoreCase(schema)) { + found = true; + break; + } + } + return found; + } + catch (SQLException e) { + throw new RuntimeException(e); + } + } + } + + @Override + public void createSchema(ConnectorSession session, String schemaName) + { + throw new TrinoException(SpannerErrorCode.SPANNER_ERROR_CODE, "Spanner connector does not support creating schemas"); + } + + @Override + public void dropSchema(ConnectorSession session, String schemaName) + { + throw new TrinoException(SpannerErrorCode.SPANNER_ERROR_CODE, "Spanner connector does not support dropping schemas"); + } + + @Override + public List getTableNames(ConnectorSession session, Optional schema) + { + System.out.println("Called get table names " + schema); + List tables = new ArrayList<>(); + try { + Connection connection = connectionFactory.openConnection(session); + ResultSet resultSet = getTablesFromSpanner(connection); + while (resultSet.next()) { + tables.add(new SchemaTableName(DEFAULT_SCHEMA, resultSet.getString(1))); + } + } + catch (SQLException e) { + e.printStackTrace(); + } + return tables; + } + + private ResultSet getTablesFromSpanner(Connection connection) + throws SQLException + { + return connection.createStatement().executeQuery("SELECT\n" + + " TABLE_NAME\n" + + "FROM\n" + + " INFORMATION_SCHEMA.TABLES\n" + + "WHERE\n" + + " TABLE_CATALOG = '' and TABLE_SCHEMA = ''"); + } + + @Override + protected String createTableSql(RemoteTableName remoteTableName, List columns, ConnectorTableMetadata tableMetadata) + { + System.out.println("Called create table sql"); + Map properties = tableMetadata.getProperties(); + String primaryKey = SpannerTableProperties.getPrimaryKey(properties); + Preconditions.checkArgument(primaryKey != null, "Primary key is required to create a table in spanner"); + String interleaveTable = SpannerTableProperties.getInterleaveInParent(properties); + boolean onDeleteCascade = SpannerTableProperties.getOnDeleteCascade(properties); + String interleaveClause = ""; + String onDeleteClause = ""; + if (interleaveTable != null) { + interleaveClause = String.format(", INTERLEAVE IN PARENT %s", quoted(interleaveTable)); + onDeleteClause = onDeleteCascade ? " ON DELETE CASCADE" : " ON DELETE NO ACTION"; + } + System.out.println("primary key " + primaryKey); + String format = format("CREATE TABLE %s (%s) PRIMARY KEY (%s) %s %s", + quoted(remoteTableName.getTableName()), join(", ", columns), quoted(primaryKey), + interleaveClause, onDeleteClause); + System.out.println(format); + return format; + } + + public boolean checkTableExists(ConnectorSession session, String tableName) + throws SQLException + { + return checkTableExists(connectionFactory.openConnection(session), tableName); + } + + public boolean checkTableExists(Connection connection, String tableName) + throws SQLException + { + ResultSet tablesFromSpanner = getTablesFromSpanner(connection); + boolean exists = false; + while (tablesFromSpanner.next()) { + String table = tablesFromSpanner.getString(1); + if (table.equalsIgnoreCase(tableName)) { + exists = true; + break; + } + } + return exists; + } + + @Override + public Optional getTableHandle(ConnectorSession session, SchemaTableName schemaTableName) + { + boolean tableExists = false; + try { + tableExists = checkTableExists(session, schemaTableName.getTableName()); + if (tableExists) { + return Optional.of(new JdbcTableHandle(new SchemaTableName(DEFAULT_SCHEMA, schemaTableName.getTableName()), + new RemoteTableName(Optional.empty(), + Optional.empty(), schemaTableName.getTableName()), + Optional.empty())); + } + else { + return Optional.empty(); + } + } + catch (SQLException e) { + throw new TrinoException(SpannerErrorCode.SPANNER_ERROR_CODE, e); + } + } + + @Override + public void dropTable(ConnectorSession session, JdbcTableHandle handle) + { + System.out.println("Drop table "); + SchemaTableName schemaTableName = handle.getRequiredNamedRelation().getSchemaTableName(); + try (Connection connection = connectionFactory.openConnection(session)) { + String format = format("DROP TABLE %s", schemaTableName.getTableName()); + System.out.println("EEX " + format); + connection.createStatement().executeUpdate(format); + } + catch (SQLException e) { + e.printStackTrace(); + } + } + + @Override + public Map getTableProperties(ConnectorSession session, JdbcTableHandle tableHandle) + { + System.out.println("PROPS WAS CALLED "); + return new HashMap<>(); + } + + public enum SpannerErrorCode + implements ErrorCodeSupplier + { + SPANNER_ERROR_CODE(1, INTERNAL_ERROR); + + private final ErrorCode errorCode; + + SpannerErrorCode(int code, ErrorType type) + { + errorCode = new ErrorCode(code + 0x0506_0000, name(), type); + } + + @Override + public ErrorCode toErrorCode() + { + return errorCode; + } + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java new file mode 100644 index 000000000000..99133940a1b7 --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java @@ -0,0 +1,19 @@ +package io.trino.plugin.spanner; + +import io.airlift.configuration.Config; + +public class SpannerConfig +{ + private String credentialsFile; + + public String getCredentialsFile() + { + return credentialsFile; + } + + @Config("spanner.credentials.file") + public void setCredentialsFile(String credentialsFile) + { + this.credentialsFile = credentialsFile; + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java new file mode 100644 index 000000000000..7a254a600bef --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java @@ -0,0 +1,98 @@ +package io.trino.plugin.spanner; + +import com.google.cloud.spanner.jdbc.JdbcDriver; +import com.google.inject.Binder; +import com.google.inject.Provides; +import com.google.inject.Scopes; +import com.google.inject.Singleton; +import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.trino.plugin.jdbc.BaseJdbcConfig; +import io.trino.plugin.jdbc.ConfiguringConnectionFactory; +import io.trino.plugin.jdbc.ConnectionFactory; +import io.trino.plugin.jdbc.DecimalModule; +import io.trino.plugin.jdbc.DriverConnectionFactory; +import io.trino.plugin.jdbc.ForBaseJdbc; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcJoinPushdownSupportModule; +import io.trino.plugin.jdbc.JdbcStatisticsConfig; +import io.trino.plugin.jdbc.QueryBuilder; +import io.trino.plugin.jdbc.RemoteQueryCancellationModule; +import io.trino.plugin.jdbc.TablePropertiesProvider; +import io.trino.plugin.jdbc.credential.CredentialProvider; +import io.trino.plugin.jdbc.logging.RemoteQueryModifier; +import io.trino.plugin.jdbc.mapping.IdentifierMapping; +import io.trino.plugin.jdbc.ptf.Query; +import io.trino.spi.ptf.ConnectorTableFunction; +import io.trino.spi.type.TypeManager; + +import java.io.File; +import java.util.Properties; + +import static com.google.inject.multibindings.Multibinder.newSetBinder; +import static io.airlift.configuration.ConfigBinder.configBinder; + +public class SpannerModule + extends AbstractConfigurationAwareModule +{ + + @Provides + @Singleton + @ForBaseJdbc + public static ConnectionFactory createConnectionFactory(BaseJdbcConfig config, + CredentialProvider credentialProvider, + SpannerConfig spannerConfig) + throws ClassNotFoundException + { + Class.forName("com.google.cloud.spanner.jdbc.JdbcDriver"); + Properties connectionProperties = new Properties(); + String connectionUrl = config.getConnectionUrl(); + JdbcDriver driver = new JdbcDriver(); + File credentials = new File(spannerConfig.getCredentialsFile()); + if (!driver.acceptsURL(connectionUrl)) { + throw new RuntimeException(config.getConnectionUrl() + " is incorrect"); + } + connectionProperties.put("credentials", credentials.getAbsolutePath()); + return new ConfiguringConnectionFactory(new DriverConnectionFactory( + driver, + config.getConnectionUrl(), + connectionProperties, + credentialProvider), + connection -> { + }); + } + + @Provides + @Singleton + public static SpannerClient getSpannerClient( + BaseJdbcConfig config, + SpannerConfig spannerConfig, + JdbcStatisticsConfig statisticsConfig, + ConnectionFactory connectionFactory, + QueryBuilder queryBuilder, + TypeManager typeManager, + IdentifierMapping identifierMapping, + RemoteQueryModifier queryModifier) + { + return new SpannerClient(config, + spannerConfig, + statisticsConfig, + connectionFactory, + queryBuilder, + typeManager, + identifierMapping, + queryModifier); + } + + @Override + protected void setup(Binder binder) + { + binder.bind(JdbcClient.class).annotatedWith(ForBaseJdbc.class).to(SpannerClient.class).in(Scopes.SINGLETON); + configBinder(binder).bindConfig(SpannerConfig.class); + configBinder(binder).bindConfig(JdbcStatisticsConfig.class); + newSetBinder(binder, TablePropertiesProvider.class).addBinding().to(SpannerTableProperties.class); + install(new DecimalModule()); + install(new JdbcJoinPushdownSupportModule()); + install(new RemoteQueryCancellationModule()); + newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(Query.class).in(Scopes.SINGLETON); + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerPlugin.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerPlugin.java new file mode 100644 index 000000000000..eb65e24db19c --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerPlugin.java @@ -0,0 +1,28 @@ +/* + * 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.trino.plugin.spanner; + +import io.trino.plugin.jdbc.JdbcPlugin; +import io.trino.spi.Plugin; + +public class SpannerPlugin + extends JdbcPlugin + implements Plugin +{ + + public SpannerPlugin() + { + super("spanner", new SpannerModule()); + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java new file mode 100644 index 000000000000..7da60bfb2c2f --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java @@ -0,0 +1,65 @@ +package io.trino.plugin.spanner; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.jdbc.TablePropertiesProvider; +import io.trino.spi.session.PropertyMetadata; + +import javax.inject.Inject; + +import java.util.List; +import java.util.Map; + +import static io.trino.spi.session.PropertyMetadata.booleanProperty; +import static io.trino.spi.session.PropertyMetadata.stringProperty; +import static java.util.Objects.requireNonNull; + +public class SpannerTableProperties + implements TablePropertiesProvider +{ + public static final String PRIMARY_KEY = "primary_key"; + public static final String INTERLEAVE_IN_PARENT = "interleave_in_parent"; + public static final String ON_DELETE_CASCADE = "on_delete_cascade"; + private final ImmutableList> sessionProperties; + + @Inject + public SpannerTableProperties() + { + System.out.println("CALLED TABLE PROPERTIES "); + sessionProperties = ImmutableList.>builder() + .add(stringProperty( + PRIMARY_KEY, + "Primary key for the table being created", + null, + false)) + .add(stringProperty(INTERLEAVE_IN_PARENT, + "Table name which needs to be interleaved with this table", null, false)) + .add(booleanProperty(ON_DELETE_CASCADE, + "Boolean property to cascade on delete. ON DELETE CASCADE if set to true or ON DELETE NO ACTION if false", + false, false)) + .build(); + } + + public static String getPrimaryKey(Map tableProperties) + { + requireNonNull(tableProperties, "tableProperties is null"); + return (String) tableProperties.get(PRIMARY_KEY); + } + + public static String getInterleaveInParent(Map tableProperties) + { + requireNonNull(tableProperties, "tableProperties is null"); + return (String) tableProperties.get(INTERLEAVE_IN_PARENT); + } + + public static boolean getOnDeleteCascade(Map tableProperties) + { + requireNonNull(tableProperties, "tableProperties is null"); + return (boolean) tableProperties.get(ON_DELETE_CASCADE); + } + + @Override + public List> getTableProperties() + { + return sessionProperties; + } +} diff --git a/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java b/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java new file mode 100644 index 000000000000..73af4f490ecc --- /dev/null +++ b/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java @@ -0,0 +1,158 @@ +/* + * 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. + */ + +import com.google.common.collect.ImmutableMap; +import io.airlift.log.Logger; +import io.trino.Session; +import io.trino.plugin.base.security.FileBasedSystemAccessControl; +import io.trino.plugin.jmx.JmxPlugin; +import io.trino.plugin.spanner.SpannerPlugin; +import io.trino.plugin.tpch.TpchPlugin; +import io.trino.spi.security.SystemAccessControl; +import io.trino.testing.DistributedQueryRunner; +import io.trino.testing.MaterializedResult; +import io.trino.testing.QueryRunner; +import io.trino.tpch.TpchTable; + +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.function.Consumer; + +import static io.airlift.testing.Closeables.closeAllSuppress; +import static io.trino.testing.TestingSession.testSessionBuilder; +import static java.util.Objects.requireNonNull; + +public final class SpannerSqlQueryRunner +{ + private static final String TPCH_SCHEMA = "tpch"; + private static final String JDBC_URL = ""; + private static final String USER = ""; + private static final String PASSWORD = ""; + private static final String SCHEMA = ""; + private static Optional warehouse; + private static Optional catalog; + + private SpannerSqlQueryRunner() + { + } + + private static String getResourcePath(String resourceName) + { + return requireNonNull(SpannerSqlQueryRunner.class.getClassLoader().getResource(resourceName), "Resource does not exist: " + resourceName).getPath(); + } + + private static SystemAccessControl newFileBasedSystemAccessControl(String resourceName) + { + return newFileBasedSystemAccessControl(ImmutableMap.of("security.config-file", getResourcePath(resourceName))); + } + + private static SystemAccessControl newFileBasedSystemAccessControl(ImmutableMap config) + { + return new FileBasedSystemAccessControl.Factory().create(config); + } + + public static DistributedQueryRunner createSpannerSqlQueryRunner( + Map extraProperties, + Map coordinatorProperties, + Map connectorProperties, + Iterable> tables, + Consumer moreSetup) + throws Exception + { + DistributedQueryRunner queryRunner = null; + try { + DistributedQueryRunner.Builder builder = DistributedQueryRunner.builder(createSession()) + .setExtraProperties(extraProperties) + .setCoordinatorProperties(coordinatorProperties) + .setAdditionalSetup(moreSetup); + queryRunner = builder.build(); + + queryRunner.installPlugin(new TpchPlugin()); + queryRunner.createCatalog("tpch", "tpch"); + connectorProperties = new HashMap<>(ImmutableMap.copyOf(connectorProperties)); + connectorProperties.putIfAbsent("connection-url", JDBC_URL); + connectorProperties.putIfAbsent("connection-user", USER); + connectorProperties.putIfAbsent("connection-password", PASSWORD); + connectorProperties.putIfAbsent("spanner.credentials.file", "credentials.json"); + queryRunner.installPlugin(new SpannerPlugin()); + queryRunner.createCatalog("spanner", + "spanner", connectorProperties); + return queryRunner; + } + catch (Throwable e) { + closeAllSuppress(e, queryRunner); + throw e; + } + } + + public static Session createSession() + { + return testSessionBuilder() + .setCatalog("spanner") + .setSchema(TPCH_SCHEMA) + .build(); + } + + public static DistributedQueryRunner createSpannerSqlQueryRunner() + throws Exception + { + DistributedQueryRunner queryRunner = createSpannerSqlQueryRunner( + ImmutableMap.of("http-server.http.port", "8080" + //Property to skip columns from select * which user does not access to + //,"hide-inaccessible-columns", "true" + ), + ImmutableMap.of(), + ImmutableMap.of(), + TpchTable.getTables(), + r -> {} + ); + + queryRunner.installPlugin(new JmxPlugin()); + queryRunner.createCatalog("jmx", "jmx"); + return queryRunner; + } + + public static void main(String[] args) + throws Exception + { + DistributedQueryRunner queryRunner = createSpannerSqlQueryRunner( + ImmutableMap.of("http-server.http.port", "8080"), + ImmutableMap.of(), + ImmutableMap.of(), + TpchTable.getTables() + , xr -> {}); + + queryRunner.installPlugin(new JmxPlugin()); + queryRunner.createCatalog("jmx", "jmx"); + MaterializedResult schemas = queryRunner.execute("SHOW SCHEMAS FROM spanner"); + System.out.println(schemas); + MaterializedResult execute = queryRunner.execute("SHOW TABLES FROM spanner.public"); + System.out.println(execute); + MaterializedResult create = queryRunner.execute("create table spanner.public.dept" + + "( id int,name varchar" + + ")WITH (primary_key = 'id'," + + "interleave_in_parent='emp'," + + "on_delete_cascade=true)"); + System.out.println(create); + MaterializedResult drop = queryRunner.execute("DROP TABLE spanner.public.dept"); + System.out.println(drop); + + System.out.println("DONE"); + + Logger log = Logger.get(SpannerSqlQueryRunner.class); + log.info("======== SERVER STARTED ========"); + log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl()); + } +} diff --git a/pom.xml b/pom.xml index ce7c5f994e61..b2a321630d0b 100644 --- a/pom.xml +++ b/pom.xml @@ -196,6 +196,7 @@ testing/trino-testing-resources testing/trino-testing-services testing/trino-tests + plugin/trino-spanner From 5deb3f22f08c5d396d4936e9bf079eb33a1c0395 Mon Sep 17 00:00:00 2001 From: taherk77 Date: Mon, 27 Mar 2023 18:50:24 +0530 Subject: [PATCH 2/8] Added more datatypes parsing. Added create table features like primary keys,not null fields, timestamp field options etc. --- .../trino/plugin/spanner/SpannerClient.java | 422 +++++++++++++++++- .../trino/plugin/spanner/SpannerModule.java | 4 +- .../spanner/SpannerTableProperties.java | 61 ++- .../java/io/trino/plugin/spanner/Test.java | 44 ++ .../src/main/resources/docker-compose.yaml | 30 ++ .../src/test/java/SpannerSqlQueryRunner.java | 45 +- .../src/test/java/TestSpannerDataTypes.java | 42 ++ 7 files changed, 608 insertions(+), 40 deletions(-) create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/Test.java create mode 100644 plugin/trino-spanner/src/main/resources/docker-compose.yaml create mode 100644 plugin/trino-spanner/src/test/java/TestSpannerDataTypes.java diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java index bbf640613f68..01a3554aedb9 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java @@ -1,16 +1,26 @@ package io.trino.plugin.spanner; +import com.google.cloud.Timestamp; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; import io.trino.plugin.jdbc.BaseJdbcClient; import io.trino.plugin.jdbc.BaseJdbcConfig; +import io.trino.plugin.jdbc.CaseSensitivity; import io.trino.plugin.jdbc.ColumnMapping; import io.trino.plugin.jdbc.ConnectionFactory; +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; import io.trino.plugin.jdbc.JdbcStatisticsConfig; import io.trino.plugin.jdbc.JdbcTableHandle; import io.trino.plugin.jdbc.JdbcTypeHandle; +import io.trino.plugin.jdbc.LongReadFunction; +import io.trino.plugin.jdbc.LongWriteFunction; +import io.trino.plugin.jdbc.ObjectWriteFunction; import io.trino.plugin.jdbc.QueryBuilder; import io.trino.plugin.jdbc.RemoteTableName; import io.trino.plugin.jdbc.StandardColumnMappings; +import io.trino.plugin.jdbc.UnsupportedTypeHandling; +import io.trino.plugin.jdbc.WriteFunction; import io.trino.plugin.jdbc.WriteMapping; import io.trino.plugin.jdbc.logging.RemoteQueryModifier; import io.trino.plugin.jdbc.mapping.IdentifierMapping; @@ -18,62 +28,248 @@ import io.trino.spi.ErrorCodeSupplier; import io.trino.spi.ErrorType; import io.trino.spi.TrinoException; +import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorTableMetadata; import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.connector.TableNotFoundException; +import io.trino.spi.security.ConnectorIdentity; +import io.trino.spi.statistics.TableStatistics; +import io.trino.spi.type.BigintType; +import io.trino.spi.type.DateType; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.TimestampType; import io.trino.spi.type.Type; import io.trino.spi.type.TypeManager; import io.trino.spi.type.VarcharType; import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; +import java.sql.Types; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Locale; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.stream.Collectors; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Strings.emptyToNull; +import static com.google.common.base.Verify.verify; +import static io.trino.plugin.jdbc.JdbcErrorCode.JDBC_ERROR; +import static io.trino.plugin.jdbc.PredicatePushdownController.FULL_PUSHDOWN; import static io.trino.plugin.jdbc.StandardColumnMappings.bigintWriteFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.booleanWriteFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.dateReadFunctionUsingLocalDate; +import static io.trino.plugin.jdbc.StandardColumnMappings.defaultVarcharColumnMapping; +import static io.trino.plugin.jdbc.StandardColumnMappings.doubleColumnMapping; +import static io.trino.plugin.jdbc.StandardColumnMappings.doubleWriteFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.longDecimalWriteFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.timestampReadFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.varbinaryReadFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.varbinaryWriteFunction; import static io.trino.plugin.jdbc.StandardColumnMappings.varcharWriteFunction; +import static io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties.getUnsupportedTypeHandling; +import static io.trino.plugin.jdbc.UnsupportedTypeHandling.IGNORE; import static io.trino.spi.ErrorType.INTERNAL_ERROR; import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.BooleanType.BOOLEAN; +import static io.trino.spi.type.DateType.DATE; +import static io.trino.spi.type.DoubleType.DOUBLE; import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.RealType.REAL; import static io.trino.spi.type.SmallintType.SMALLINT; import static io.trino.spi.type.TinyintType.TINYINT; +import static io.trino.spi.type.VarbinaryType.VARBINARY; import static java.lang.String.format; import static java.lang.String.join; +import static java.sql.DatabaseMetaData.columnNoNulls; +import static java.time.format.DateTimeFormatter.ISO_DATE; +import static java.util.stream.Collectors.joining; public class SpannerClient extends BaseJdbcClient { - private final String DEFAULT_SCHEMA = "public"; + // Maps to Spanner's default empty schema + public static final String DEFAULT_SCHEMA = "default"; private final SpannerConfig config; + private final IdentifierMapping identifierMapping; public SpannerClient(BaseJdbcConfig config, SpannerConfig spannerConfig, JdbcStatisticsConfig statisticsConfig, ConnectionFactory connectionFactory, QueryBuilder queryBuilder, TypeManager typeManager, IdentifierMapping identifierMapping, RemoteQueryModifier queryModifier) { super("`", connectionFactory, queryBuilder, config.getJdbcTypesMappedToVarchar(), identifierMapping, queryModifier, true); this.config = spannerConfig; + this.identifierMapping = identifierMapping; + } + + private static RemoteTableName getRemoteTable(ResultSet resultSet) + throws SQLException + { + String schema = resultSet.getString("TABLE_SCHEM"); + if (schema != null && schema.equals("")) { + schema = null; + } + return new RemoteTableName( + Optional.ofNullable(null), + Optional.ofNullable(schema), + resultSet.getString("TABLE_NAME")); + } + + private static ColumnMetadata getPageSinkIdColumn(List otherColumnNames) + { + // While it's unlikely this column name will collide with client table columns, + // guarantee it will not by appending a deterministic suffix to it. + String baseColumnName = "trino_page_sink_id"; + String columnName = baseColumnName; + int suffix = 1; + while (otherColumnNames.contains(columnName)) { + columnName = baseColumnName + "_" + suffix; + suffix++; + } + return new ColumnMetadata(columnName, BigintType.BIGINT); } @Override public Optional toColumnMapping(ConnectorSession session, Connection connection, JdbcTypeHandle typeHandle) { - return Optional.of(StandardColumnMappings.varcharColumnMapping(VarcharType.VARCHAR, true)); + int jdbcType = typeHandle.getJdbcType(); + System.out.println("Column mapping for type " + typeHandle); + Optional mapping = getForcedMappingToVarchar(typeHandle); + if (mapping.isPresent()) { + return mapping; + } + switch (jdbcType) { + case Types.SMALLINT: + case Types.INTEGER: + case Types.TINYINT: + case Types.BIGINT: + return Optional.of(StandardColumnMappings.bigintColumnMapping()); + case Types.NUMERIC: + case Types.DECIMAL: + return Optional.of(StandardColumnMappings.decimalColumnMapping(DecimalType.createDecimalType(9, 38))); + case Types.REAL: + case Types.FLOAT: + case Types.DOUBLE: + return Optional.of(doubleColumnMapping()); + case Types.CHAR: + case Types.VARCHAR: + case Types.NVARCHAR: + case Types.LONGVARCHAR: + case Types.LONGNVARCHAR: + return Optional.of(defaultVarcharColumnMapping(typeHandle.getRequiredColumnSize(), false)); + case Types.BINARY: + case Types.VARBINARY: + case Types.LONGVARBINARY: + return Optional.of(ColumnMapping.sliceMapping(VARBINARY, varbinaryReadFunction(), varbinaryWriteFunction(), FULL_PUSHDOWN)); + case Types.DATE: + return Optional.of(ColumnMapping.longMapping( + DATE, + dateReadFunctionUsingLocalDate(), + spannerDateWriteFunctionUsingLocalDate())); + case Types.TIMESTAMP: + return Optional.of(ColumnMapping.longMapping( + TimestampType.TIMESTAMP_MILLIS, + (resultSet, columnIndex) -> { + java.sql.Timestamp timestamp = resultSet.getTimestamp(columnIndex); + return timestamp.toInstant().toEpochMilli() * 1000; + }, + (statement, index, value) -> statement.setTimestamp(index, new java.sql.Timestamp(value / 1000)))); + case Types.BOOLEAN: + return Optional.of(StandardColumnMappings.booleanColumnMapping()); + default: + throw new TrinoException(SpannerErrorCode.SPANNER_ERROR_CODE, "Spanner type mapper cannot build type mapping for JDBC type " + typeHandle.getJdbcType()); + } + } + + private LongWriteFunction spannerDateWriteFunctionUsingLocalDate() + { + return new LongWriteFunction() + { + @Override + public String getBindExpression() + { + return "CAST(? AS DATE)"; + } + + @Override + public void set(PreparedStatement statement, int index, long epochDay) + throws SQLException + { + statement.setString(index, LocalDate.ofEpochDay(epochDay).format(ISO_DATE)); + } + }; + } + + private LongWriteFunction spannerTimestampWriteFunction() + { + return new LongWriteFunction() + { + @Override + public void set(PreparedStatement statement, int index, long value) + throws SQLException + { + Timestamp timestamp = Timestamp.parseTimestamp(Instant.ofEpochMilli(value).toString()); + statement.setObject(index, timestamp); + } + + @Override + public String getBindExpression() + { + return "CAST(? AS TIMESTAMP)"; + } + }; } @Override public WriteMapping toWriteMapping(ConnectorSession session, Type type) { + //Spanner handles all types int and long types as INT64 if (type == TINYINT || type == SMALLINT || type == INTEGER || type == BIGINT) { return WriteMapping.longMapping("INT64", bigintWriteFunction()); } - return WriteMapping.sliceMapping("STRING(MAX)", varcharWriteFunction()); + if (type == BOOLEAN) { + return WriteMapping.booleanMapping("BOOL", booleanWriteFunction()); + } + if (type instanceof DecimalType) { + return WriteMapping.objectMapping("NUMERIC", longDecimalWriteFunction(DecimalType.createDecimalType(9, 38))); + } + if (type == REAL || type == DOUBLE) { + return WriteMapping.doubleMapping("FLOAT64", doubleWriteFunction()); + } + if (type instanceof TimestampType) { + return WriteMapping.objectMapping("TIMESTAMP", + tsWrite()); + } + if (type instanceof VarcharType) { + return WriteMapping.sliceMapping("STRING(MAX)", varcharWriteFunction()); + } + if (type instanceof DateType) { + return WriteMapping.sliceMapping("DATE", varcharWriteFunction()); + } + + throw new RuntimeException("Dont know type " + type); + } + + private ObjectWriteFunction tsWrite() + { + return ObjectWriteFunction.of( + String.class, + (statement, index, value) -> statement.setTimestamp(index, + com.google.cloud.Timestamp.parseTimestamp(value).toSqlTimestamp())); } @Override @@ -102,7 +298,6 @@ public Collection listSchemas(Connection connection) @Override public boolean schemaExists(ConnectorSession session, String schema) { - System.out.println("Called schemaExists " + schema); if (schema.equalsIgnoreCase(DEFAULT_SCHEMA)) { return true; } @@ -140,7 +335,6 @@ public void dropSchema(ConnectorSession session, String schemaName) @Override public List getTableNames(ConnectorSession session, Optional schema) { - System.out.println("Called get table names " + schema); List tables = new ArrayList<>(); try { Connection connection = connectionFactory.openConnection(session); @@ -167,26 +361,44 @@ private ResultSet getTablesFromSpanner(Connection connection) } @Override - protected String createTableSql(RemoteTableName remoteTableName, List columns, ConnectorTableMetadata tableMetadata) + public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata) { - System.out.println("Called create table sql"); + Map columnAndDataTypeMap = tableMetadata.getTableSchema().getColumns() + .stream() + .collect(Collectors.toMap(k -> k.getName().toUpperCase(Locale.ENGLISH), + v -> toWriteMapping(session, v.getType()).getDataType(), + (k, v) -> v, LinkedHashMap::new)); + Map properties = tableMetadata.getProperties(); - String primaryKey = SpannerTableProperties.getPrimaryKey(properties); - Preconditions.checkArgument(primaryKey != null, "Primary key is required to create a table in spanner"); + List primaryKeys = SpannerTableProperties.getPrimaryKey(properties); + List notNullFields = SpannerTableProperties.getNotNullFields(properties); + List commitTimestampFields = SpannerTableProperties.getCommitTimestampFields(properties); + Preconditions.checkArgument(primaryKeys != null && !primaryKeys.isEmpty(), "Primary key is required to create a table in spanner"); + Map columns = new LinkedHashMap<>(); + columnAndDataTypeMap.forEach((column, dataType) -> { + columns.put(column, join(" ", quoted(column), dataType)); + if (notNullFields.contains(column)) { + String columnWithConstraint = String.format("%s NOT NULL", columns.get(column)); + columns.put(column, columnWithConstraint); + } + if (commitTimestampFields.contains(column)) { + String columnWithConstraint = String.format("%s OPTIONS(allow_commit_timestamp=true)", columns.get(column)); + columns.put(column, columnWithConstraint); + } + }); String interleaveTable = SpannerTableProperties.getInterleaveInParent(properties); boolean onDeleteCascade = SpannerTableProperties.getOnDeleteCascade(properties); String interleaveClause = ""; String onDeleteClause = ""; if (interleaveTable != null) { - interleaveClause = String.format(", INTERLEAVE IN PARENT %s", quoted(interleaveTable)); - onDeleteClause = onDeleteCascade ? " ON DELETE CASCADE" : " ON DELETE NO ACTION"; + interleaveClause = String.format(", INTERLEAVE IN PARENT %s ", quoted(interleaveTable)); + onDeleteClause = onDeleteCascade ? " ON DELETE CASCADE " : " ON DELETE NO ACTION "; } - System.out.println("primary key " + primaryKey); - String format = format("CREATE TABLE %s (%s) PRIMARY KEY (%s) %s %s", - quoted(remoteTableName.getTableName()), join(", ", columns), quoted(primaryKey), + String sql = format("CREATE TABLE %s (%s) PRIMARY KEY (%s) %s %s", + quoted(tableMetadata.getTable().getTableName()), String.join(", ", columns.values()), + quoted(join(", ", primaryKeys)), interleaveClause, onDeleteClause); - System.out.println(format); - return format; + execute(session, sql); } public boolean checkTableExists(ConnectorSession session, String tableName) @@ -195,6 +407,175 @@ public boolean checkTableExists(ConnectorSession session, String tableName) return checkTableExists(connectionFactory.openConnection(session), tableName); } + @Override + public ResultSet getTables(Connection connection, Optional schemaName, Optional tableName) + throws SQLException + { + DatabaseMetaData metadata = connection.getMetaData(); + return metadata.getTables( + schemaName.orElse(null), + null, + escapeObjectNameForMetadataQuery(tableName, metadata.getSearchStringEscape()).orElse(null), + getTableTypes().map(types -> types.toArray(String[]::new)).orElse(null)); + } + + @Override + protected String getTableSchemaName(ResultSet resultSet) + throws SQLException + { + return null; + } + + @Override + protected ResultSet getColumns(JdbcTableHandle tableHandle, DatabaseMetaData metadata) + throws SQLException + { + RemoteTableName remoteTableName = tableHandle.getRequiredNamedRelation().getRemoteTableName(); + String schema = remoteTableName.getSchemaName().orElse(DEFAULT_SCHEMA) + .equalsIgnoreCase(DEFAULT_SCHEMA) ? null : escapeObjectNameForMetadataQuery(remoteTableName.getSchemaName(), metadata.getSearchStringEscape()).orElse(null); + return metadata.getColumns( + null, + schema, + escapeObjectNameForMetadataQuery(remoteTableName.getTableName(), metadata.getSearchStringEscape()), + null); + } + + @Override + public boolean supportsRetries() + { + return false; + } + + @Override + public String buildInsertSql(JdbcOutputTableHandle handle, List columnWriters) + { + boolean hasPageSinkIdColumn = handle.getPageSinkIdColumnName().isPresent(); + checkArgument(handle.getColumnNames().size() == columnWriters.size(), "handle and columnWriters mismatch: %s, %s", handle, columnWriters); + return format( + "INSERT INTO %s (%s%s) VALUES (%s%s)", + quoted(null, null, handle.getTemporaryTableName().orElseGet(handle::getTableName)), + handle.getColumnNames().stream() + .map(this::quoted) + .collect(joining(", ")), + hasPageSinkIdColumn ? ", " + quoted(handle.getPageSinkIdColumnName().get()) : "", + columnWriters.stream() + .map(WriteFunction::getBindExpression) + .collect(joining(",")), + hasPageSinkIdColumn ? ", ?" : ""); + } + + @Override + public JdbcOutputTableHandle beginInsertTable(ConnectorSession session, JdbcTableHandle tableHandle, List columns) + { + SchemaTableName schemaTableName = tableHandle.asPlainTable().getSchemaTableName(); + ConnectorIdentity identity = session.getIdentity(); + + verify(tableHandle.getAuthorization().isEmpty(), "Unexpected authorization is required for table: %s".formatted(tableHandle)); + try (Connection connection = connectionFactory.openConnection(session)) { + verify(connection.getAutoCommit()); + String remoteSchema = identifierMapping.toRemoteSchemaName(identity, connection, schemaTableName.getSchemaName()); + String remoteTable = identifierMapping.toRemoteTableName(identity, connection, remoteSchema, schemaTableName.getTableName()); + + ImmutableList.Builder columnNames = ImmutableList.builder(); + ImmutableList.Builder columnTypes = ImmutableList.builder(); + ImmutableList.Builder jdbcColumnTypes = ImmutableList.builder(); + for (JdbcColumnHandle column : columns) { + columnNames.add(column.getColumnName()); + columnTypes.add(column.getColumnType()); + jdbcColumnTypes.add(column.getJdbcTypeHandle()); + } + return new JdbcOutputTableHandle( + null, + DEFAULT_SCHEMA, + remoteTable, + columnNames.build(), + columnTypes.build(), + Optional.of(jdbcColumnTypes.build()), + Optional.empty(), + Optional.empty()); + } + catch (SQLException e) { + throw new TrinoException(JDBC_ERROR, e); + } + } + + private boolean shouldUseFaultTolerantExecution(ConnectorSession session) + { + //Does not support fault-tolerant exec + return false; + } + + @Override + public List getColumns(ConnectorSession session, JdbcTableHandle tableHandle) + { + if (tableHandle.getColumns().isPresent()) { + return tableHandle.getColumns().get(); + } + checkArgument(tableHandle.isNamedRelation(), "Cannot get columns for %s", tableHandle); + verify(tableHandle.getAuthorization().isEmpty(), "Unexpected authorization is required for table: %s".formatted(tableHandle)); + SchemaTableName schemaTableName = tableHandle.getRequiredNamedRelation().getSchemaTableName(); + RemoteTableName remoteTableName = tableHandle.getRequiredNamedRelation().getRemoteTableName(); + try (Connection connection = connectionFactory.openConnection(session); + ResultSet resultSet = getColumns(tableHandle, connection.getMetaData())) { + Map caseSensitivityMapping = getCaseSensitivityForColumns(session, connection, tableHandle); + int allColumns = 0; + List columns = new ArrayList<>(); + while (resultSet.next()) { + // skip if table doesn't match expected + RemoteTableName remoteTable = getRemoteTable(resultSet); + if (!(Objects.equals(remoteTableName, remoteTable))) { + continue; + } + allColumns++; + String columnName = resultSet.getString("COLUMN_NAME"); + JdbcTypeHandle typeHandle = new JdbcTypeHandle( + getInteger(resultSet, "DATA_TYPE").orElseThrow(() -> new IllegalStateException("DATA_TYPE is null")), + Optional.ofNullable(resultSet.getString("TYPE_NAME")), + getInteger(resultSet, "COLUMN_SIZE"), + getInteger(resultSet, "DECIMAL_DIGITS"), + Optional.empty(), + Optional.ofNullable(caseSensitivityMapping.get(columnName))); + Optional columnMapping = toColumnMapping(session, connection, typeHandle); + //log.debug("Mapping data type of '%s' column '%s': %s mapped to %s", schemaTableName, columnName, typeHandle, columnMapping); + boolean nullable = (resultSet.getInt("NULLABLE") != columnNoNulls); + // Note: some databases (e.g. SQL Server) do not return column remarks/comment here. + Optional comment = Optional.ofNullable(emptyToNull(resultSet.getString("REMARKS"))); + // skip unsupported column types + columnMapping.ifPresent(mapping -> columns.add(JdbcColumnHandle.builder() + .setColumnName(columnName) + .setJdbcTypeHandle(typeHandle) + .setColumnType(mapping.getType()) + .setNullable(nullable) + .setComment(comment) + .build())); + if (columnMapping.isEmpty()) { + UnsupportedTypeHandling unsupportedTypeHandling = getUnsupportedTypeHandling(session); + verify( + unsupportedTypeHandling == IGNORE, + "Unsupported type handling is set to %s, but toColumnMapping() returned empty for %s", + unsupportedTypeHandling, + typeHandle); + } + } + if (columns.isEmpty()) { + throw new TableNotFoundException( + schemaTableName, + format("Table '%s' has no supported columns (all %s columns are not supported)", schemaTableName, allColumns)); + } + return ImmutableList.copyOf(columns); + } + catch (SQLException e) { + throw new TrinoException(JDBC_ERROR, e); + } + } + + @Override + public TableStatistics getTableStatistics(ConnectorSession session, JdbcTableHandle handle) + { + System.out.println("Called table statistics"); + return super.getTableStatistics(session, handle); + } + public boolean checkTableExists(Connection connection, String tableName) throws SQLException { @@ -210,6 +591,12 @@ public boolean checkTableExists(Connection connection, String tableName) return exists; } + @Override + public void finishInsertTable(ConnectorSession session, JdbcOutputTableHandle handle, Set pageSinkIds) + { + //Nothing to do after insert + } + @Override public Optional getTableHandle(ConnectorSession session, SchemaTableName schemaTableName) { @@ -238,7 +625,6 @@ public void dropTable(ConnectorSession session, JdbcTableHandle handle) SchemaTableName schemaTableName = handle.getRequiredNamedRelation().getSchemaTableName(); try (Connection connection = connectionFactory.openConnection(session)) { String format = format("DROP TABLE %s", schemaTableName.getTableName()); - System.out.println("EEX " + format); connection.createStatement().executeUpdate(format); } catch (SQLException e) { @@ -249,7 +635,7 @@ public void dropTable(ConnectorSession session, JdbcTableHandle handle) @Override public Map getTableProperties(ConnectorSession session, JdbcTableHandle tableHandle) { - System.out.println("PROPS WAS CALLED "); + // System.out.println("PROPS WAS CALLED "); return new HashMap<>(); } diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java index 7a254a600bef..883e2d3fa7a7 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java @@ -47,11 +47,11 @@ public static ConnectionFactory createConnectionFactory(BaseJdbcConfig config, Properties connectionProperties = new Properties(); String connectionUrl = config.getConnectionUrl(); JdbcDriver driver = new JdbcDriver(); - File credentials = new File(spannerConfig.getCredentialsFile()); + //File credentials = new File(spannerConfig.getCredentialsFile()); if (!driver.acceptsURL(connectionUrl)) { throw new RuntimeException(config.getConnectionUrl() + " is incorrect"); } - connectionProperties.put("credentials", credentials.getAbsolutePath()); + //connectionProperties.put("credentials", spannerConfig.getCredentialsFile()); return new ConfiguringConnectionFactory(new DriverConnectionFactory( driver, config.getConnectionUrl(), diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java index 7da60bfb2c2f..888219e6adb1 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java @@ -3,20 +3,26 @@ import com.google.common.collect.ImmutableList; import io.trino.plugin.jdbc.TablePropertiesProvider; import io.trino.spi.session.PropertyMetadata; +import io.trino.spi.type.ArrayType; import javax.inject.Inject; import java.util.List; +import java.util.Locale; import java.util.Map; +import java.util.stream.Collectors; import static io.trino.spi.session.PropertyMetadata.booleanProperty; import static io.trino.spi.session.PropertyMetadata.stringProperty; +import static io.trino.spi.type.VarcharType.VARCHAR; import static java.util.Objects.requireNonNull; public class SpannerTableProperties implements TablePropertiesProvider { - public static final String PRIMARY_KEY = "primary_key"; + public static final String PRIMARY_KEYS = "primary_keys"; + public static final String NOT_NULL_FIELDS = "not_null_fields"; + public static final String COMMIT_TIMESTAMP_FIELDS = "commit_timestamp_fields"; public static final String INTERLEAVE_IN_PARENT = "interleave_in_parent"; public static final String ON_DELETE_CASCADE = "on_delete_cascade"; private final ImmutableList> sessionProperties; @@ -26,11 +32,33 @@ public SpannerTableProperties() { System.out.println("CALLED TABLE PROPERTIES "); sessionProperties = ImmutableList.>builder() - .add(stringProperty( - PRIMARY_KEY, - "Primary key for the table being created", - null, - false)) + .add(new PropertyMetadata<>( + PRIMARY_KEYS, + "Primary keys for the table being created", + new ArrayType(VARCHAR), + List.class, + ImmutableList.of(), + false, + value -> (List) value, + value -> value)) + .add(new PropertyMetadata<>( + NOT_NULL_FIELDS, + "Array of fields that should have NOT NULL constraints set on them in Spanner", + new ArrayType(VARCHAR), + List.class, + ImmutableList.of(), + false, + value -> (List) value, + value -> value)) + .add(new PropertyMetadata<>( + COMMIT_TIMESTAMP_FIELDS, + "Array of timestamp fields that should have 'OPTIONS (allow_commit_timestamp=true)' constraints set on them in Spanner", + new ArrayType(VARCHAR), + List.class, + ImmutableList.of(), + false, + value -> (List) value, + value -> value)) .add(stringProperty(INTERLEAVE_IN_PARENT, "Table name which needs to be interleaved with this table", null, false)) .add(booleanProperty(ON_DELETE_CASCADE, @@ -39,10 +67,10 @@ public SpannerTableProperties() .build(); } - public static String getPrimaryKey(Map tableProperties) + public static List getPrimaryKey(Map tableProperties) { requireNonNull(tableProperties, "tableProperties is null"); - return (String) tableProperties.get(PRIMARY_KEY); + return toUpperCase((List) tableProperties.get(PRIMARY_KEYS)); } public static String getInterleaveInParent(Map tableProperties) @@ -51,12 +79,29 @@ public static String getInterleaveInParent(Map tableProperties) return (String) tableProperties.get(INTERLEAVE_IN_PARENT); } + public static List getNotNullFields(Map tableProperties) + { + requireNonNull(tableProperties, "tableProperties is null"); + return toUpperCase((List) tableProperties.get(NOT_NULL_FIELDS)); + } + + public static List getCommitTimestampFields(Map tableProperties) + { + requireNonNull(tableProperties, "tableProperties is null"); + return toUpperCase((List) tableProperties.get(COMMIT_TIMESTAMP_FIELDS)); + } + public static boolean getOnDeleteCascade(Map tableProperties) { requireNonNull(tableProperties, "tableProperties is null"); return (boolean) tableProperties.get(ON_DELETE_CASCADE); } + private static List toUpperCase(List collection) + { + return collection.stream().map(f -> f.toUpperCase(Locale.ENGLISH)).collect(Collectors.toList()); + } + @Override public List> getTableProperties() { diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/Test.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/Test.java new file mode 100644 index 000000000000..6fc0af1ec465 --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/Test.java @@ -0,0 +1,44 @@ +package io.trino.plugin.spanner; + +import com.google.cloud.spanner.SpannerApiFutures; +import com.google.spanner.v1.SpannerGrpc; +import org.joda.time.DateTime; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.text.SimpleDateFormat; +import java.time.Instant; +import java.time.LocalDateTime; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import java.util.Calendar; +import java.util.Date; + +public class Test +{ + public static void main(String[] args) + throws SQLException + { + Connection connection = DriverManager.getConnection("jdbc:cloudspanner://0.0.0.0:9010/projects/spanner-project/instances/spanner-instance/databases/spanner-database;autoConfigEmulator=true"); + Statement statement = connection.createStatement(); + statement.execute("drop table t1"); + //1679963300421000 + //1591142320347 + //2020-06-02T23:58:40.347847393Z + Instant parse = Instant.parse("2020-06-02T23:58:40.347847393Z"); + //2023-03-28T00:42:31.756+00:00 + System.out.println(parse.toEpochMilli()); + System.out.println(Instant.ofEpochMilli(1679963300421L)); + LocalDateTime localDateTime = LocalDateTime.ofEpochSecond(1679963300421000L, 0, ZoneOffset.UTC); + System.out.println(localDateTime); + statement.execute("create table t1 (id int64,ts TIMESTAMP NOT NULL OPTIONS(allow_commit_timestamp=true))PRIMARY KEY (id)"); + statement.execute("select * from "); + /*PreparedStatement preparedStatement = connection.prepareStatement("insert into t1 values(?,?)"); + preparedStatement.setInt(1,1); + preparedStatement.setTimestamp(1,1);*/ + } +} diff --git a/plugin/trino-spanner/src/main/resources/docker-compose.yaml b/plugin/trino-spanner/src/main/resources/docker-compose.yaml new file mode 100644 index 000000000000..f9305954ee75 --- /dev/null +++ b/plugin/trino-spanner/src/main/resources/docker-compose.yaml @@ -0,0 +1,30 @@ +version: '3' +services: + spanner: + image: gcr.io/cloud-spanner-emulator/emulator:latest + ports: + - "9010:9010" + - "9020:9020" + + gcloud-spanner-init: + image: gcr.io/google.com/cloudsdktool/cloud-sdk:latest + environment: + PROJECT_ID: "spanner-project" + SPANNER_EMULATOR_URL: "http://localhost:9020/" + INSTANCE_NAME: "spanner-instance" + DATABASE_NAME: "spanner-database" + command: > + bash -c 'gcloud config configurations create emulator && + gcloud config set auth/disable_credentials true && + gcloud config set project $${PROJECT_ID} && + gcloud config set api_endpoint_overrides/spanner $${SPANNER_EMULATOR_URL} && + gcloud config set auth/disable_credentials true && + gcloud spanner instances create $${INSTANCE_NAME} --config=emulator-config --description=Emulator --nodes=1 + gcloud spanner databases create $${DATABASE_NAME} --instance=$${INSTANCE_NAME}' + spanner-cli: + image: sjdaws/spanner-cli:latest + environment: + SPANNER_EMULATOR_HOST: "spanner:9010" + depends_on: + - "gcloud-spanner-init" + command: ['sh', '-c', 'echo spanner client.... && tail -f /dev/null'] diff --git a/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java b/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java index 73af4f490ecc..a4538fd5cfff 100644 --- a/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java +++ b/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java @@ -24,6 +24,7 @@ import io.trino.testing.MaterializedResult; import io.trino.testing.QueryRunner; import io.trino.tpch.TpchTable; +import org.testng.annotations.Test; import java.util.HashMap; import java.util.Map; @@ -37,7 +38,8 @@ public final class SpannerSqlQueryRunner { private static final String TPCH_SCHEMA = "tpch"; - private static final String JDBC_URL = ""; + //url for emulated spanner host + private static final String JDBC_URL = "jdbc:cloudspanner://0.0.0.0:9010/projects/spanner-project/instances/spanner-instance/databases/spanner-database;autoConfigEmulator=true"; private static final String USER = ""; private static final String PASSWORD = ""; private static final String SCHEMA = ""; @@ -127,26 +129,31 @@ public static DistributedQueryRunner createSpannerSqlQueryRunner() public static void main(String[] args) throws Exception { - DistributedQueryRunner queryRunner = createSpannerSqlQueryRunner( - ImmutableMap.of("http-server.http.port", "8080"), - ImmutableMap.of(), - ImmutableMap.of(), - TpchTable.getTables() - , xr -> {}); + DistributedQueryRunner queryRunner = getQueryRunner(); queryRunner.installPlugin(new JmxPlugin()); queryRunner.createCatalog("jmx", "jmx"); MaterializedResult schemas = queryRunner.execute("SHOW SCHEMAS FROM spanner"); System.out.println(schemas); - MaterializedResult execute = queryRunner.execute("SHOW TABLES FROM spanner.public"); + MaterializedResult execute = queryRunner.execute("SHOW TABLES FROM spanner.default"); System.out.println(execute); - MaterializedResult create = queryRunner.execute("create table spanner.public.dept" + + MaterializedResult emp = queryRunner.execute("create table if not exists spanner.default.emp" + "( id int,name varchar" + - ")WITH (primary_key = 'id'," + + ") WITH(primary_keys = ARRAY['id'])"); + System.out.println(emp); + MaterializedResult create = queryRunner.execute("create table if not exists spanner.default.dept" + + "( id int,name varchar" + + ")WITH (primary_keys = ARRAY['id']," + "interleave_in_parent='emp'," + - "on_delete_cascade=true)"); + "on_delete_cascade=true" + + ")"); System.out.println(create); - MaterializedResult drop = queryRunner.execute("DROP TABLE spanner.public.dept"); + MaterializedResult insert = queryRunner.execute("insert into spanner.default.emp values(1,'Tom')"); + System.out.println(insert); + MaterializedResult count = queryRunner.execute("select count(*) from spanner.default.emp"); + System.out.println(count); + + MaterializedResult drop = queryRunner.execute("DROP TABLE spanner.default.dept"); System.out.println(drop); System.out.println("DONE"); @@ -155,4 +162,18 @@ public static void main(String[] args) log.info("======== SERVER STARTED ========"); log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl()); } + + static DistributedQueryRunner getQueryRunner() + throws Exception + { + DistributedQueryRunner queryRunner = createSpannerSqlQueryRunner( + ImmutableMap.of("http-server.http.port", "8080"), + ImmutableMap.of(), + ImmutableMap.of(), + TpchTable.getTables() + , xr -> {}); + return queryRunner; + } + + } diff --git a/plugin/trino-spanner/src/test/java/TestSpannerDataTypes.java b/plugin/trino-spanner/src/test/java/TestSpannerDataTypes.java new file mode 100644 index 000000000000..4ca21fb8dc5f --- /dev/null +++ b/plugin/trino-spanner/src/test/java/TestSpannerDataTypes.java @@ -0,0 +1,42 @@ +import io.trino.testing.DistributedQueryRunner; +import io.trino.testing.MaterializedResult; +import org.testng.annotations.Test; + +public class TestSpannerDataTypes +{ + @Test + public void testDataTypes() + throws Exception + { + DistributedQueryRunner queryRunner = SpannerSqlQueryRunner.getQueryRunner(); + queryRunner.execute("DROP TABLE IF EXISTS spanner.default.dTest"); + queryRunner.execute("create table spanner.default.dTest" + + "( id int,name varchar,is_active boolean)" + + "WITH (primary_keys = ARRAY['id']," + + "not_null_fields=ARRAY['name','is_active'])"); + queryRunner.execute("insert into spanner.default.dtest values(1,'Tom',true)"); + MaterializedResult execute = queryRunner.execute("select \"id\" from spanner.default.dtest"); + System.out.println(execute); + } + @Test + public void testTimestamp() + throws Exception + { + DistributedQueryRunner queryRunner = SpannerSqlQueryRunner.getQueryRunner(); + queryRunner.execute("DROP TABLE IF EXISTS spanner.default.dTest"); + queryRunner.execute("create table spanner.default.dTest" + + "( id int,name varchar,is_active boolean,ts timestamp," + + "LastContactDate DATE,PopularityScore DOUBLE)" + + "WITH (primary_keys = ARRAY['id']," + + "not_null_fields=ARRAY['name','is_active','ts'])"); + + queryRunner.execute("insert into spanner.default.dtest values(1,'Tom',true," + + "CURRENT_TIMESTAMP,CURRENT_DATE,1.111)"); + queryRunner.execute("insert into spanner.default.dtest values(2,'Tom cat',true," + + "NULL,CURRENT_DATE,1.111)"); + + MaterializedResult execute = queryRunner.execute("select * from spanner.default.dtest"); + System.out.println(execute); + + } +} From 8b1045e725b2515e7e46f86fb1b8f82556919b27 Mon Sep 17 00:00:00 2001 From: taherk77 Date: Mon, 27 Mar 2023 18:53:55 +0530 Subject: [PATCH 3/8] Removed useless file --- .../java/io/trino/plugin/spanner/Test.java | 44 ------------------- 1 file changed, 44 deletions(-) delete mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/Test.java diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/Test.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/Test.java deleted file mode 100644 index 6fc0af1ec465..000000000000 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/Test.java +++ /dev/null @@ -1,44 +0,0 @@ -package io.trino.plugin.spanner; - -import com.google.cloud.spanner.SpannerApiFutures; -import com.google.spanner.v1.SpannerGrpc; -import org.joda.time.DateTime; - -import java.sql.Connection; -import java.sql.DriverManager; -import java.sql.PreparedStatement; -import java.sql.ResultSet; -import java.sql.SQLException; -import java.sql.Statement; -import java.text.SimpleDateFormat; -import java.time.Instant; -import java.time.LocalDateTime; -import java.time.ZoneOffset; -import java.time.ZonedDateTime; -import java.util.Calendar; -import java.util.Date; - -public class Test -{ - public static void main(String[] args) - throws SQLException - { - Connection connection = DriverManager.getConnection("jdbc:cloudspanner://0.0.0.0:9010/projects/spanner-project/instances/spanner-instance/databases/spanner-database;autoConfigEmulator=true"); - Statement statement = connection.createStatement(); - statement.execute("drop table t1"); - //1679963300421000 - //1591142320347 - //2020-06-02T23:58:40.347847393Z - Instant parse = Instant.parse("2020-06-02T23:58:40.347847393Z"); - //2023-03-28T00:42:31.756+00:00 - System.out.println(parse.toEpochMilli()); - System.out.println(Instant.ofEpochMilli(1679963300421L)); - LocalDateTime localDateTime = LocalDateTime.ofEpochSecond(1679963300421000L, 0, ZoneOffset.UTC); - System.out.println(localDateTime); - statement.execute("create table t1 (id int64,ts TIMESTAMP NOT NULL OPTIONS(allow_commit_timestamp=true))PRIMARY KEY (id)"); - statement.execute("select * from "); - /*PreparedStatement preparedStatement = connection.prepareStatement("insert into t1 values(?,?)"); - preparedStatement.setInt(1,1); - preparedStatement.setTimestamp(1,1);*/ - } -} From 14cf842408ca49955519f7702f71da7200883391 Mon Sep 17 00:00:00 2001 From: taherk77 Date: Fri, 14 Apr 2023 16:20:33 +0530 Subject: [PATCH 4/8] Added spanner sink which uses spanner Mutations. Added spanner write modes as insert and upsert --- plugin/trino-spanner/pom.xml | 19 ++ .../trino/plugin/spanner/SpannerClient.java | 212 ++++++++++-------- .../trino/plugin/spanner/SpannerConfig.java | 17 ++ .../trino/plugin/spanner/SpannerModule.java | 177 ++++++++++++++- .../spanner/SpannerSessionProperties.java | 36 +++ .../io/trino/plugin/spanner/SpannerSink.java | 159 +++++++++++++ .../plugin/spanner/SpannerSinkProvider.java | 72 ++++++ .../spanner/SpannerTableProperties.java | 19 +- .../src/test/java/SpannerSqlQueryRunner.java | 179 --------------- .../src/test/java/TestSpannerDataTypes.java | 42 ---- .../plugin/spanner/SpannerQueryRunner.java | 155 +++++++++++++ .../plugin/spanner/TestSpannerPlugin.java | 22 ++ .../spanner/TestingSpannerInstance.java | 114 ++++++++++ 13 files changed, 893 insertions(+), 330 deletions(-) create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSessionProperties.java create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java delete mode 100644 plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java delete mode 100644 plugin/trino-spanner/src/test/java/TestSpannerDataTypes.java create mode 100644 plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java create mode 100644 plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerPlugin.java create mode 100644 plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestingSpannerInstance.java diff --git a/plugin/trino-spanner/pom.xml b/plugin/trino-spanner/pom.xml index af2be234f422..639a2e7cc660 100644 --- a/plugin/trino-spanner/pom.xml +++ b/plugin/trino-spanner/pom.xml @@ -122,6 +122,25 @@ trino-testing test + + io.trino + trino-testing-containers + test + + + + org.testcontainers + gcloud + 1.17.3 + test + + + + + io.trino + trino-testing-services + test + io.trino diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java index 01a3554aedb9..89d0045153f0 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java @@ -1,6 +1,18 @@ +/* + * 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.trino.plugin.spanner; -import com.google.cloud.Timestamp; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import io.trino.plugin.jdbc.BaseJdbcClient; @@ -13,11 +25,11 @@ import io.trino.plugin.jdbc.JdbcStatisticsConfig; import io.trino.plugin.jdbc.JdbcTableHandle; import io.trino.plugin.jdbc.JdbcTypeHandle; -import io.trino.plugin.jdbc.LongReadFunction; import io.trino.plugin.jdbc.LongWriteFunction; import io.trino.plugin.jdbc.ObjectWriteFunction; import io.trino.plugin.jdbc.QueryBuilder; import io.trino.plugin.jdbc.RemoteTableName; +import io.trino.plugin.jdbc.SliceWriteFunction; import io.trino.plugin.jdbc.StandardColumnMappings; import io.trino.plugin.jdbc.UnsupportedTypeHandling; import io.trino.plugin.jdbc.WriteFunction; @@ -35,7 +47,6 @@ import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.statistics.TableStatistics; -import io.trino.spi.type.BigintType; import io.trino.spi.type.DateType; import io.trino.spi.type.DecimalType; import io.trino.spi.type.TimestampType; @@ -49,9 +60,8 @@ import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Types; -import java.time.Instant; +import java.time.Duration; import java.time.LocalDate; -import java.time.LocalDateTime; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; @@ -79,7 +89,6 @@ import static io.trino.plugin.jdbc.StandardColumnMappings.doubleColumnMapping; import static io.trino.plugin.jdbc.StandardColumnMappings.doubleWriteFunction; import static io.trino.plugin.jdbc.StandardColumnMappings.longDecimalWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.timestampReadFunction; import static io.trino.plugin.jdbc.StandardColumnMappings.varbinaryReadFunction; import static io.trino.plugin.jdbc.StandardColumnMappings.varbinaryWriteFunction; import static io.trino.plugin.jdbc.StandardColumnMappings.varcharWriteFunction; @@ -108,6 +117,7 @@ public class SpannerClient public static final String DEFAULT_SCHEMA = "default"; private final SpannerConfig config; private final IdentifierMapping identifierMapping; + private final String tableTypes[] = {"BASE TABLE", "VIEW"}; public SpannerClient(BaseJdbcConfig config, SpannerConfig spannerConfig, JdbcStatisticsConfig statisticsConfig, ConnectionFactory connectionFactory, QueryBuilder queryBuilder, TypeManager typeManager, IdentifierMapping identifierMapping, RemoteQueryModifier queryModifier) { @@ -129,30 +139,31 @@ private static RemoteTableName getRemoteTable(ResultSet resultSet) resultSet.getString("TABLE_NAME")); } - private static ColumnMetadata getPageSinkIdColumn(List otherColumnNames) + public static void sleep() { - // While it's unlikely this column name will collide with client table columns, - // guarantee it will not by appending a deterministic suffix to it. - String baseColumnName = "trino_page_sink_id"; - String columnName = baseColumnName; - int suffix = 1; - while (otherColumnNames.contains(columnName)) { - columnName = baseColumnName + "_" + suffix; - suffix++; + try { + Thread.sleep(Duration.ofSeconds(1).toMillis()); + } + catch (InterruptedException e) { + e.printStackTrace(); } - return new ColumnMetadata(columnName, BigintType.BIGINT); } @Override public Optional toColumnMapping(ConnectorSession session, Connection connection, JdbcTypeHandle typeHandle) { int jdbcType = typeHandle.getJdbcType(); + String jdbcTypeName = typeHandle.getJdbcTypeName() + .orElseThrow(() -> new TrinoException(JDBC_ERROR, "Type name is missing: " + typeHandle)); System.out.println("Column mapping for type " + typeHandle); + System.out.println("JDBC TYPE NAME " + jdbcTypeName + " " + jdbcType); Optional mapping = getForcedMappingToVarchar(typeHandle); if (mapping.isPresent()) { return mapping; } switch (jdbcType) { + case Types.BOOLEAN: + return Optional.of(StandardColumnMappings.booleanColumnMapping()); case Types.SMALLINT: case Types.INTEGER: case Types.TINYINT: @@ -168,12 +179,8 @@ public Optional toColumnMapping(ConnectorSession session, Connect case Types.CHAR: case Types.VARCHAR: case Types.NVARCHAR: - case Types.LONGVARCHAR: - case Types.LONGNVARCHAR: return Optional.of(defaultVarcharColumnMapping(typeHandle.getRequiredColumnSize(), false)); case Types.BINARY: - case Types.VARBINARY: - case Types.LONGVARBINARY: return Optional.of(ColumnMapping.sliceMapping(VARBINARY, varbinaryReadFunction(), varbinaryWriteFunction(), FULL_PUSHDOWN)); case Types.DATE: return Optional.of(ColumnMapping.longMapping( @@ -188,8 +195,6 @@ public Optional toColumnMapping(ConnectorSession session, Connect return timestamp.toInstant().toEpochMilli() * 1000; }, (statement, index, value) -> statement.setTimestamp(index, new java.sql.Timestamp(value / 1000)))); - case Types.BOOLEAN: - return Optional.of(StandardColumnMappings.booleanColumnMapping()); default: throw new TrinoException(SpannerErrorCode.SPANNER_ERROR_CODE, "Spanner type mapper cannot build type mapping for JDBC type " + typeHandle.getJdbcType()); } @@ -214,24 +219,16 @@ public void set(PreparedStatement statement, int index, long epochDay) }; } - private LongWriteFunction spannerTimestampWriteFunction() + @Override + protected void execute(ConnectorSession session, String query) { - return new LongWriteFunction() - { - @Override - public void set(PreparedStatement statement, int index, long value) - throws SQLException - { - Timestamp timestamp = Timestamp.parseTimestamp(Instant.ofEpochMilli(value).toString()); - statement.setObject(index, timestamp); - } + } - @Override - public String getBindExpression() - { - return "CAST(? AS TIMESTAMP)"; - } - }; + @Override + protected void execute(ConnectorSession session, Connection connection, String query) + throws SQLException + { + super.execute(session, connection, query); } @Override @@ -250,21 +247,38 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type) if (type == REAL || type == DOUBLE) { return WriteMapping.doubleMapping("FLOAT64", doubleWriteFunction()); } + if (type == VARBINARY) { + return WriteMapping.sliceMapping("BYTES(MAX)", + SliceWriteFunction.of(Types.LONGVARBINARY, + (statement, index, value) -> statement.setBytes(index, value.byteArray()))); + } if (type instanceof TimestampType) { return WriteMapping.objectMapping("TIMESTAMP", - tsWrite()); + spannerTimestampWriteFunction()); } - if (type instanceof VarcharType) { - return WriteMapping.sliceMapping("STRING(MAX)", varcharWriteFunction()); + if (type instanceof VarcharType varcharType) { + String dataType = "STRING(MAX)"; + if (!varcharType.isUnbounded() && varcharType.getBoundedLength() <= 16777215) { + dataType = String.format("STRING(%s)", varcharType.getBoundedLength()); + } + return WriteMapping.sliceMapping(dataType, varcharWriteFunction()); } if (type instanceof DateType) { - return WriteMapping.sliceMapping("DATE", varcharWriteFunction()); + return WriteMapping.longMapping("DATE", + new LongWriteFunction() + { + @Override + public void set(PreparedStatement statement, int index, long value) + throws SQLException + { + statement.setDate(index, java.sql.Date.valueOf(LocalDate.ofEpochDay(value))); + } + }); } - - throw new RuntimeException("Dont know type " + type); + return WriteMapping.sliceMapping("STRING(MAX)", varcharWriteFunction()); } - private ObjectWriteFunction tsWrite() + private ObjectWriteFunction spannerTimestampWriteFunction() { return ObjectWriteFunction.of( String.class, @@ -275,16 +289,14 @@ private ObjectWriteFunction tsWrite() @Override protected Optional> getTableTypes() { - return Optional.of(Arrays.asList("BASE TABLE", "VIEW")); + return Optional.of(Arrays.asList(tableTypes)); } @Override public Collection listSchemas(Connection connection) { Set schemas = new HashSet<>(Collections.singleton(DEFAULT_SCHEMA)); - - try { - ResultSet resultSet = connection.getMetaData().getSchemas(null, null); + try (ResultSet resultSet = connection.getMetaData().getSchemas(null, null)) { while (resultSet.next()) { schemas.add(resultSet.getString(1)); } @@ -298,26 +310,8 @@ public Collection listSchemas(Connection connection) @Override public boolean schemaExists(ConnectorSession session, String schema) { - if (schema.equalsIgnoreCase(DEFAULT_SCHEMA)) { - return true; - } - else { - try { - Connection connection = connectionFactory.openConnection(session); - ResultSet schemas = connection.getMetaData().getSchemas(null, null); - boolean found = false; - while (schemas.next()) { - if (schemas.getString(1).equalsIgnoreCase(schema)) { - found = true; - break; - } - } - return found; - } - catch (SQLException e) { - throw new RuntimeException(e); - } - } + //There are no schemas in Spanner except for the default one that we have added and INFORMATION_SCHEMA + return schema.equalsIgnoreCase(DEFAULT_SCHEMA) || schema.equalsIgnoreCase("INFORMATION_SCHEMA"); } @Override @@ -336,11 +330,10 @@ public void dropSchema(ConnectorSession session, String schemaName) public List getTableNames(ConnectorSession session, Optional schema) { List tables = new ArrayList<>(); - try { - Connection connection = connectionFactory.openConnection(session); - ResultSet resultSet = getTablesFromSpanner(connection); + try (Connection connection = connectionFactory.openConnection(session); + ResultSet resultSet = getTables(connection, Optional.empty(), Optional.empty())) { while (resultSet.next()) { - tables.add(new SchemaTableName(DEFAULT_SCHEMA, resultSet.getString(1))); + tables.add(new SchemaTableName(DEFAULT_SCHEMA, resultSet.getString("TABLE_NAME"))); } } catch (SQLException e) { @@ -349,19 +342,27 @@ public List getTableNames(ConnectorSession session, Optional pageSinkIds) + { + //Do nothing + } + + @Override + public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata) + { + return super.beginCreateTable(session, tableMetadata); + } + + private String createSpannerTable(ConnectorSession session, ConnectorTableMetadata tableMetadata) { Map columnAndDataTypeMap = tableMetadata.getTableSchema().getColumns() .stream() @@ -373,7 +374,7 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe List primaryKeys = SpannerTableProperties.getPrimaryKey(properties); List notNullFields = SpannerTableProperties.getNotNullFields(properties); List commitTimestampFields = SpannerTableProperties.getCommitTimestampFields(properties); - Preconditions.checkArgument(primaryKeys != null && !primaryKeys.isEmpty(), "Primary key is required to create a table in spanner"); + Preconditions.checkArgument(!primaryKeys.isEmpty(), "Primary key is required to create a table in spanner"); Map columns = new LinkedHashMap<>(); columnAndDataTypeMap.forEach((column, dataType) -> { columns.put(column, join(" ", quoted(column), dataType)); @@ -398,7 +399,14 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe quoted(tableMetadata.getTable().getTableName()), String.join(", ", columns.values()), quoted(join(", ", primaryKeys)), interleaveClause, onDeleteClause); - execute(session, sql); + System.out.println(sql); + return sql; + } + + @Override + protected String createTableSql(RemoteTableName remoteTableName, List columns, ConnectorTableMetadata tableMetadata) + { + return createSpannerTable(null, tableMetadata); } public boolean checkTableExists(ConnectorSession session, String tableName) @@ -413,10 +421,10 @@ public ResultSet getTables(Connection connection, Optional schemaName, O { DatabaseMetaData metadata = connection.getMetaData(); return metadata.getTables( - schemaName.orElse(null), + null, null, escapeObjectNameForMetadataQuery(tableName, metadata.getSearchStringEscape()).orElse(null), - getTableTypes().map(types -> types.toArray(String[]::new)).orElse(null)); + null); } @Override @@ -446,14 +454,28 @@ public boolean supportsRetries() return false; } + @Override + protected JdbcOutputTableHandle createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, String targetTableName, Optional pageSinkIdColumn) + throws SQLException + { + return super.createTable(session, tableMetadata, targetTableName, pageSinkIdColumn); + } + + @Override + protected JdbcOutputTableHandle createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, String targetTableName) + throws SQLException + { + return super.createTable(session, tableMetadata, targetTableName); + } + @Override public String buildInsertSql(JdbcOutputTableHandle handle, List columnWriters) { boolean hasPageSinkIdColumn = handle.getPageSinkIdColumnName().isPresent(); checkArgument(handle.getColumnNames().size() == columnWriters.size(), "handle and columnWriters mismatch: %s, %s", handle, columnWriters); - return format( + String sql = format( "INSERT INTO %s (%s%s) VALUES (%s%s)", - quoted(null, null, handle.getTemporaryTableName().orElseGet(handle::getTableName)), + quoted(null, null, handle.getTableName()), handle.getColumnNames().stream() .map(this::quoted) .collect(joining(", ")), @@ -462,6 +484,8 @@ public String buildInsertSql(JdbcOutputTableHandle handle, List c .map(WriteFunction::getBindExpression) .collect(joining(",")), hasPageSinkIdColumn ? ", ?" : ""); + System.out.println("INSERT SQL " + sql); + return sql; } @Override @@ -499,12 +523,6 @@ public JdbcOutputTableHandle beginInsertTable(ConnectorSession session, JdbcTabl } } - private boolean shouldUseFaultTolerantExecution(ConnectorSession session) - { - //Does not support fault-tolerant exec - return false; - } - @Override public List getColumns(ConnectorSession session, JdbcTableHandle tableHandle) { @@ -579,10 +597,10 @@ public TableStatistics getTableStatistics(ConnectorSession session, JdbcTableHan public boolean checkTableExists(Connection connection, String tableName) throws SQLException { - ResultSet tablesFromSpanner = getTablesFromSpanner(connection); + ResultSet tablesFromSpanner = getTables(connection, Optional.empty(), Optional.empty()); boolean exists = false; while (tablesFromSpanner.next()) { - String table = tablesFromSpanner.getString(1); + String table = tablesFromSpanner.getString("TABLE_NAME"); if (table.equalsIgnoreCase(tableName)) { exists = true; break; @@ -635,7 +653,7 @@ public void dropTable(ConnectorSession session, JdbcTableHandle handle) @Override public Map getTableProperties(ConnectorSession session, JdbcTableHandle tableHandle) { - // System.out.println("PROPS WAS CALLED "); + // System.out.println("PROPS WAS CALLED "); return new HashMap<>(); } diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java index 99133940a1b7..3196fe52bf34 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java @@ -1,3 +1,16 @@ +/* + * 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.trino.plugin.spanner; import io.airlift.configuration.Config; @@ -5,6 +18,10 @@ public class SpannerConfig { private String credentialsFile; + /*private int minSessions = 100; + private int maxSessions = 400; + private int numChannels = 4; + private boolean retryAbortsInternally = true;*/ public String getCredentialsFile() { diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java index 883e2d3fa7a7..4ec12cf3c5af 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java @@ -1,35 +1,89 @@ +/* + * 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.trino.plugin.spanner; import com.google.cloud.spanner.jdbc.JdbcDriver; +import com.google.common.util.concurrent.MoreExecutors; import com.google.inject.Binder; +import com.google.inject.Key; import com.google.inject.Provides; import com.google.inject.Scopes; import com.google.inject.Singleton; +import com.google.inject.multibindings.Multibinder; import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.trino.plugin.base.session.SessionPropertiesProvider; import io.trino.plugin.jdbc.BaseJdbcConfig; +import io.trino.plugin.jdbc.CachingJdbcClient; import io.trino.plugin.jdbc.ConfiguringConnectionFactory; import io.trino.plugin.jdbc.ConnectionFactory; -import io.trino.plugin.jdbc.DecimalModule; +import io.trino.plugin.jdbc.DefaultJdbcMetadataFactory; +import io.trino.plugin.jdbc.DefaultQueryBuilder; import io.trino.plugin.jdbc.DriverConnectionFactory; +import io.trino.plugin.jdbc.DynamicFilteringStats; import io.trino.plugin.jdbc.ForBaseJdbc; +import io.trino.plugin.jdbc.ForJdbcDynamicFiltering; +import io.trino.plugin.jdbc.ForLazyConnectionFactory; +import io.trino.plugin.jdbc.ForRecordCursor; import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcConnector; +import io.trino.plugin.jdbc.JdbcDynamicFilteringConfig; +import io.trino.plugin.jdbc.JdbcDynamicFilteringSessionProperties; +import io.trino.plugin.jdbc.JdbcDynamicFilteringSplitManager; import io.trino.plugin.jdbc.JdbcJoinPushdownSupportModule; +import io.trino.plugin.jdbc.JdbcMetadataConfig; +import io.trino.plugin.jdbc.JdbcMetadataFactory; +import io.trino.plugin.jdbc.JdbcMetadataSessionProperties; +import io.trino.plugin.jdbc.JdbcRecordSetProvider; +import io.trino.plugin.jdbc.JdbcSplitManager; import io.trino.plugin.jdbc.JdbcStatisticsConfig; +import io.trino.plugin.jdbc.JdbcTransactionManager; +import io.trino.plugin.jdbc.JdbcWriteConfig; +import io.trino.plugin.jdbc.JdbcWriteSessionProperties; +import io.trino.plugin.jdbc.LazyConnectionFactory; +import io.trino.plugin.jdbc.MaxDomainCompactionThreshold; import io.trino.plugin.jdbc.QueryBuilder; -import io.trino.plugin.jdbc.RemoteQueryCancellationModule; +import io.trino.plugin.jdbc.QueryConfig; +import io.trino.plugin.jdbc.ReusableConnectionFactoryModule; +import io.trino.plugin.jdbc.StatsCollecting; import io.trino.plugin.jdbc.TablePropertiesProvider; +import io.trino.plugin.jdbc.TypeHandlingJdbcConfig; +import io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties; import io.trino.plugin.jdbc.credential.CredentialProvider; import io.trino.plugin.jdbc.logging.RemoteQueryModifier; +import io.trino.plugin.jdbc.logging.RemoteQueryModifierModule; import io.trino.plugin.jdbc.mapping.IdentifierMapping; -import io.trino.plugin.jdbc.ptf.Query; +import io.trino.plugin.jdbc.mapping.IdentifierMappingModule; +import io.trino.plugin.jdbc.procedure.FlushJdbcMetadataCacheProcedure; +import io.trino.spi.connector.ConnectorAccessControl; +import io.trino.spi.connector.ConnectorPageSinkProvider; +import io.trino.spi.connector.ConnectorRecordSetProvider; +import io.trino.spi.connector.ConnectorSplitManager; +import io.trino.spi.procedure.Procedure; import io.trino.spi.ptf.ConnectorTableFunction; import io.trino.spi.type.TypeManager; -import java.io.File; +import javax.annotation.PreDestroy; +import javax.inject.Provider; + import java.util.Properties; +import java.util.concurrent.ExecutorService; import static com.google.inject.multibindings.Multibinder.newSetBinder; +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; +import static io.airlift.configuration.ConditionalModule.conditionalModule; import static io.airlift.configuration.ConfigBinder.configBinder; +import static io.trino.plugin.jdbc.JdbcModule.bindSessionPropertiesProvider; public class SpannerModule extends AbstractConfigurationAwareModule @@ -52,6 +106,7 @@ public static ConnectionFactory createConnectionFactory(BaseJdbcConfig config, throw new RuntimeException(config.getConnectionUrl() + " is incorrect"); } //connectionProperties.put("credentials", spannerConfig.getCredentialsFile()); + connectionProperties.setProperty("retryAbortsInternally", "true"); return new ConfiguringConnectionFactory(new DriverConnectionFactory( driver, config.getConnectionUrl(), @@ -83,16 +138,120 @@ public static SpannerClient getSpannerClient( queryModifier); } + @Provides + public static SpannerTableProperties tableProperties() + { + return new SpannerTableProperties(); + } + @Singleton + public static SpannerSinkProvider configureSnowflakeSink( + ConnectionFactory factory, + JdbcClient client, + RemoteQueryModifier modifier, + SpannerClient snowflakeConfig, + SpannerTableProperties propertiesProvider) + { + return new SpannerSinkProvider(factory, client, modifier, snowflakeConfig, propertiesProvider); + } + + public static Multibinder sessionPropertiesProviderBinder(Binder binder) + { + return newSetBinder(binder, SessionPropertiesProvider.class); + } + + public static void bindSessionPropertiesProvider(Binder binder, Class type) + { + sessionPropertiesProviderBinder(binder).addBinding().to(type).in(Scopes.SINGLETON); + } + + public static Multibinder procedureBinder(Binder binder) + { + return newSetBinder(binder, Procedure.class); + } + + public static void bindProcedure(Binder binder, Class> type) + { + procedureBinder(binder).addBinding().toProvider(type).in(Scopes.SINGLETON); + } + + public static Multibinder tablePropertiesProviderBinder(Binder binder) + { + return newSetBinder(binder, TablePropertiesProvider.class); + } + + public static void bindTablePropertiesProvider(Binder binder, Class type) + { + tablePropertiesProviderBinder(binder).addBinding().to(type).in(Scopes.SINGLETON); + } + @Override - protected void setup(Binder binder) + public void setup(Binder binder) { binder.bind(JdbcClient.class).annotatedWith(ForBaseJdbc.class).to(SpannerClient.class).in(Scopes.SINGLETON); configBinder(binder).bindConfig(SpannerConfig.class); + configBinder(binder).bindConfig(TypeHandlingJdbcConfig.class); configBinder(binder).bindConfig(JdbcStatisticsConfig.class); - newSetBinder(binder, TablePropertiesProvider.class).addBinding().to(SpannerTableProperties.class); - install(new DecimalModule()); + install(new IdentifierMappingModule()); + install(new RemoteQueryModifierModule()); install(new JdbcJoinPushdownSupportModule()); - install(new RemoteQueryCancellationModule()); - newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(Query.class).in(Scopes.SINGLETON); + + newOptionalBinder(binder, ConnectorAccessControl.class); + newOptionalBinder(binder, QueryBuilder.class).setDefault().to(DefaultQueryBuilder.class).in(Scopes.SINGLETON); + + procedureBinder(binder); + tablePropertiesProviderBinder(binder); + + newOptionalBinder(binder, JdbcMetadataFactory.class).setBinding().to(DefaultJdbcMetadataFactory.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, Key.get(ConnectorSplitManager.class, ForJdbcDynamicFiltering.class)).setDefault().to(JdbcSplitManager.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, ConnectorSplitManager.class).setBinding().to(JdbcDynamicFilteringSplitManager.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, ConnectorRecordSetProvider.class).setBinding().to(JdbcRecordSetProvider.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, ConnectorPageSinkProvider.class).setBinding().to(SpannerSinkProvider.class).in(Scopes.SINGLETON); + + binder.bind(JdbcTransactionManager.class).in(Scopes.SINGLETON); + binder.bind(JdbcConnector.class).in(Scopes.SINGLETON); + + configBinder(binder).bindConfig(JdbcMetadataConfig.class); + configBinder(binder).bindConfig(JdbcWriteConfig.class); + configBinder(binder).bindConfig(BaseJdbcConfig.class); + configBinder(binder).bindConfig(SpannerConfig.class); + configBinder(binder).bindConfig(JdbcDynamicFilteringConfig.class); + bindTablePropertiesProvider(binder, SpannerTableProperties.class); + + configBinder(binder).bindConfig(TypeHandlingJdbcConfig.class); + bindSessionPropertiesProvider(binder, TypeHandlingJdbcSessionProperties.class); + bindSessionPropertiesProvider(binder, JdbcMetadataSessionProperties.class); + bindSessionPropertiesProvider(binder, JdbcWriteSessionProperties.class); + bindSessionPropertiesProvider(binder, JdbcDynamicFilteringSessionProperties.class); + bindSessionPropertiesProvider(binder, SpannerSessionProperties.class); + binder.bind(DynamicFilteringStats.class).in(Scopes.SINGLETON); + + binder.bind(CachingJdbcClient.class).in(Scopes.SINGLETON); + binder.bind(JdbcClient.class).to(Key.get(CachingJdbcClient.class)).in(Scopes.SINGLETON); + + newSetBinder(binder, Procedure.class).addBinding().toProvider(FlushJdbcMetadataCacheProcedure.class).in(Scopes.SINGLETON); + newSetBinder(binder, ConnectorTableFunction.class); + + binder.bind(ConnectionFactory.class) + .annotatedWith(ForLazyConnectionFactory.class) + .to(Key.get(ConnectionFactory.class, StatsCollecting.class)) + .in(Scopes.SINGLETON); + install(conditionalModule( + QueryConfig.class, + QueryConfig::isReuseConnection, + new ReusableConnectionFactoryModule(), + innerBinder -> innerBinder.bind(ConnectionFactory.class).to(LazyConnectionFactory.class).in(Scopes.SINGLETON))); + + newOptionalBinder(binder, Key.get(int.class, MaxDomainCompactionThreshold.class)); + + newOptionalBinder(binder, Key.get(ExecutorService.class, ForRecordCursor.class)) + .setBinding() + .toProvider(MoreExecutors::newDirectExecutorService) + .in(Scopes.SINGLETON); + } + + @PreDestroy + public void shutdownRecordCursorExecutor(@ForRecordCursor ExecutorService executor) + { + executor.shutdownNow(); } } diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSessionProperties.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSessionProperties.java new file mode 100644 index 000000000000..f78c09f5c51c --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSessionProperties.java @@ -0,0 +1,36 @@ +package io.trino.plugin.spanner; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.base.session.SessionPropertiesProvider; +import io.trino.spi.session.PropertyMetadata; + +import java.util.List; + +public class SpannerSessionProperties + implements SessionPropertiesProvider +{ + public static final String WRITE_MODE = "write_mode"; + private final ImmutableList> sessionProperties; + + public SpannerSessionProperties() + { + sessionProperties = ImmutableList.>builder() + .add(PropertyMetadata.enumProperty(WRITE_MODE, + "On write mode INSERT spanner throws an error on insert with duplicate primary keys," + + "on write mode UPSERT spanner updates the record with existing primary key with the new record being", + Mode.class, + Mode.INSERT, + false)).build(); + } + + @Override + public List> getSessionProperties() + { + return sessionProperties; + } + + enum Mode + { + INSERT, UPSERT + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java new file mode 100644 index 000000000000..580d14b89dc6 --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java @@ -0,0 +1,159 @@ +package io.trino.plugin.spanner; + +import com.google.cloud.NoCredentials; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.DatabaseId; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.Mutation; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.SpannerOptions; +import com.google.common.collect.ImmutableList; +import io.airlift.slice.Slice; +import io.airlift.slice.Slices; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.plugin.jdbc.logging.RemoteQueryModifier; +import io.trino.spi.Page; +import io.trino.spi.TrinoException; +import io.trino.spi.block.Block; +import io.trino.spi.connector.ConnectorPageSink; +import io.trino.spi.connector.ConnectorPageSinkId; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.type.Type; + +import java.time.LocalDate; +import java.util.Collection; +import java.util.LinkedList; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import static io.trino.plugin.jdbc.JdbcWriteSessionProperties.getWriteBatchSize; +import static java.time.format.DateTimeFormatter.ISO_DATE; +import static java.util.concurrent.CompletableFuture.completedFuture; + +public class SpannerSink + implements ConnectorPageSink +{ + private final SpannerOptions options; + private final int maxBatchSize; + private final DatabaseClient client; + private final List columnTypes; + private final List columnNames; + private final String project = "spanner-project"; + private final String instance = "spanner-instance"; + private final String database = "spanner-database"; + private final String table; + private final ConnectorPageSinkId pageSinkId; + private final SpannerSessionProperties.Mode writeMode; + private List mutations = new LinkedList<>(); + + public SpannerSink(ConnectorSession session, JdbcOutputTableHandle handle, SpannerClient spannerClient, ConnectorPageSinkId pageSinkId, RemoteQueryModifier remoteQueryModifier) + { + this.options = SpannerOptions + .newBuilder() + .setEmulatorHost("0.0.0.0:9010") + .setCredentials(NoCredentials.getInstance()) + .setProjectId(project) + .build(); + this.pageSinkId = pageSinkId; + this.maxBatchSize = getWriteBatchSize(session); + + this.client = options.getService().getDatabaseClient(DatabaseId.of(project, instance, database)); + columnTypes = handle.getColumnTypes(); + columnNames = handle.getColumnNames(); + table = handle.getTableName(); + writeMode = session.getProperty(SpannerSessionProperties.WRITE_MODE, SpannerSessionProperties.Mode.class); + } + + public Mutation.WriteBuilder createWriteBuilder() + { + if (writeMode.equals(SpannerSessionProperties.Mode.UPSERT)) { + return Mutation.newInsertOrUpdateBuilder(table); + } + else { + return Mutation.newInsertBuilder(table); + } + } + + @Override + public CompletableFuture appendPage(Page page) + { + for (int position = 0; position < page.getPositionCount(); position++) { + Mutation.WriteBuilder writeBuilder = createWriteBuilder(); + for (int channel = 0; channel < page.getChannelCount(); channel++) { + Block block = page.getBlock(channel); + Type type = columnTypes.get(channel); + String columnName = columnNames.get(channel); + if (!block.isNull(position)) { + Class javaType = type.getJavaType(); + if (javaType == boolean.class) { + writeBuilder.set(columnName).to(type.getBoolean(block, position)); + } + else if (javaType == long.class) { + if (type.getDisplayName().equalsIgnoreCase("DATE")) { + String date = LocalDate.ofEpochDay(type.getLong(block, position)).format(ISO_DATE); + writeBuilder.set(columnName).to(date); + } + else { + writeBuilder.set(columnName).to(type.getLong(block, position)); + } + } + else if (javaType == double.class) { + writeBuilder.set(columnName).to(type.getDouble(block, position)); + } + else if (javaType == Slice.class) { + writeBuilder.set(columnName).to(type.getSlice(block, position).toStringUtf8()); + } + else { + System.out.println("TYPE CLASS " + javaType); + System.out.println("TYPE Display NAME " + type.getDisplayName()); + System.out.println("TYPE Base NAME " + type.getBaseName()); + System.out.println("TYPE ID " + type.getTypeId()); + System.out.println("TYPE Signature " + type.getTypeSignature()); + throw new RuntimeException("Unknown type"); + } + } + mutations.add(writeBuilder.build()); + if (mutations.size() >= maxBatchSize) { + write(); + } + } + } + return NOT_BLOCKED; + } + + private void write() + { + if (!mutations.isEmpty()) { + try { + Timestamp write = client.write(mutations); + System.out.println("Batch write completed " + write + " " + mutations.size() + " records flushed"); + } + catch (Exception e) { + System.out.println(e); + if (e instanceof SpannerException spannerEx) { + if (spannerEx.getErrorCode().equals(ErrorCode.ALREADY_EXISTS)) { + throw new TrinoException(SpannerClient.SpannerErrorCode.SPANNER_ERROR_CODE, String.format("%s Try changing %s to %s and retry this query ", spannerEx.getMessage(), SpannerSessionProperties.WRITE_MODE, + SpannerSessionProperties.Mode.UPSERT)); + } + } + } + mutations = new LinkedList<>(); + } + } + + @Override + public CompletableFuture> finish() + { + write(); + return completedFuture(ImmutableList.of(Slices.wrappedLongArray(pageSinkId.getId()))); + } + + @Override + public void abort() + { + mutations = null; + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java new file mode 100644 index 000000000000..7196902f9a35 --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java @@ -0,0 +1,72 @@ +package io.trino.plugin.spanner; + +import com.google.inject.Inject; +import io.trino.plugin.jdbc.ConnectionFactory; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.plugin.jdbc.logging.RemoteQueryModifier; +import io.trino.spi.connector.ConnectorInsertTableHandle; +import io.trino.spi.connector.ConnectorMergeSink; +import io.trino.spi.connector.ConnectorMergeTableHandle; +import io.trino.spi.connector.ConnectorOutputTableHandle; +import io.trino.spi.connector.ConnectorPageSink; +import io.trino.spi.connector.ConnectorPageSinkId; +import io.trino.spi.connector.ConnectorPageSinkProvider; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTableExecuteHandle; +import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.session.PropertyMetadata; + +import java.util.List; + +public class SpannerSinkProvider + implements ConnectorPageSinkProvider +{ + private final SpannerClient client; + private final RemoteQueryModifier modifier; + private final SpannerTableProperties spannerTableProperties; + + @Inject + public SpannerSinkProvider( + ConnectionFactory connectionFactory, + JdbcClient jdbcClient, + RemoteQueryModifier modifier, + SpannerClient client, + SpannerTableProperties propertiesProvider) + { + System.out.println("Called Spanner Sink provider"); + this.client = client; + this.modifier = modifier; + System.out.println("TABLE PROPS in SINK " + propertiesProvider.getTableProperties()); + this.spannerTableProperties = propertiesProvider; + PropertyMetadata propertyMetadata = propertiesProvider.getTableProperties().get(0); + System.out.println(propertyMetadata.getDefaultValue()); + } + + @Override + public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorOutputTableHandle outputTableHandle, ConnectorPageSinkId pageSinkId) + { + return new SpannerSink(session, (JdbcOutputTableHandle) outputTableHandle, client, pageSinkId, modifier); + } + + @Override + public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorInsertTableHandle insertTableHandle, ConnectorPageSinkId pageSinkId) + { + System.out.println("SinkProvider 2"); + return null; + } + + @Override + public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle, ConnectorPageSinkId pageSinkId) + { + System.out.println("SinkProvider 3"); + return null; + } + + @Override + public ConnectorMergeSink createMergeSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorMergeTableHandle mergeHandle, ConnectorPageSinkId pageSinkId) + { + System.out.println("SinkProvider 4"); + return null; + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java index 888219e6adb1..94530811b78f 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerTableProperties.java @@ -1,3 +1,16 @@ +/* + * 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.trino.plugin.spanner; import com.google.common.collect.ImmutableList; @@ -70,7 +83,7 @@ public SpannerTableProperties() public static List getPrimaryKey(Map tableProperties) { requireNonNull(tableProperties, "tableProperties is null"); - return toUpperCase((List) tableProperties.get(PRIMARY_KEYS)); + return ImmutableList.copyOf(toUpperCase((List) tableProperties.get(PRIMARY_KEYS))); } public static String getInterleaveInParent(Map tableProperties) @@ -82,13 +95,13 @@ public static String getInterleaveInParent(Map tableProperties) public static List getNotNullFields(Map tableProperties) { requireNonNull(tableProperties, "tableProperties is null"); - return toUpperCase((List) tableProperties.get(NOT_NULL_FIELDS)); + return ImmutableList.copyOf(toUpperCase((List) tableProperties.get(NOT_NULL_FIELDS))); } public static List getCommitTimestampFields(Map tableProperties) { requireNonNull(tableProperties, "tableProperties is null"); - return toUpperCase((List) tableProperties.get(COMMIT_TIMESTAMP_FIELDS)); + return ImmutableList.copyOf(toUpperCase((List) tableProperties.get(COMMIT_TIMESTAMP_FIELDS))); } public static boolean getOnDeleteCascade(Map tableProperties) diff --git a/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java b/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java deleted file mode 100644 index a4538fd5cfff..000000000000 --- a/plugin/trino-spanner/src/test/java/SpannerSqlQueryRunner.java +++ /dev/null @@ -1,179 +0,0 @@ -/* - * 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. - */ - -import com.google.common.collect.ImmutableMap; -import io.airlift.log.Logger; -import io.trino.Session; -import io.trino.plugin.base.security.FileBasedSystemAccessControl; -import io.trino.plugin.jmx.JmxPlugin; -import io.trino.plugin.spanner.SpannerPlugin; -import io.trino.plugin.tpch.TpchPlugin; -import io.trino.spi.security.SystemAccessControl; -import io.trino.testing.DistributedQueryRunner; -import io.trino.testing.MaterializedResult; -import io.trino.testing.QueryRunner; -import io.trino.tpch.TpchTable; -import org.testng.annotations.Test; - -import java.util.HashMap; -import java.util.Map; -import java.util.Optional; -import java.util.function.Consumer; - -import static io.airlift.testing.Closeables.closeAllSuppress; -import static io.trino.testing.TestingSession.testSessionBuilder; -import static java.util.Objects.requireNonNull; - -public final class SpannerSqlQueryRunner -{ - private static final String TPCH_SCHEMA = "tpch"; - //url for emulated spanner host - private static final String JDBC_URL = "jdbc:cloudspanner://0.0.0.0:9010/projects/spanner-project/instances/spanner-instance/databases/spanner-database;autoConfigEmulator=true"; - private static final String USER = ""; - private static final String PASSWORD = ""; - private static final String SCHEMA = ""; - private static Optional warehouse; - private static Optional catalog; - - private SpannerSqlQueryRunner() - { - } - - private static String getResourcePath(String resourceName) - { - return requireNonNull(SpannerSqlQueryRunner.class.getClassLoader().getResource(resourceName), "Resource does not exist: " + resourceName).getPath(); - } - - private static SystemAccessControl newFileBasedSystemAccessControl(String resourceName) - { - return newFileBasedSystemAccessControl(ImmutableMap.of("security.config-file", getResourcePath(resourceName))); - } - - private static SystemAccessControl newFileBasedSystemAccessControl(ImmutableMap config) - { - return new FileBasedSystemAccessControl.Factory().create(config); - } - - public static DistributedQueryRunner createSpannerSqlQueryRunner( - Map extraProperties, - Map coordinatorProperties, - Map connectorProperties, - Iterable> tables, - Consumer moreSetup) - throws Exception - { - DistributedQueryRunner queryRunner = null; - try { - DistributedQueryRunner.Builder builder = DistributedQueryRunner.builder(createSession()) - .setExtraProperties(extraProperties) - .setCoordinatorProperties(coordinatorProperties) - .setAdditionalSetup(moreSetup); - queryRunner = builder.build(); - - queryRunner.installPlugin(new TpchPlugin()); - queryRunner.createCatalog("tpch", "tpch"); - connectorProperties = new HashMap<>(ImmutableMap.copyOf(connectorProperties)); - connectorProperties.putIfAbsent("connection-url", JDBC_URL); - connectorProperties.putIfAbsent("connection-user", USER); - connectorProperties.putIfAbsent("connection-password", PASSWORD); - connectorProperties.putIfAbsent("spanner.credentials.file", "credentials.json"); - queryRunner.installPlugin(new SpannerPlugin()); - queryRunner.createCatalog("spanner", - "spanner", connectorProperties); - return queryRunner; - } - catch (Throwable e) { - closeAllSuppress(e, queryRunner); - throw e; - } - } - - public static Session createSession() - { - return testSessionBuilder() - .setCatalog("spanner") - .setSchema(TPCH_SCHEMA) - .build(); - } - - public static DistributedQueryRunner createSpannerSqlQueryRunner() - throws Exception - { - DistributedQueryRunner queryRunner = createSpannerSqlQueryRunner( - ImmutableMap.of("http-server.http.port", "8080" - //Property to skip columns from select * which user does not access to - //,"hide-inaccessible-columns", "true" - ), - ImmutableMap.of(), - ImmutableMap.of(), - TpchTable.getTables(), - r -> {} - ); - - queryRunner.installPlugin(new JmxPlugin()); - queryRunner.createCatalog("jmx", "jmx"); - return queryRunner; - } - - public static void main(String[] args) - throws Exception - { - DistributedQueryRunner queryRunner = getQueryRunner(); - - queryRunner.installPlugin(new JmxPlugin()); - queryRunner.createCatalog("jmx", "jmx"); - MaterializedResult schemas = queryRunner.execute("SHOW SCHEMAS FROM spanner"); - System.out.println(schemas); - MaterializedResult execute = queryRunner.execute("SHOW TABLES FROM spanner.default"); - System.out.println(execute); - MaterializedResult emp = queryRunner.execute("create table if not exists spanner.default.emp" + - "( id int,name varchar" + - ") WITH(primary_keys = ARRAY['id'])"); - System.out.println(emp); - MaterializedResult create = queryRunner.execute("create table if not exists spanner.default.dept" + - "( id int,name varchar" + - ")WITH (primary_keys = ARRAY['id']," + - "interleave_in_parent='emp'," + - "on_delete_cascade=true" + - ")"); - System.out.println(create); - MaterializedResult insert = queryRunner.execute("insert into spanner.default.emp values(1,'Tom')"); - System.out.println(insert); - MaterializedResult count = queryRunner.execute("select count(*) from spanner.default.emp"); - System.out.println(count); - - MaterializedResult drop = queryRunner.execute("DROP TABLE spanner.default.dept"); - System.out.println(drop); - - System.out.println("DONE"); - - Logger log = Logger.get(SpannerSqlQueryRunner.class); - log.info("======== SERVER STARTED ========"); - log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl()); - } - - static DistributedQueryRunner getQueryRunner() - throws Exception - { - DistributedQueryRunner queryRunner = createSpannerSqlQueryRunner( - ImmutableMap.of("http-server.http.port", "8080"), - ImmutableMap.of(), - ImmutableMap.of(), - TpchTable.getTables() - , xr -> {}); - return queryRunner; - } - - -} diff --git a/plugin/trino-spanner/src/test/java/TestSpannerDataTypes.java b/plugin/trino-spanner/src/test/java/TestSpannerDataTypes.java deleted file mode 100644 index 4ca21fb8dc5f..000000000000 --- a/plugin/trino-spanner/src/test/java/TestSpannerDataTypes.java +++ /dev/null @@ -1,42 +0,0 @@ -import io.trino.testing.DistributedQueryRunner; -import io.trino.testing.MaterializedResult; -import org.testng.annotations.Test; - -public class TestSpannerDataTypes -{ - @Test - public void testDataTypes() - throws Exception - { - DistributedQueryRunner queryRunner = SpannerSqlQueryRunner.getQueryRunner(); - queryRunner.execute("DROP TABLE IF EXISTS spanner.default.dTest"); - queryRunner.execute("create table spanner.default.dTest" + - "( id int,name varchar,is_active boolean)" + - "WITH (primary_keys = ARRAY['id']," + - "not_null_fields=ARRAY['name','is_active'])"); - queryRunner.execute("insert into spanner.default.dtest values(1,'Tom',true)"); - MaterializedResult execute = queryRunner.execute("select \"id\" from spanner.default.dtest"); - System.out.println(execute); - } - @Test - public void testTimestamp() - throws Exception - { - DistributedQueryRunner queryRunner = SpannerSqlQueryRunner.getQueryRunner(); - queryRunner.execute("DROP TABLE IF EXISTS spanner.default.dTest"); - queryRunner.execute("create table spanner.default.dTest" + - "( id int,name varchar,is_active boolean,ts timestamp," + - "LastContactDate DATE,PopularityScore DOUBLE)" + - "WITH (primary_keys = ARRAY['id']," + - "not_null_fields=ARRAY['name','is_active','ts'])"); - - queryRunner.execute("insert into spanner.default.dtest values(1,'Tom',true," + - "CURRENT_TIMESTAMP,CURRENT_DATE,1.111)"); - queryRunner.execute("insert into spanner.default.dtest values(2,'Tom cat',true," + - "NULL,CURRENT_DATE,1.111)"); - - MaterializedResult execute = queryRunner.execute("select * from spanner.default.dtest"); - System.out.println(execute); - - } -} diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java new file mode 100644 index 000000000000..74983d37f87c --- /dev/null +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java @@ -0,0 +1,155 @@ +package io.trino.plugin.spanner;/* + * 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. + */ + +import com.google.common.collect.ImmutableMap; +import io.airlift.log.Logger; +import io.trino.Session; +import io.trino.metadata.QualifiedObjectName; +import io.trino.plugin.jmx.JmxPlugin; +import io.trino.plugin.tpch.TpchPlugin; +import io.trino.testing.DistributedQueryRunner; +import io.trino.testing.MaterializedResult; +import io.trino.testing.QueryRunner; +import io.trino.tpch.TpchTable; +import org.intellij.lang.annotations.Language; + +import java.util.HashMap; +import java.util.Map; +import java.util.function.Consumer; + +import static io.airlift.testing.Closeables.closeAllSuppress; +import static io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME; +import static io.trino.testing.TestingSession.testSessionBuilder; +import static java.lang.String.format; + +public final class SpannerQueryRunner +{ + private static final Logger LOG = Logger.get(SpannerQueryRunner.class); + private static final String TPCH_SCHEMA = "tpch"; + + private SpannerQueryRunner() {} + + public static DistributedQueryRunner createSpannerQueryRunner( + TestingSpannerInstance instance, + Map extraProperties, + Map connectorProperties, + Iterable> tables) + throws Exception + { + return createSpannerQueryRunner(instance, extraProperties, Map.of(), connectorProperties, tables, runner -> {}); + } + + public static DistributedQueryRunner createSpannerQueryRunner( + TestingSpannerInstance instance, + Map extraProperties, + Map coordinatorProperties, + Map connectorProperties, + Iterable> tables, + Consumer moreSetup) + throws Exception + { + DistributedQueryRunner queryRunner = null; + try { + queryRunner = DistributedQueryRunner.builder(createSession()) + .setExtraProperties(extraProperties) + .setCoordinatorProperties(coordinatorProperties) + .setAdditionalSetup(moreSetup) + .build(); + + queryRunner.installPlugin(new TpchPlugin()); + queryRunner.createCatalog("tpch", "tpch"); + + // note: additional copy via ImmutableList so that if fails on nulls + connectorProperties = new HashMap<>(ImmutableMap.copyOf(connectorProperties)); + //connectorProperties.putIfAbsent("connection-url", instance.getJdbcUrl()); + connectorProperties.putIfAbsent("connection-url", "jdbc:cloudspanner://0.0.0.0:9010/projects/spanner-project/instances/spanner-instance/databases/spanner-database;autoConfigEmulator=true"); + connectorProperties.putIfAbsent("connection-user", ""); + connectorProperties.putIfAbsent("connection-password", ""); + connectorProperties.putIfAbsent("spanner.credentials.file", "credentials.json"); + queryRunner.installPlugin(new SpannerPlugin()); + queryRunner.createCatalog("spanner", "spanner", connectorProperties); + copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME,createSession(), tables); + MaterializedResult execute = queryRunner.execute("SHOW TABLES FROM spanner.default"); + System.out.println(execute); +/* + MaterializedResult rows = queryRunner.execute("SELECT * FROM spanner.default.customer"); + System.out.println(rows); +*/ + return queryRunner; + } + catch (Throwable e) { + closeAllSuppress(e, queryRunner, instance); + throw e; + } + } + + public static Session createSession() + { + return testSessionBuilder() + .setCatalog("spanner") + .setSchema("default") + .setCatalogSessionProperty("spanner",SpannerSessionProperties.WRITE_MODE, SpannerSessionProperties.Mode.UPSERT.name()) + .build(); + } + + public static void main(String[] args) + throws Exception + { + DistributedQueryRunner queryRunner = createSpannerQueryRunner( + new TestingSpannerInstance(), + ImmutableMap.of("http-server.http.port", "8080"), + ImmutableMap.of(), + TpchTable.getTables()); + + queryRunner.installPlugin(new JmxPlugin()); + queryRunner.createCatalog("jmx", "jmx"); + + Logger log = Logger.get(SpannerQueryRunner.class); + log.info("======== SERVER STARTED ========"); + log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl()); + } + + private static void copyTpchTables( + QueryRunner queryRunner, + String sourceCatalog, + String sourceSchema, + Session session, + Iterable> tables) + { + LOG.debug("Loading data from %s.%s...", sourceCatalog, sourceSchema); + for (TpchTable table : tables) { + copyTable(queryRunner, sourceCatalog, session, sourceSchema, table); + } + } + + private static void copyTable( + QueryRunner queryRunner, + String catalog, + Session session, + String schema, + TpchTable table) + { + QualifiedObjectName source = new QualifiedObjectName(catalog, schema, table.getTableName()); + String target = table.getTableName(); + String primaryKey = table.getColumns().get(0).getSimplifiedColumnName(); + String tableProperties = String.format("WITH (PRIMARY_KEYS = ARRAY['%s'])", primaryKey); + @Language("SQL") + String sql = format("CREATE TABLE IF NOT EXISTS %s %s AS SELECT * FROM %s", + target, tableProperties, source, primaryKey, primaryKey, source); + System.out.println(sql); + LOG.debug("Running import for %s %s", target, sql); + long rows = queryRunner.execute(session, sql).getUpdateCount().getAsLong(); + LOG.debug("%s rows loaded into %s", rows, target); + } +} diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerPlugin.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerPlugin.java new file mode 100644 index 000000000000..a3cdd1f7ceaf --- /dev/null +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerPlugin.java @@ -0,0 +1,22 @@ +package io.trino.plugin.spanner; + +import com.google.common.collect.ImmutableMap; +import io.trino.plugin.spanner.SpannerPlugin; +import io.trino.spi.Plugin; +import io.trino.spi.connector.ConnectorFactory; +import io.trino.testing.TestingConnectorContext; +import org.testcontainers.utility.DockerImageName; +import org.testng.annotations.Test; + +import static com.google.common.collect.Iterables.getOnlyElement; + +public class TestSpannerPlugin +{ + @Test + public void testCreateConnector() + { + Plugin plugin = new SpannerPlugin(); + ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories()); + factory.create("test", ImmutableMap.of("connection-url", "jdbc:cloudspanner://0.0.0.0:9010/projects/spanner-project/instances/spanner-instance/databases/spanner-database;autoConfigEmulator=true"), new TestingConnectorContext()).shutdown(); + } +} diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestingSpannerInstance.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestingSpannerInstance.java new file mode 100644 index 000000000000..323ccfe4bf17 --- /dev/null +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestingSpannerInstance.java @@ -0,0 +1,114 @@ +package io.trino.plugin.spanner; + +import com.google.cloud.NoCredentials; +import com.google.cloud.spanner.Database; +import com.google.cloud.spanner.DatabaseAdminClient; +import com.google.cloud.spanner.DatabaseId; +import com.google.cloud.spanner.InstanceAdminClient; +import com.google.cloud.spanner.InstanceConfigId; +import com.google.cloud.spanner.InstanceId; +import com.google.cloud.spanner.InstanceInfo; +import com.google.cloud.spanner.Spanner; +import com.google.cloud.spanner.SpannerOptions; +import org.intellij.lang.annotations.Language; +import org.testcontainers.containers.SpannerEmulatorContainer; +import org.testcontainers.utility.DockerImageName; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Properties; +import java.util.concurrent.ExecutionException; + +public class TestingSpannerInstance + implements AutoCloseable +{ + private final String SPANNER_IMAGE = "gcr.io/cloud-spanner-emulator/emulator:latest"; + + private final SpannerEmulatorContainer emulatorContainer; + private final SpannerOptions options=null; + private final String PROJECT = "test-project"; + private final String INSTANCE = "test-instance"; + private final String DATABASE = "trinodb"; + private final Spanner spanner=null; + private final DatabaseId databaseId=null; + private final InstanceId instanceId=null; + + public TestingSpannerInstance() + throws ExecutionException, InterruptedException + { + this.emulatorContainer = new SpannerEmulatorContainer(DockerImageName.parse(SPANNER_IMAGE)); + emulatorContainer.setExposedPorts(Arrays.asList(9010)); + /*emulatorContainer.start(); + options = SpannerOptions + .newBuilder() + .setEmulatorHost(emulatorContainer.getEmulatorGrpcEndpoint()) + .setCredentials(NoCredentials.getInstance()) + .setProjectId(PROJECT) + .build(); + this.spanner = options.getService(); + this.instanceId = createInstance(); + Database database = createDatabase(); + this.databaseId = DatabaseId.of(instanceId, DATABASE);*/ + } + + private static void execute(String url, String sql) + { + try (Connection connection = DriverManager.getConnection(url, new Properties()); + Statement statement = connection.createStatement()) { + statement.execute(sql); + } + catch (SQLException e) { + throw new RuntimeException(e); + } + } + + private Database createDatabase() + throws InterruptedException, ExecutionException + { + DatabaseAdminClient dbAdminClient = options.getService().getDatabaseAdminClient(); + return dbAdminClient + .createDatabase( + INSTANCE, + DATABASE, new ArrayList<>()) + .get(); + } + + private InstanceId createInstance() + throws InterruptedException, ExecutionException + { + InstanceConfigId instanceConfig = InstanceConfigId.of(PROJECT, "emulator-config"); + InstanceId instanceId = InstanceId.of(PROJECT, INSTANCE); + InstanceAdminClient insAdminClient = spanner.getInstanceAdminClient(); + return insAdminClient + .createInstance( + InstanceInfo + .newBuilder(instanceId) + .setNodeCount(1) + .setDisplayName("Test instance") + .setInstanceConfigId(instanceConfig) + .build()) + .get().getId(); + } + + public void execute(@Language("SQL") String sql) + { + execute(getJdbcUrl(), sql); + } + + public String getJdbcUrl() + { + return String.format("jdbc:cloudspanner://%s/projects/%s/instances/%s/databases/%s;autoConfigEmulator=true", + emulatorContainer.getEmulatorGrpcEndpoint(), PROJECT, INSTANCE, DATABASE); + } + + @Override + public void close() + throws Exception + { + emulatorContainer.stop(); + } +} From dbe61fc6e154b134eb7860abb32817625dac21c5 Mon Sep 17 00:00:00 2001 From: taherk77 Date: Fri, 14 Apr 2023 16:29:59 +0530 Subject: [PATCH 5/8] Added licenses and changed sink --- .../spanner/SpannerSessionProperties.java | 13 ++++++ .../io/trino/plugin/spanner/SpannerSink.java | 16 ++++++- .../plugin/spanner/SpannerSinkProvider.java | 42 +++++++------------ 3 files changed, 44 insertions(+), 27 deletions(-) diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSessionProperties.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSessionProperties.java index f78c09f5c51c..0aec148c9971 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSessionProperties.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSessionProperties.java @@ -1,3 +1,16 @@ +/* + * 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.trino.plugin.spanner; import com.google.common.collect.ImmutableList; diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java index 580d14b89dc6..85af7c0906a2 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java @@ -1,3 +1,16 @@ +/* + * 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.trino.plugin.spanner; import com.google.cloud.NoCredentials; @@ -49,7 +62,8 @@ public class SpannerSink private final SpannerSessionProperties.Mode writeMode; private List mutations = new LinkedList<>(); - public SpannerSink(ConnectorSession session, JdbcOutputTableHandle handle, SpannerClient spannerClient, ConnectorPageSinkId pageSinkId, RemoteQueryModifier remoteQueryModifier) + public SpannerSink(ConnectorSession session, JdbcOutputTableHandle handle, + ConnectorPageSinkId pageSinkId) { this.options = SpannerOptions .newBuilder() diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java index 7196902f9a35..8ce297070b0f 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java @@ -1,3 +1,16 @@ +/* + * 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.trino.plugin.spanner; import com.google.inject.Inject; @@ -6,18 +19,12 @@ import io.trino.plugin.jdbc.JdbcOutputTableHandle; import io.trino.plugin.jdbc.logging.RemoteQueryModifier; import io.trino.spi.connector.ConnectorInsertTableHandle; -import io.trino.spi.connector.ConnectorMergeSink; -import io.trino.spi.connector.ConnectorMergeTableHandle; import io.trino.spi.connector.ConnectorOutputTableHandle; import io.trino.spi.connector.ConnectorPageSink; import io.trino.spi.connector.ConnectorPageSinkId; import io.trino.spi.connector.ConnectorPageSinkProvider; import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.ConnectorTableExecuteHandle; import io.trino.spi.connector.ConnectorTransactionHandle; -import io.trino.spi.session.PropertyMetadata; - -import java.util.List; public class SpannerSinkProvider implements ConnectorPageSinkProvider @@ -39,34 +46,17 @@ public SpannerSinkProvider( this.modifier = modifier; System.out.println("TABLE PROPS in SINK " + propertiesProvider.getTableProperties()); this.spannerTableProperties = propertiesProvider; - PropertyMetadata propertyMetadata = propertiesProvider.getTableProperties().get(0); - System.out.println(propertyMetadata.getDefaultValue()); } @Override public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorOutputTableHandle outputTableHandle, ConnectorPageSinkId pageSinkId) { - return new SpannerSink(session, (JdbcOutputTableHandle) outputTableHandle, client, pageSinkId, modifier); - } - - @Override - public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorInsertTableHandle insertTableHandle, ConnectorPageSinkId pageSinkId) - { - System.out.println("SinkProvider 2"); - return null; - } - - @Override - public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle, ConnectorPageSinkId pageSinkId) - { - System.out.println("SinkProvider 3"); - return null; + return new SpannerSink(session, (JdbcOutputTableHandle) outputTableHandle, pageSinkId); } @Override - public ConnectorMergeSink createMergeSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorMergeTableHandle mergeHandle, ConnectorPageSinkId pageSinkId) + public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorInsertTableHandle tableHandle, ConnectorPageSinkId pageSinkId) { - System.out.println("SinkProvider 4"); - return null; + return new SpannerSink(session, (JdbcOutputTableHandle) tableHandle, pageSinkId); } } From c8e5f5f6328cbcbb235a5b12014137d8b0cbc764 Mon Sep 17 00:00:00 2001 From: taherk77 Date: Tue, 6 Jun 2023 12:53:03 +0530 Subject: [PATCH 6/8] Added spanner transactions --- .../trino/plugin/spanner/SpannerClient.java | 11 +- .../trino/plugin/spanner/SpannerConfig.java | 61 +++++ .../plugin/spanner/SpannerConnector.java | 110 ++++++++ .../spanner/SpannerConnectorFactory.java | 60 +++++ .../spanner/SpannerConnectorMetadata.java | 252 ++++++++++++++++++ .../trino/plugin/spanner/SpannerModule.java | 161 +++++++---- .../trino/plugin/spanner/SpannerPlugin.java | 8 +- .../io/trino/plugin/spanner/SpannerSink.java | 75 ++++-- .../plugin/spanner/SpannerSinkProvider.java | 9 +- .../plugin/spanner/SpannerQueryRunner.java | 32 ++- .../spanner/TestingSpannerInstance.java | 52 +++- 11 files changed, 716 insertions(+), 115 deletions(-) create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnector.java create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnectorFactory.java create mode 100644 plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnectorMetadata.java diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java index 89d0045153f0..ce93d3e2ec5c 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerClient.java @@ -54,6 +54,8 @@ import io.trino.spi.type.TypeManager; import io.trino.spi.type.VarcharType; +import javax.inject.Inject; + import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.PreparedStatement; @@ -69,6 +71,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Locale; import java.util.Map; @@ -119,9 +122,13 @@ public class SpannerClient private final IdentifierMapping identifierMapping; private final String tableTypes[] = {"BASE TABLE", "VIEW"}; - public SpannerClient(BaseJdbcConfig config, SpannerConfig spannerConfig, JdbcStatisticsConfig statisticsConfig, ConnectionFactory connectionFactory, QueryBuilder queryBuilder, TypeManager typeManager, IdentifierMapping identifierMapping, RemoteQueryModifier queryModifier) + @Inject + public SpannerClient( + SpannerConfig spannerConfig, JdbcStatisticsConfig statisticsConfig, ConnectionFactory connectionFactory, QueryBuilder queryBuilder, TypeManager typeManager, IdentifierMapping identifierMapping, RemoteQueryModifier queryModifier) { - super("`", connectionFactory, queryBuilder, config.getJdbcTypesMappedToVarchar(), identifierMapping, queryModifier, true); + super("`", connectionFactory, queryBuilder, + new LinkedHashSet<>(), + identifierMapping, queryModifier, true); this.config = spannerConfig; this.identifierMapping = identifierMapping; } diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java index 3196fe52bf34..0755b89816d3 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConfig.java @@ -18,6 +18,67 @@ public class SpannerConfig { private String credentialsFile; + private String projectId; + + public String getHost() + { + return host; + } + @Config("spanner.emulated.host") + public void setHost(String host) + { + this.host = host; + } + + private String host = ""; + + public boolean isEmulator() + { + return isEmulator; + } + + @Config("spanner.emulated") + public void setEmulator(boolean emulator) + { + isEmulator = emulator; + } + + private boolean isEmulator = false; + private String instanceId; + private String database; + + public String getProjectId() + { + return projectId; + } + + @Config("spanner.projectId") + public void setProjectId(String projectId) + { + this.projectId = projectId; + } + + + public String getInstanceId() + { + return instanceId; + } + @Config("spanner.instanceId") + public void setInstanceId(String instanceId) + { + this.instanceId = instanceId; + } + + public String getDatabase() + { + return database; + } + + @Config("spanner.database") + public void setDatabase(String database) + { + this.database = database; + } /*private int minSessions = 100; private int maxSessions = 400; private int numChannels = 4; diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnector.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnector.java new file mode 100644 index 000000000000..83a77380f9e8 --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnector.java @@ -0,0 +1,110 @@ +package io.trino.plugin.spanner; + +import com.google.inject.Inject; +import io.airlift.bootstrap.LifeCycleManager; +import io.trino.plugin.base.session.SessionPropertiesProvider; +import io.trino.plugin.jdbc.JdbcTransactionHandle; +import io.trino.plugin.jdbc.TablePropertiesProvider; +import io.trino.spi.connector.Connector; +import io.trino.spi.connector.ConnectorMetadata; +import io.trino.spi.connector.ConnectorPageSinkProvider; +import io.trino.spi.connector.ConnectorRecordSetProvider; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorSplitManager; +import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.session.PropertyMetadata; +import io.trino.spi.transaction.IsolationLevel; + +import java.util.Collections; +import java.util.List; +import java.util.Set; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.util.Objects.requireNonNull; + +public class SpannerConnector implements Connector +{ + private final LifeCycleManager lifeCycleManager; + private final ConnectorMetadata metadata; + private final ConnectorSplitManager splitManager; + private final ConnectorRecordSetProvider recordSetProvider; + private final ConnectorPageSinkProvider pageSinkProvider; + private final List> tableProperties; + private final List> sessionProperties; + + @Inject + public SpannerConnector( + LifeCycleManager lifeCycleManager, + SpannerConnectorMetadata metadata, + ConnectorSplitManager splitManager, + ConnectorRecordSetProvider recordSetProvider, + ConnectorPageSinkProvider pageSinkProvider, + Set tableProperties, + Set sessionProperties) + { + this.lifeCycleManager = requireNonNull(lifeCycleManager, "lifeCycleManager is null"); + this.metadata = requireNonNull(metadata, "metadata is null"); + this.splitManager = requireNonNull(splitManager, "splitManager is null"); + this.recordSetProvider = requireNonNull(recordSetProvider, "recordSetProvider is null"); + this.pageSinkProvider = requireNonNull(pageSinkProvider, "pageSinkProvider is null"); + this.tableProperties = tableProperties.stream() + .flatMap(tablePropertiesProvider -> tablePropertiesProvider.getTableProperties().stream()) + .collect(toImmutableList()); + this.sessionProperties = sessionProperties.stream() + .flatMap(sessionPropertiesProvider -> sessionPropertiesProvider.getSessionProperties().stream()) + .collect(toImmutableList()); + } + @Override + public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel, boolean readOnly, boolean autoCommit) + { + return new JdbcTransactionHandle(); + } + + @Override + public ConnectorMetadata getMetadata(ConnectorSession session, ConnectorTransactionHandle transaction) + { + return metadata; + } + + @Override + public ConnectorSplitManager getSplitManager() + { + return splitManager; + } + + @Override + public ConnectorRecordSetProvider getRecordSetProvider() + { + return recordSetProvider; + } + + @Override + public ConnectorPageSinkProvider getPageSinkProvider() + { + return pageSinkProvider; + } + + @Override + public List> getTableProperties() + { + return tableProperties; + } + + @Override + public List> getColumnProperties() + { + return Collections.emptyList(); + } + + @Override + public List> getSessionProperties() + { + return sessionProperties; + } + + @Override + public final void shutdown() + { + lifeCycleManager.stop(); + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnectorFactory.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnectorFactory.java new file mode 100644 index 000000000000..d15fdebf4c41 --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnectorFactory.java @@ -0,0 +1,60 @@ +package io.trino.plugin.spanner; + +import com.google.inject.Injector; +import io.airlift.bootstrap.Bootstrap; +import io.airlift.json.JsonModule; +import io.trino.plugin.base.CatalogName; +import io.trino.spi.NodeManager; +import io.trino.spi.classloader.ThreadContextClassLoader; +import io.trino.spi.connector.Connector; +import io.trino.spi.connector.ConnectorContext; +import io.trino.spi.connector.ConnectorFactory; +import io.trino.spi.type.TypeManager; + +import java.util.Map; + +import static io.trino.plugin.base.Versions.checkSpiVersion; +import static java.util.Objects.requireNonNull; + +public class SpannerConnectorFactory + implements ConnectorFactory +{ + private final ClassLoader classLoader; + + public SpannerConnectorFactory(ClassLoader classLoader) + { + this.classLoader = requireNonNull(classLoader, "classLoader is null"); + } + + @Override + public String getName() + { + return "spanner"; + } + + @Override + public Connector create(String catalogName, Map requiredConfig, ConnectorContext context) + { + requireNonNull(requiredConfig, "requiredConfig is null"); + checkSpiVersion(context, this); + + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + Bootstrap app = new Bootstrap( + new JsonModule(), + new SpannerModule(catalogName), + binder -> { + binder.bind(CatalogName.class).toInstance(new CatalogName(catalogName)); + binder.bind(ClassLoader.class).toInstance(SpannerConnectorFactory.class.getClassLoader()); + binder.bind(TypeManager.class).toInstance(context.getTypeManager()); + binder.bind(NodeManager.class).toInstance(context.getNodeManager()); + }); + + Injector injector = app + .doNotInitializeLogging() + .setRequiredConfigurationProperties(requiredConfig) + .initialize(); + + return injector.getInstance(SpannerConnector.class); + } + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnectorMetadata.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnectorMetadata.java new file mode 100644 index 000000000000..a16ecf6160fc --- /dev/null +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerConnectorMetadata.java @@ -0,0 +1,252 @@ +package io.trino.plugin.spanner; + +import com.google.common.collect.ImmutableList; +import io.airlift.slice.Slice; +import io.trino.plugin.jdbc.DefaultJdbcMetadata; +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.JdbcNamedRelationHandle; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.plugin.jdbc.JdbcQueryEventListener; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.plugin.jdbc.RemoteTableName; +import io.trino.plugin.jdbc.mapping.IdentifierMapping; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.AggregateFunction; +import io.trino.spi.connector.AggregationApplicationResult; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorInsertTableHandle; +import io.trino.spi.connector.ConnectorOutputMetadata; +import io.trino.spi.connector.ConnectorOutputTableHandle; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTableHandle; +import io.trino.spi.connector.ConnectorTableLayout; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.ConnectorTableProperties; +import io.trino.spi.connector.ConnectorTableSchema; +import io.trino.spi.connector.LocalProperty; +import io.trino.spi.connector.RetryMode; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.connector.SortingProperty; +import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.security.TrinoPrincipal; +import io.trino.spi.statistics.ComputedStatistics; + +import javax.annotation.Nullable; +import javax.inject.Inject; + +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.connector.RetryMode.NO_RETRIES; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class SpannerConnectorMetadata + extends DefaultJdbcMetadata +{ + + public static final String DEFAULT_SCHEMA = "default"; + private final SpannerClient spannerClient; + private final IdentifierMapping identifierMapping; + + @Inject + public SpannerConnectorMetadata(SpannerClient spannerClient, + IdentifierMapping identifierMapping, + Set jdbcQueryEventListeners) + { + super(spannerClient, false, jdbcQueryEventListeners); + this.spannerClient = requireNonNull(spannerClient, "spannerClient is null"); + this.identifierMapping = requireNonNull(identifierMapping, "identifierMapping is null"); + } + + public static @Nullable + String toTrinoSchemaName(@Nullable String schema) + { + return "".equals(schema) ? DEFAULT_SCHEMA : schema; + } + + @Override + public JdbcTableHandle getTableHandle(ConnectorSession session, SchemaTableName schemaTableName) + { + return spannerClient.getTableHandle(session, schemaTableName) + .map(JdbcTableHandle::asPlainTable) + .map(JdbcNamedRelationHandle::getRemoteTableName) + .map(remoteTableName -> new JdbcTableHandle( + schemaTableName, + new RemoteTableName(remoteTableName.getCatalogName(), + Optional.ofNullable(toTrinoSchemaName(remoteTableName.getSchemaName().orElse(null))), remoteTableName.getTableName()), + Optional.empty())) + .orElse(null); + } + + @Override + public ConnectorTableProperties getTableProperties(ConnectorSession session, ConnectorTableHandle table) + { + JdbcTableHandle tableHandle = (JdbcTableHandle) table; + List> sortingProperties = tableHandle.getSortOrder() + .map(properties -> properties + .stream() + .map(item -> (LocalProperty) new SortingProperty( + item.getColumn(), + item.getSortOrder())) + .collect(toImmutableList())) + .orElse(ImmutableList.of()); + + return new ConnectorTableProperties(TupleDomain.all(), Optional.empty(), Optional.empty(), Optional.empty(), sortingProperties); + } + + @Override + public ConnectorTableSchema getTableSchema(ConnectorSession session, ConnectorTableHandle table) + { + JdbcTableHandle handle = (JdbcTableHandle) table; + return new ConnectorTableSchema( + getSchemaTableName(handle), + getColumnMetadata(session, handle).stream() + .map(ColumnMetadata::getColumnSchema) + .collect(toImmutableList())); + } + + @Override + public ConnectorTableMetadata getTableMetadata(ConnectorSession session, ConnectorTableHandle table) + { + JdbcTableHandle handle = (JdbcTableHandle) table; + return new ConnectorTableMetadata( + getSchemaTableName(handle), + getColumnMetadata(session, handle), + spannerClient.getTableProperties(session, handle)); + } + + private List getColumnMetadata(ConnectorSession session, JdbcTableHandle handle) + { + return spannerClient.getColumns(session, handle).stream() + .map(JdbcColumnHandle::getColumnMetadata) + .collect(toImmutableList()); + } + + @Override + public void createSchema(ConnectorSession session, String schemaName, Map properties, TrinoPrincipal owner) + { + throw new TrinoException(NOT_SUPPORTED, "Can't create schemas in Spanner"); + } + + @Override + public void dropSchema(ConnectorSession session, String schemaName) + { + throw new TrinoException(NOT_SUPPORTED, "Can't drop schema which in spanner"); + } + + @Override + public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean ignoreExisting) + { + spannerClient.beginCreateTable(session, tableMetadata); + } + + @Override + public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Optional layout, RetryMode retryMode) + { + if (retryMode != NO_RETRIES) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support query retries"); + } + return spannerClient.beginCreateTable(session, tableMetadata); + } + + @Override + public Optional finishCreateTable(ConnectorSession session, ConnectorOutputTableHandle tableHandle, Collection fragments, Collection computedStatistics) + { + return Optional.empty(); + } + + @Override + public boolean supportsMissingColumnsOnInsert() + { + return true; + } + + @Override + public ConnectorInsertTableHandle beginInsert(ConnectorSession session, ConnectorTableHandle tableHandle, List columns, RetryMode retryMode) + { + if (retryMode != NO_RETRIES) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support query retries"); + } + JdbcTableHandle handle = (JdbcTableHandle) tableHandle; + + List columnHandles = columns.stream() + .map(JdbcColumnHandle.class::cast) + .collect(toImmutableList()); + + RemoteTableName remoteTableName = handle.asPlainTable().getRemoteTableName(); + return new JdbcOutputTableHandle( + "spanner", + remoteTableName.getSchemaName().orElse(null), + remoteTableName.getTableName(), + columnHandles.stream().map(JdbcColumnHandle::getColumnName).collect(toImmutableList()), + columnHandles.stream().map(JdbcColumnHandle::getColumnType).collect(toImmutableList()), + Optional.of(columnHandles.stream().map(JdbcColumnHandle::getJdbcTypeHandle).collect(toImmutableList())), + Optional.empty(), + Optional.empty()); + } + + @Override + public Optional finishInsert(ConnectorSession session, ConnectorInsertTableHandle insertHandle, Collection fragments, Collection computedStatistics) + { + return Optional.empty(); + } + + @Override + public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnMetadata column) + { + if (column.getComment() != null) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support adding columns with comments"); + } + + JdbcTableHandle handle = (JdbcTableHandle) tableHandle; + RemoteTableName remoteTableName = handle.asPlainTable().getRemoteTableName(); + spannerClient.execute(session, format( + "ALTER TABLE %s ADD %s %s", + remoteTableName.getTableName(), + spannerClient.quoted(column.getName()), + spannerClient.toWriteMapping(session, column.getType()).getDataType())); + } + + @Override + public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column) + { + JdbcTableHandle handle = (JdbcTableHandle) tableHandle; + JdbcColumnHandle columnHandle = (JdbcColumnHandle) column; + RemoteTableName remoteTableName = handle.asPlainTable().getRemoteTableName(); + spannerClient.execute(session, format( + "ALTER TABLE %s DROP COLUMN %s", + remoteTableName.getTableName(), + spannerClient.quoted(columnHandle.getColumnName()))); + } + + @Override + public void dropTable(ConnectorSession session, ConnectorTableHandle tableHandle) + { + spannerClient.dropTable(session, (JdbcTableHandle) tableHandle); + } + + @Override + public void truncateTable(ConnectorSession session, ConnectorTableHandle tableHandle) + { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support truncating tables"); + } + + @Override + public Optional> applyAggregation( + ConnectorSession session, + ConnectorTableHandle table, + List aggregates, + Map assignments, + List> groupingSets) + { + // TODO support aggregation pushdown + return Optional.empty(); + } +} diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java index 4ec12cf3c5af..1100f484414a 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerModule.java @@ -14,7 +14,6 @@ package io.trino.plugin.spanner; import com.google.cloud.spanner.jdbc.JdbcDriver; -import com.google.common.util.concurrent.MoreExecutors; import com.google.inject.Binder; import com.google.inject.Key; import com.google.inject.Provides; @@ -22,11 +21,12 @@ import com.google.inject.Singleton; import com.google.inject.multibindings.Multibinder; import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.trino.plugin.base.classloader.ForClassLoaderSafe; import io.trino.plugin.base.session.SessionPropertiesProvider; import io.trino.plugin.jdbc.BaseJdbcConfig; -import io.trino.plugin.jdbc.CachingJdbcClient; import io.trino.plugin.jdbc.ConfiguringConnectionFactory; import io.trino.plugin.jdbc.ConnectionFactory; +import io.trino.plugin.jdbc.DecimalModule; import io.trino.plugin.jdbc.DefaultJdbcMetadataFactory; import io.trino.plugin.jdbc.DefaultQueryBuilder; import io.trino.plugin.jdbc.DriverConnectionFactory; @@ -36,14 +36,15 @@ import io.trino.plugin.jdbc.ForLazyConnectionFactory; import io.trino.plugin.jdbc.ForRecordCursor; import io.trino.plugin.jdbc.JdbcClient; -import io.trino.plugin.jdbc.JdbcConnector; +import io.trino.plugin.jdbc.JdbcDiagnosticModule; import io.trino.plugin.jdbc.JdbcDynamicFilteringConfig; import io.trino.plugin.jdbc.JdbcDynamicFilteringSessionProperties; import io.trino.plugin.jdbc.JdbcDynamicFilteringSplitManager; -import io.trino.plugin.jdbc.JdbcJoinPushdownSupportModule; import io.trino.plugin.jdbc.JdbcMetadataConfig; import io.trino.plugin.jdbc.JdbcMetadataFactory; import io.trino.plugin.jdbc.JdbcMetadataSessionProperties; +import io.trino.plugin.jdbc.JdbcPageSinkProvider; +import io.trino.plugin.jdbc.JdbcQueryEventListener; import io.trino.plugin.jdbc.JdbcRecordSetProvider; import io.trino.plugin.jdbc.JdbcSplitManager; import io.trino.plugin.jdbc.JdbcStatisticsConfig; @@ -53,73 +54,96 @@ import io.trino.plugin.jdbc.LazyConnectionFactory; import io.trino.plugin.jdbc.MaxDomainCompactionThreshold; import io.trino.plugin.jdbc.QueryBuilder; -import io.trino.plugin.jdbc.QueryConfig; import io.trino.plugin.jdbc.ReusableConnectionFactoryModule; import io.trino.plugin.jdbc.StatsCollecting; import io.trino.plugin.jdbc.TablePropertiesProvider; import io.trino.plugin.jdbc.TypeHandlingJdbcConfig; import io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties; -import io.trino.plugin.jdbc.credential.CredentialProvider; +import io.trino.plugin.jdbc.credential.EmptyCredentialProvider; import io.trino.plugin.jdbc.logging.RemoteQueryModifier; import io.trino.plugin.jdbc.logging.RemoteQueryModifierModule; import io.trino.plugin.jdbc.mapping.IdentifierMapping; import io.trino.plugin.jdbc.mapping.IdentifierMappingModule; -import io.trino.plugin.jdbc.procedure.FlushJdbcMetadataCacheProcedure; -import io.trino.spi.connector.ConnectorAccessControl; +import io.trino.spi.connector.ConnectorMetadata; import io.trino.spi.connector.ConnectorPageSinkProvider; import io.trino.spi.connector.ConnectorRecordSetProvider; import io.trino.spi.connector.ConnectorSplitManager; import io.trino.spi.procedure.Procedure; -import io.trino.spi.ptf.ConnectorTableFunction; import io.trino.spi.type.TypeManager; import javax.annotation.PreDestroy; import javax.inject.Provider; import java.util.Properties; +import java.util.Set; import java.util.concurrent.ExecutorService; +import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; import static com.google.inject.multibindings.Multibinder.newSetBinder; import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; import static io.airlift.configuration.ConditionalModule.conditionalModule; import static io.airlift.configuration.ConfigBinder.configBinder; -import static io.trino.plugin.jdbc.JdbcModule.bindSessionPropertiesProvider; +import static java.util.Objects.requireNonNull; +import static org.weakref.jmx.guice.ExportBinder.newExporter; public class SpannerModule extends AbstractConfigurationAwareModule { + private final String catalogName; + + public SpannerModule(String catalogName) + { + this.catalogName = requireNonNull(catalogName, "catalogName is null"); + } @Provides @Singleton @ForBaseJdbc - public static ConnectionFactory createConnectionFactory(BaseJdbcConfig config, - CredentialProvider credentialProvider, + public static ConnectionFactory createConnectionFactory( SpannerConfig spannerConfig) throws ClassNotFoundException { Class.forName("com.google.cloud.spanner.jdbc.JdbcDriver"); Properties connectionProperties = new Properties(); - String connectionUrl = config.getConnectionUrl(); + String connectionUrlTemplate = "jdbc:cloudspanner:/%s/projects/%s/instances/%s/databases/%s%s"; + String host = "/"; + String configureEmulator = ";autoConfigEmulator=true"; + if (spannerConfig.isEmulator()) { + host = host + spannerConfig.getHost(); + } + String url = String.format(connectionUrlTemplate, host, spannerConfig.getProjectId(), spannerConfig.getInstanceId(), spannerConfig.getDatabase(), configureEmulator); JdbcDriver driver = new JdbcDriver(); + System.out.println("USING connection URL " + url); //File credentials = new File(spannerConfig.getCredentialsFile()); - if (!driver.acceptsURL(connectionUrl)) { - throw new RuntimeException(config.getConnectionUrl() + " is incorrect"); + if (!driver.acceptsURL(url)) { + throw new RuntimeException( + url + " is incorrect"); } + //jdbc:cloudspanner://0.0.0.0:9010/projects/test-project/instances/test-instance/databases/trinodb;autoConfigEmulator=true + //jdbc:cloudspanner://localhost:9010/projects/test-project/instances/test-instance/databases/test-db;usePlainText=true //connectionProperties.put("credentials", spannerConfig.getCredentialsFile()); connectionProperties.setProperty("retryAbortsInternally", "true"); return new ConfiguringConnectionFactory(new DriverConnectionFactory( driver, - config.getConnectionUrl(), + url, connectionProperties, - credentialProvider), + new EmptyCredentialProvider()), connection -> { }); } + @Provides + @Singleton + public static JdbcStatisticsConfig getConf() + { + JdbcStatisticsConfig jdbcStatisticsConfig = new JdbcStatisticsConfig(); + jdbcStatisticsConfig.setEnabled(true); + return jdbcStatisticsConfig; + } + @Provides @Singleton public static SpannerClient getSpannerClient( - BaseJdbcConfig config, SpannerConfig spannerConfig, JdbcStatisticsConfig statisticsConfig, ConnectionFactory connectionFactory, @@ -128,7 +152,7 @@ public static SpannerClient getSpannerClient( IdentifierMapping identifierMapping, RemoteQueryModifier queryModifier) { - return new SpannerClient(config, + return new SpannerClient( spannerConfig, statisticsConfig, connectionFactory, @@ -143,15 +167,16 @@ public static SpannerTableProperties tableProperties() { return new SpannerTableProperties(); } + @Singleton - public static SpannerSinkProvider configureSnowflakeSink( + public static SpannerSinkProvider configureSpannerSink( ConnectionFactory factory, - JdbcClient client, RemoteQueryModifier modifier, - SpannerClient snowflakeConfig, + SpannerClient spannerClient, + SpannerConfig config, SpannerTableProperties propertiesProvider) { - return new SpannerSinkProvider(factory, client, modifier, snowflakeConfig, propertiesProvider); + return new SpannerSinkProvider(factory, modifier, spannerClient, config, propertiesProvider); } public static Multibinder sessionPropertiesProviderBinder(Binder binder) @@ -184,22 +209,39 @@ public static void bindTablePropertiesProvider(Binder binder, Class config.setDomainCompactionThreshold(DEFAULT_DOMAIN_COMPACTION_THRESHOLD)); - newOptionalBinder(binder, ConnectorAccessControl.class); - newOptionalBinder(binder, QueryBuilder.class).setDefault().to(DefaultQueryBuilder.class).in(Scopes.SINGLETON); + configBinder(binder).bindConfig(TypeHandlingJdbcConfig.class); + bindSessionPropertiesProvider(binder, TypeHandlingJdbcSessionProperties.class); + bindSessionPropertiesProvider(binder, JdbcMetadataSessionProperties.class); + bindSessionPropertiesProvider(binder, JdbcWriteSessionProperties.class); + bindSessionPropertiesProvider(binder, SpannerSessionProperties.class); + bindSessionPropertiesProvider(binder, JdbcDynamicFilteringSessionProperties.class); + newOptionalBinder(binder, Key.get(ConnectorSplitManager.class, ForJdbcDynamicFiltering.class)).setDefault().to(JdbcSplitManager.class).in(Scopes.SINGLETON); - procedureBinder(binder); - tablePropertiesProviderBinder(binder); + binder.bind(DynamicFilteringStats.class).in(Scopes.SINGLETON); + newExporter(binder).export(DynamicFilteringStats.class) + .as(generator -> generator.generatedNameOf(DynamicFilteringStats.class, catalogName)); newOptionalBinder(binder, JdbcMetadataFactory.class).setBinding().to(DefaultJdbcMetadataFactory.class).in(Scopes.SINGLETON); newOptionalBinder(binder, Key.get(ConnectorSplitManager.class, ForJdbcDynamicFiltering.class)).setDefault().to(JdbcSplitManager.class).in(Scopes.SINGLETON); @@ -208,45 +250,29 @@ public void setup(Binder binder) newOptionalBinder(binder, ConnectorPageSinkProvider.class).setBinding().to(SpannerSinkProvider.class).in(Scopes.SINGLETON); binder.bind(JdbcTransactionManager.class).in(Scopes.SINGLETON); - binder.bind(JdbcConnector.class).in(Scopes.SINGLETON); - configBinder(binder).bindConfig(JdbcMetadataConfig.class); configBinder(binder).bindConfig(JdbcWriteConfig.class); - configBinder(binder).bindConfig(BaseJdbcConfig.class); - configBinder(binder).bindConfig(SpannerConfig.class); configBinder(binder).bindConfig(JdbcDynamicFilteringConfig.class); - bindTablePropertiesProvider(binder, SpannerTableProperties.class); - - configBinder(binder).bindConfig(TypeHandlingJdbcConfig.class); - bindSessionPropertiesProvider(binder, TypeHandlingJdbcSessionProperties.class); - bindSessionPropertiesProvider(binder, JdbcMetadataSessionProperties.class); - bindSessionPropertiesProvider(binder, JdbcWriteSessionProperties.class); - bindSessionPropertiesProvider(binder, JdbcDynamicFilteringSessionProperties.class); - bindSessionPropertiesProvider(binder, SpannerSessionProperties.class); - binder.bind(DynamicFilteringStats.class).in(Scopes.SINGLETON); - - binder.bind(CachingJdbcClient.class).in(Scopes.SINGLETON); - binder.bind(JdbcClient.class).to(Key.get(CachingJdbcClient.class)).in(Scopes.SINGLETON); - - newSetBinder(binder, Procedure.class).addBinding().toProvider(FlushJdbcMetadataCacheProcedure.class).in(Scopes.SINGLETON); - newSetBinder(binder, ConnectorTableFunction.class); + binder.bind(JdbcClient.class).annotatedWith(ForBaseJdbc.class).to(Key.get(SpannerClient.class)).in(Scopes.SINGLETON); + binder.bind(JdbcClient.class).to(Key.get(JdbcClient.class, StatsCollecting.class)).in(Scopes.SINGLETON); + binder.bind(ConnectorMetadata.class).annotatedWith(ForClassLoaderSafe.class).to(SpannerConnectorMetadata.class).in(Scopes.SINGLETON); binder.bind(ConnectionFactory.class) .annotatedWith(ForLazyConnectionFactory.class) .to(Key.get(ConnectionFactory.class, StatsCollecting.class)) .in(Scopes.SINGLETON); install(conditionalModule( - QueryConfig.class, - QueryConfig::isReuseConnection, + SpannerConfig.class, + (p) -> true, new ReusableConnectionFactoryModule(), innerBinder -> innerBinder.bind(ConnectionFactory.class).to(LazyConnectionFactory.class).in(Scopes.SINGLETON))); - newOptionalBinder(binder, Key.get(int.class, MaxDomainCompactionThreshold.class)); + bindTablePropertiesProvider(binder, SpannerTableProperties.class); - newOptionalBinder(binder, Key.get(ExecutorService.class, ForRecordCursor.class)) - .setBinding() - .toProvider(MoreExecutors::newDirectExecutorService) - .in(Scopes.SINGLETON); + binder.bind(SpannerConnector.class).in(Scopes.SINGLETON); + install(new JdbcDiagnosticModule()); + install(new IdentifierMappingModule()); + install(new DecimalModule()); } @PreDestroy @@ -254,4 +280,21 @@ public void shutdownRecordCursorExecutor(@ForRecordCursor ExecutorService execut { executor.shutdownNow(); } + + @Singleton + @ForRecordCursor + @Provides + public ExecutorService createRecordCursorExecutor() + { + return newDirectExecutorService(); + } + + @Singleton + @Provides + public SpannerConnectorMetadata getSpannerMetadata( + SpannerClient client, IdentifierMapping mapping, + Set listeners) + { + return new SpannerConnectorMetadata(client, mapping, listeners); + } } diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerPlugin.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerPlugin.java index eb65e24db19c..d72b06046207 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerPlugin.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerPlugin.java @@ -13,16 +13,18 @@ */ package io.trino.plugin.spanner; +import com.google.common.collect.ImmutableList; import io.trino.plugin.jdbc.JdbcPlugin; import io.trino.spi.Plugin; +import io.trino.spi.connector.ConnectorFactory; public class SpannerPlugin - extends JdbcPlugin implements Plugin { - public SpannerPlugin() + @Override + public Iterable getConnectorFactories() { - super("spanner", new SpannerModule()); + return ImmutableList.of(new SpannerConnectorFactory(SpannerPlugin.class.getClassLoader())); } } diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java index 85af7c0906a2..6e9fb8380a39 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSink.java @@ -14,33 +14,32 @@ package io.trino.plugin.spanner; import com.google.cloud.NoCredentials; -import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AbortedException; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.DatabaseId; -import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; -import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerOptions; +import com.google.cloud.spanner.TransactionContext; +import com.google.cloud.spanner.TransactionManager; import com.google.common.collect.ImmutableList; import io.airlift.slice.Slice; import io.airlift.slice.Slices; -import io.grpc.Status; -import io.grpc.StatusRuntimeException; import io.trino.plugin.jdbc.JdbcOutputTableHandle; -import io.trino.plugin.jdbc.logging.RemoteQueryModifier; import io.trino.spi.Page; -import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.connector.ConnectorPageSink; import io.trino.spi.connector.ConnectorPageSinkId; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.Type; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.time.LocalDate; import java.util.Collection; import java.util.LinkedList; import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; import static io.trino.plugin.jdbc.JdbcWriteSessionProperties.getWriteBatchSize; import static java.time.format.DateTimeFormatter.ISO_DATE; @@ -54,27 +53,33 @@ public class SpannerSink private final DatabaseClient client; private final List columnTypes; private final List columnNames; - private final String project = "spanner-project"; - private final String instance = "spanner-instance"; - private final String database = "spanner-database"; private final String table; private final ConnectorPageSinkId pageSinkId; private final SpannerSessionProperties.Mode writeMode; + private final boolean isEmulatedHost; + private final Logger LOG = LoggerFactory.getLogger(SpannerSink.class); + int maxRetries = 3; + int retry = 0; private List mutations = new LinkedList<>(); - public SpannerSink(ConnectorSession session, JdbcOutputTableHandle handle, + public SpannerSink(SpannerConfig config, ConnectorSession session, JdbcOutputTableHandle handle, ConnectorPageSinkId pageSinkId) { - this.options = SpannerOptions + isEmulatedHost = config.isEmulator(); + SpannerOptions.Builder builder = SpannerOptions .newBuilder() - .setEmulatorHost("0.0.0.0:9010") - .setCredentials(NoCredentials.getInstance()) - .setProjectId(project) - .build(); + .setProjectId(config.getProjectId()); + if (isEmulatedHost) { + builder.setEmulatorHost(config.getHost()) + .setCredentials(NoCredentials.getInstance()); + } + else { + //builder.setCredentials(Credentials); + } + this.options = builder.build(); this.pageSinkId = pageSinkId; this.maxBatchSize = getWriteBatchSize(session); - - this.client = options.getService().getDatabaseClient(DatabaseId.of(project, instance, database)); + this.client = options.getService().getDatabaseClient(DatabaseId.of(config.getProjectId(), config.getInstanceId(), config.getDatabase())); columnTypes = handle.getColumnTypes(); columnNames = handle.getColumnNames(); table = handle.getTableName(); @@ -141,20 +146,32 @@ else if (javaType == Slice.class) { private void write() { if (!mutations.isEmpty()) { - try { - Timestamp write = client.write(mutations); - System.out.println("Batch write completed " + write + " " + mutations.size() + " records flushed"); - } - catch (Exception e) { - System.out.println(e); - if (e instanceof SpannerException spannerEx) { - if (spannerEx.getErrorCode().equals(ErrorCode.ALREADY_EXISTS)) { - throw new TrinoException(SpannerClient.SpannerErrorCode.SPANNER_ERROR_CODE, String.format("%s Try changing %s to %s and retry this query ", spannerEx.getMessage(), SpannerSessionProperties.WRITE_MODE, - SpannerSessionProperties.Mode.UPSERT)); + try (TransactionManager manager = client.transactionManager()) { + TransactionContext transaction = manager.begin(); + while (true) { + try { + transaction.buffer(mutations); + manager.commit(); + break; + } + catch (AbortedException e) { + blockFor(e.getRetryDelayInMillis()); + transaction = manager.resetForRetry(); } } } - mutations = new LinkedList<>(); + } + mutations = new LinkedList<>(); + } + + private void blockFor(long delay) + { + try { + if (delay > 0L) { + TimeUnit.MILLISECONDS.sleep(delay); + } + } + catch (InterruptedException ignored) { } } diff --git a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java index 8ce297070b0f..f029875de974 100644 --- a/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java +++ b/plugin/trino-spanner/src/main/java/io/trino/plugin/spanner/SpannerSinkProvider.java @@ -15,7 +15,6 @@ import com.google.inject.Inject; import io.trino.plugin.jdbc.ConnectionFactory; -import io.trino.plugin.jdbc.JdbcClient; import io.trino.plugin.jdbc.JdbcOutputTableHandle; import io.trino.plugin.jdbc.logging.RemoteQueryModifier; import io.trino.spi.connector.ConnectorInsertTableHandle; @@ -32,13 +31,14 @@ public class SpannerSinkProvider private final SpannerClient client; private final RemoteQueryModifier modifier; private final SpannerTableProperties spannerTableProperties; + private final SpannerConfig config; @Inject public SpannerSinkProvider( ConnectionFactory connectionFactory, - JdbcClient jdbcClient, RemoteQueryModifier modifier, SpannerClient client, + SpannerConfig config, SpannerTableProperties propertiesProvider) { System.out.println("Called Spanner Sink provider"); @@ -46,17 +46,18 @@ public SpannerSinkProvider( this.modifier = modifier; System.out.println("TABLE PROPS in SINK " + propertiesProvider.getTableProperties()); this.spannerTableProperties = propertiesProvider; + this.config = config; } @Override public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorOutputTableHandle outputTableHandle, ConnectorPageSinkId pageSinkId) { - return new SpannerSink(session, (JdbcOutputTableHandle) outputTableHandle, pageSinkId); + return new SpannerSink(config, session, (JdbcOutputTableHandle) outputTableHandle, pageSinkId); } @Override public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorInsertTableHandle tableHandle, ConnectorPageSinkId pageSinkId) { - return new SpannerSink(session, (JdbcOutputTableHandle) tableHandle, pageSinkId); + return new SpannerSink(config, session, (JdbcOutputTableHandle) tableHandle, pageSinkId); } } diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java index 74983d37f87c..65c91f7c3a72 100644 --- a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java @@ -47,7 +47,7 @@ public static DistributedQueryRunner createSpannerQueryRunner( Iterable> tables) throws Exception { - return createSpannerQueryRunner(instance, extraProperties, Map.of(), connectorProperties, tables, runner -> {}); + return createSpannerQueryRunner(instance, extraProperties, ImmutableMap.of(), connectorProperties, tables, runner -> {}); } public static DistributedQueryRunner createSpannerQueryRunner( @@ -61,7 +61,8 @@ public static DistributedQueryRunner createSpannerQueryRunner( { DistributedQueryRunner queryRunner = null; try { - queryRunner = DistributedQueryRunner.builder(createSession()) + queryRunner = DistributedQueryRunner.builder( + createSession()) .setExtraProperties(extraProperties) .setCoordinatorProperties(coordinatorProperties) .setAdditionalSetup(moreSetup) @@ -72,14 +73,26 @@ public static DistributedQueryRunner createSpannerQueryRunner( // note: additional copy via ImmutableList so that if fails on nulls connectorProperties = new HashMap<>(ImmutableMap.copyOf(connectorProperties)); - //connectorProperties.putIfAbsent("connection-url", instance.getJdbcUrl()); - connectorProperties.putIfAbsent("connection-url", "jdbc:cloudspanner://0.0.0.0:9010/projects/spanner-project/instances/spanner-instance/databases/spanner-database;autoConfigEmulator=true"); - connectorProperties.putIfAbsent("connection-user", ""); - connectorProperties.putIfAbsent("connection-password", ""); connectorProperties.putIfAbsent("spanner.credentials.file", "credentials.json"); + connectorProperties.putIfAbsent("spanner.instanceId", instance.getInstanceId()); + connectorProperties.putIfAbsent("spanner.projectId", instance.getProjectId()); + connectorProperties.putIfAbsent("spanner.database", instance.getDatabaseId()); + connectorProperties.putIfAbsent("spanner.emulated", "true"); + connectorProperties.putIfAbsent("spanner.emulated.host", instance.getHost()); + /* connectorProperties = new HashMap<>(ImmutableMap.copyOf(connectorProperties)); + connectorProperties.putIfAbsent("spanner.credentials.file", "credentials.json"); + connectorProperties.putIfAbsent("spanner.instanceId", "spanner-instance"); + connectorProperties.putIfAbsent("spanner.projectId", "spanner-project"); + connectorProperties.putIfAbsent("spanner.database", "spanner-database"); + connectorProperties.putIfAbsent("spanner.emulated", "true"); + connectorProperties.putIfAbsent("spanner.emulated.host", "localhost:9010");*/ queryRunner.installPlugin(new SpannerPlugin()); queryRunner.createCatalog("spanner", "spanner", connectorProperties); - copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME,createSession(), tables); + MaterializedResult execute1 = queryRunner.execute("create table emp WITH (PRIMARY_KEYS = ARRAY['id']) as select 1 as id,'T' as name"); + System.out.println(execute1); + MaterializedResult execute2 = queryRunner.execute("select * from emp"); + System.out.println(execute2); + copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, createSession(), tables); MaterializedResult execute = queryRunner.execute("SHOW TABLES FROM spanner.default"); System.out.println(execute); /* @@ -99,7 +112,9 @@ public static Session createSession() return testSessionBuilder() .setCatalog("spanner") .setSchema("default") - .setCatalogSessionProperty("spanner",SpannerSessionProperties.WRITE_MODE, SpannerSessionProperties.Mode.UPSERT.name()) + .setCatalogSessionProperty("spanner", + SpannerSessionProperties.WRITE_MODE, + SpannerSessionProperties.Mode.UPSERT.name()) .build(); } @@ -150,6 +165,7 @@ private static void copyTable( System.out.println(sql); LOG.debug("Running import for %s %s", target, sql); long rows = queryRunner.execute(session, sql).getUpdateCount().getAsLong(); + LOG.debug("%s rows loaded into %s", rows, target); } } diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestingSpannerInstance.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestingSpannerInstance.java index 323ccfe4bf17..acd39deb708a 100644 --- a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestingSpannerInstance.java +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestingSpannerInstance.java @@ -18,8 +18,8 @@ import java.sql.DriverManager; import java.sql.SQLException; import java.sql.Statement; +import java.time.Duration; import java.util.ArrayList; -import java.util.Arrays; import java.util.Properties; import java.util.concurrent.ExecutionException; @@ -29,20 +29,24 @@ public class TestingSpannerInstance private final String SPANNER_IMAGE = "gcr.io/cloud-spanner-emulator/emulator:latest"; private final SpannerEmulatorContainer emulatorContainer; - private final SpannerOptions options=null; private final String PROJECT = "test-project"; private final String INSTANCE = "test-instance"; private final String DATABASE = "trinodb"; - private final Spanner spanner=null; - private final DatabaseId databaseId=null; - private final InstanceId instanceId=null; + private Database database; + private SpannerOptions options; + private Spanner spanner; + private DatabaseId databaseId; + private InstanceId instanceId; + private String host = null; public TestingSpannerInstance() throws ExecutionException, InterruptedException { - this.emulatorContainer = new SpannerEmulatorContainer(DockerImageName.parse(SPANNER_IMAGE)); - emulatorContainer.setExposedPorts(Arrays.asList(9010)); - /*emulatorContainer.start(); + this.emulatorContainer = new SpannerEmulatorContainer(DockerImageName.parse(SPANNER_IMAGE)) + .withExposedPorts(9010, 9020); + emulatorContainer.start(); + //this.host = "0.0.0.0:" + emulatorContainer.getEmulatorGrpcEndpoint().split(":")[1]; + this.host=emulatorContainer.getEmulatorGrpcEndpoint(); options = SpannerOptions .newBuilder() .setEmulatorHost(emulatorContainer.getEmulatorGrpcEndpoint()) @@ -51,8 +55,9 @@ public TestingSpannerInstance() .build(); this.spanner = options.getService(); this.instanceId = createInstance(); - Database database = createDatabase(); - this.databaseId = DatabaseId.of(instanceId, DATABASE);*/ + this.database = createDatabase(); + + this.databaseId = DatabaseId.of(instanceId, DATABASE); } private static void execute(String url, String sql) @@ -66,6 +71,33 @@ private static void execute(String url, String sql) } } + public static void main(String[] args) + throws ExecutionException, InterruptedException + { + TestingSpannerInstance spannerInstance = new TestingSpannerInstance(); + Thread.sleep(Duration.ofHours(1).toMillis()); + } + + public String getDatabaseId() + { + return DATABASE; + } + + public String getInstanceId() + { + return INSTANCE; + } + + public String getProjectId() + { + return PROJECT; + } + + public String getHost() + { + return host; + } + private Database createDatabase() throws InterruptedException, ExecutionException { From 6555aa8e77e40c4b08b4a70d0995451e92a4fda3 Mon Sep 17 00:00:00 2001 From: taherk77 Date: Thu, 15 Jun 2023 12:58:10 +0530 Subject: [PATCH 7/8] Added dtype test case for spanner --- .../plugin/spanner/SpannerQueryRunner.java | 38 ++++--- .../spanner/TestSpannerDataTypesMapping.java | 104 ++++++++++++++++++ .../plugin/spanner/TestSpannerPlugin.java | 17 ++- 3 files changed, 141 insertions(+), 18 deletions(-) create mode 100644 plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerDataTypesMapping.java diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java index 65c91f7c3a72..7b25a4b3f110 100644 --- a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerQueryRunner.java @@ -23,6 +23,7 @@ import io.trino.testing.QueryRunner; import io.trino.tpch.TpchTable; import org.intellij.lang.annotations.Language; +import org.jetbrains.annotations.NotNull; import java.util.HashMap; import java.util.Map; @@ -44,10 +45,10 @@ public static DistributedQueryRunner createSpannerQueryRunner( TestingSpannerInstance instance, Map extraProperties, Map connectorProperties, - Iterable> tables) + Iterable> tables, boolean addTpcDsTables) throws Exception { - return createSpannerQueryRunner(instance, extraProperties, ImmutableMap.of(), connectorProperties, tables, runner -> {}); + return createSpannerQueryRunner(instance, extraProperties, ImmutableMap.of(), connectorProperties, tables, runner -> {}, addTpcDsTables); } public static DistributedQueryRunner createSpannerQueryRunner( @@ -56,7 +57,8 @@ public static DistributedQueryRunner createSpannerQueryRunner( Map coordinatorProperties, Map connectorProperties, Iterable> tables, - Consumer moreSetup) + Consumer moreSetup, + boolean addTpcDsTables) throws Exception { DistributedQueryRunner queryRunner = null; @@ -88,13 +90,12 @@ public static DistributedQueryRunner createSpannerQueryRunner( connectorProperties.putIfAbsent("spanner.emulated.host", "localhost:9010");*/ queryRunner.installPlugin(new SpannerPlugin()); queryRunner.createCatalog("spanner", "spanner", connectorProperties); - MaterializedResult execute1 = queryRunner.execute("create table emp WITH (PRIMARY_KEYS = ARRAY['id']) as select 1 as id,'T' as name"); - System.out.println(execute1); - MaterializedResult execute2 = queryRunner.execute("select * from emp"); - System.out.println(execute2); - copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, createSession(), tables); - MaterializedResult execute = queryRunner.execute("SHOW TABLES FROM spanner.default"); - System.out.println(execute); + if (addTpcDsTables) { + copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, createSession(), tables); + MaterializedResult execute = queryRunner.execute("SHOW TABLES FROM spanner.default"); + System.out.println(execute); + } + /* MaterializedResult rows = queryRunner.execute("SELECT * FROM spanner.default.customer"); System.out.println(rows); @@ -121,11 +122,7 @@ public static Session createSession() public static void main(String[] args) throws Exception { - DistributedQueryRunner queryRunner = createSpannerQueryRunner( - new TestingSpannerInstance(), - ImmutableMap.of("http-server.http.port", "8080"), - ImmutableMap.of(), - TpchTable.getTables()); + DistributedQueryRunner queryRunner = getSpannerQueryRunner(); queryRunner.installPlugin(new JmxPlugin()); queryRunner.createCatalog("jmx", "jmx"); @@ -135,6 +132,17 @@ public static void main(String[] args) log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl()); } + @NotNull + public static DistributedQueryRunner getSpannerQueryRunner() + throws Exception + { + return createSpannerQueryRunner( + new TestingSpannerInstance(), + ImmutableMap.of("http-server.http.port", "8080"), + ImmutableMap.of(), + TpchTable.getTables(), false); + } + private static void copyTpchTables( QueryRunner queryRunner, String sourceCatalog, diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerDataTypesMapping.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerDataTypesMapping.java new file mode 100644 index 000000000000..d1ad8f537988 --- /dev/null +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerDataTypesMapping.java @@ -0,0 +1,104 @@ +package io.trino.plugin.spanner; + +import com.google.common.collect.ImmutableMap; +import io.trino.Session; +import io.trino.spi.type.BigintType; +import io.trino.testing.AbstractTestQueryFramework; +import io.trino.testing.QueryRunner; +import io.trino.testing.datatype.CreateAndInsertDataSetup; +import io.trino.testing.datatype.CreateAsSelectDataSetup; +import io.trino.testing.datatype.DataSetup; +import io.trino.testing.datatype.SqlDataTypeTest; +import io.trino.testing.sql.JdbcSqlExecutor; +import io.trino.testing.sql.TemporaryRelation; +import io.trino.testing.sql.TrinoSqlExecutor; +import io.trino.tpch.TpchTable; +import org.testng.annotations.Test; + +import java.util.Properties; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static io.trino.spi.type.BooleanType.BOOLEAN; + +public class TestSpannerDataTypesMapping + extends AbstractTestQueryFramework +{ + protected TestingSpannerInstance spannerInstance; + + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + spannerInstance = closeAfterClass(new TestingSpannerInstance()); + return SpannerQueryRunner.createSpannerQueryRunner( + spannerInstance, + ImmutableMap.of("http-server.http.port", "8080"), + ImmutableMap.of(), + TpchTable.getTables(), false); + } + + @Test + public void testBoolean() + { + SqlDataTypeTest.create() + .addRoundTrip("bool", "true", BOOLEAN) + .addRoundTrip("bool", "false", BOOLEAN) + .addRoundTrip("bool", "NULL", BOOLEAN, "CAST(NULL AS BOOLEAN)") + .addRoundTrip("int64", "1", BigintType.BIGINT, "CAST(1 as BIGINT)") + .execute(getQueryRunner(), spannerCreateAndInsert("test_boolean")) + .execute(getQueryRunner(), trinoCreateAsSelect("test_boolean")) + .execute(getQueryRunner(), trinoCreateAndInsert("test_boolean")); + } + + private DataSetup spannerCreateAndInsert(String tableNamePrefix) + { + JdbcSqlExecutor jdbcSqlExecutor = new JdbcSqlExecutor(spannerInstance.getJdbcUrl(), new Properties()); + return inputs -> { + + String primaryKey = String.format("col_%s", inputs.size() - 1); + + jdbcSqlExecutor.execute("CREATE TABLE %s (%s) PRIMARY KEY (%s)" + .formatted( + tableNamePrefix, + IntStream.range(0, inputs.size()) + .mapToObj(f -> String.format("col_%s %s", f, inputs.get(f).getDeclaredType().get())) + .collect(Collectors.joining(", ")), + primaryKey)); + return new TemporaryRelation() + { + @Override + public String getName() + { + return tableNamePrefix; + } + + @Override + public void close() + { + + } + }; + }; + } + + private DataSetup trinoCreateAsSelect(String tableNamePrefix) + { + return trinoCreateAsSelect(getSession(), tableNamePrefix); + } + + private DataSetup trinoCreateAsSelect(Session session, String tableNamePrefix) + { + return new CreateAsSelectDataSetup(new TrinoSqlExecutor(getQueryRunner(), session), tableNamePrefix); + } + + private DataSetup trinoCreateAndInsert(String tableNamePrefix) + { + return trinoCreateAndInsert(getSession(), tableNamePrefix); + } + + private DataSetup trinoCreateAndInsert(Session session, String tableNamePrefix) + { + return new CreateAndInsertDataSetup(new TrinoSqlExecutor(getQueryRunner(), session), tableNamePrefix); + } +} diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerPlugin.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerPlugin.java index a3cdd1f7ceaf..56a0c0ea299c 100644 --- a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerPlugin.java +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerPlugin.java @@ -1,22 +1,33 @@ package io.trino.plugin.spanner; import com.google.common.collect.ImmutableMap; -import io.trino.plugin.spanner.SpannerPlugin; import io.trino.spi.Plugin; import io.trino.spi.connector.ConnectorFactory; import io.trino.testing.TestingConnectorContext; -import org.testcontainers.utility.DockerImageName; import org.testng.annotations.Test; +import java.util.concurrent.ExecutionException; + import static com.google.common.collect.Iterables.getOnlyElement; public class TestSpannerPlugin { @Test public void testCreateConnector() + throws Exception { Plugin plugin = new SpannerPlugin(); ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories()); - factory.create("test", ImmutableMap.of("connection-url", "jdbc:cloudspanner://0.0.0.0:9010/projects/spanner-project/instances/spanner-instance/databases/spanner-database;autoConfigEmulator=true"), new TestingConnectorContext()).shutdown(); + TestingSpannerInstance instance = new TestingSpannerInstance(); + factory.create("test", ImmutableMap.of( + "spanner.credentials.file", "credentials.json", + "spanner.instanceId", instance.getInstanceId() + , "spanner.projectId", instance.getProjectId() + , "spanner.database", instance.getDatabaseId() + , "spanner.emulated", "true" + , "spanner.emulated.host", instance.getHost() + ), + new TestingConnectorContext()).shutdown(); + instance.close(); } } From 7097ff09a5c0785166dc8ac5d6870c3042791b37 Mon Sep 17 00:00:00 2001 From: taherk77 Date: Fri, 20 Oct 2023 13:07:48 +0530 Subject: [PATCH 8/8] Try fix test case for Spanner Datatype mapping --- .../plugin/spanner/SpannerTestTable.java | 44 ++++++++++++++ .../spanner/TestSpannerDataTypesMapping.java | 59 ++++++------------- 2 files changed, 61 insertions(+), 42 deletions(-) create mode 100644 plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerTestTable.java diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerTestTable.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerTestTable.java new file mode 100644 index 000000000000..6bd22c28b8d6 --- /dev/null +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/SpannerTestTable.java @@ -0,0 +1,44 @@ +package io.trino.plugin.spanner; + +import io.trino.testing.sql.SqlExecutor; +import io.trino.testing.sql.TestTable; + +import java.util.List; + +import static java.lang.String.format; + +public class SpannerTestTable extends TestTable +{ + + public SpannerTestTable(SqlExecutor sqlExecutor, String namePrefix, String tableDefinition) + { + super(sqlExecutor, namePrefix, tableDefinition); + } + + + @Override + public void createAndInsert(List rowsToInsert) + { + String[] fields = tableDefinition.split(","); + String pkField = fields[fields.length - 1].trim(); + String pkColumn = pkField.split(" ")[0]; + String create = "CREATE TABLE %s (%s) PRIMARY KEY (%s)" + .formatted( + this.name, + tableDefinition, + pkColumn); + sqlExecutor.execute(create); + + try { + for (String row : rowsToInsert) { + String sql = format("INSERT INTO %s VALUES (%s)", name, row); + sqlExecutor.execute(sql); + } + } + catch (Exception e) { + try (SpannerTestTable ignored = this) { + throw e; + } + } + } +} diff --git a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerDataTypesMapping.java b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerDataTypesMapping.java index d1ad8f537988..b9e69c0932d8 100644 --- a/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerDataTypesMapping.java +++ b/plugin/trino-spanner/src/test/java/io/trino/plugin/spanner/TestSpannerDataTypesMapping.java @@ -3,18 +3,22 @@ import com.google.common.collect.ImmutableMap; import io.trino.Session; import io.trino.spi.type.BigintType; +import io.trino.spi.type.IntegerType; import io.trino.testing.AbstractTestQueryFramework; import io.trino.testing.QueryRunner; +import io.trino.testing.datatype.ColumnSetup; import io.trino.testing.datatype.CreateAndInsertDataSetup; import io.trino.testing.datatype.CreateAsSelectDataSetup; import io.trino.testing.datatype.DataSetup; +import io.trino.testing.datatype.DataType; import io.trino.testing.datatype.SqlDataTypeTest; import io.trino.testing.sql.JdbcSqlExecutor; -import io.trino.testing.sql.TemporaryRelation; import io.trino.testing.sql.TrinoSqlExecutor; import io.trino.tpch.TpchTable; +import org.jetbrains.annotations.NotNull; import org.testng.annotations.Test; +import java.util.List; import java.util.Properties; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -45,41 +49,22 @@ public void testBoolean() .addRoundTrip("bool", "true", BOOLEAN) .addRoundTrip("bool", "false", BOOLEAN) .addRoundTrip("bool", "NULL", BOOLEAN, "CAST(NULL AS BOOLEAN)") - .addRoundTrip("int64", "1", BigintType.BIGINT, "CAST(1 as BIGINT)") - .execute(getQueryRunner(), spannerCreateAndInsert("test_boolean")) - .execute(getQueryRunner(), trinoCreateAsSelect("test_boolean")) - .execute(getQueryRunner(), trinoCreateAndInsert("test_boolean")); + .addRoundTrip("int64","1", BigintType.BIGINT,"CAST(1 AS BIGINT)") + .execute(getQueryRunner(), spannerCreateAndInsert("test_boolean")); } private DataSetup spannerCreateAndInsert(String tableNamePrefix) { - JdbcSqlExecutor jdbcSqlExecutor = new JdbcSqlExecutor(spannerInstance.getJdbcUrl(), new Properties()); - return inputs -> { - - String primaryKey = String.format("col_%s", inputs.size() - 1); - - jdbcSqlExecutor.execute("CREATE TABLE %s (%s) PRIMARY KEY (%s)" - .formatted( - tableNamePrefix, - IntStream.range(0, inputs.size()) - .mapToObj(f -> String.format("col_%s %s", f, inputs.get(f).getDeclaredType().get())) - .collect(Collectors.joining(", ")), - primaryKey)); - return new TemporaryRelation() - { - @Override - public String getName() - { - return tableNamePrefix; - } - - @Override - public void close() - { - - } - }; - }; + return inputs -> new SpannerTestTable(new JdbcSqlExecutor(spannerInstance.getJdbcUrl(), new Properties()) + ,tableNamePrefix, String.format("%s", getColumns(inputs))); + } + @NotNull + private String getColumns(List inputs) + { + return IntStream.range(0, inputs.size()) + .mapToObj(f -> String.format("col_%s %s", f, + inputs.get(f).getDeclaredType().get())) + .collect(Collectors.joining(", ")); } private DataSetup trinoCreateAsSelect(String tableNamePrefix) @@ -91,14 +76,4 @@ private DataSetup trinoCreateAsSelect(Session session, String tableNamePrefix) { return new CreateAsSelectDataSetup(new TrinoSqlExecutor(getQueryRunner(), session), tableNamePrefix); } - - private DataSetup trinoCreateAndInsert(String tableNamePrefix) - { - return trinoCreateAndInsert(getSession(), tableNamePrefix); - } - - private DataSetup trinoCreateAndInsert(Session session, String tableNamePrefix) - { - return new CreateAndInsertDataSetup(new TrinoSqlExecutor(getQueryRunner(), session), tableNamePrefix); - } }