diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java index f3818b4005..1189e21b0b 100644 --- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java +++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java @@ -148,7 +148,7 @@ public class ConnectionHandler implements Runnable { private Connection spannerConnection; private DatabaseId databaseId; private WellKnownClient wellKnownClient = WellKnownClient.UNSPECIFIED; - private boolean hasDeterminedClientUsingQuery; + private int autoDetectCount; /** * List of PARSE messages that we received before auto-detecting the client. This list can be used @@ -967,9 +967,8 @@ public void setWellKnownClient(WellKnownClient wellKnownClient) { * executed. */ public void maybeDetermineWellKnownClient(Statement statement) { - if (!this.hasDeterminedClientUsingQuery) { - if (this.wellKnownClient == WellKnownClient.UNSPECIFIED - && getServer().getOptions().shouldAutoDetectClient()) { + if (this.wellKnownClient == WellKnownClient.UNSPECIFIED && this.autoDetectCount < 100) { + if (getServer().getOptions().shouldAutoDetectClient()) { setWellKnownClient( ClientAutoDetector.detectClient( skippedAutoDetectParseMessages, ImmutableList.of(statement))); @@ -985,8 +984,7 @@ && getServer().getOptions().shouldAutoDetectClient()) { } maybeSetApplicationName(); skippedAutoDetectParseMessages.clear(); - // Make sure that we only try to detect the client once. - this.hasDeterminedClientUsingQuery = true; + this.autoDetectCount++; } } @@ -997,7 +995,7 @@ && getServer().getOptions().shouldAutoDetectClient()) { * messages. */ public void maybeDetermineWellKnownClient(ParseMessage parseMessage) { - if (!this.hasDeterminedClientUsingQuery) { + if (this.wellKnownClient == WellKnownClient.UNSPECIFIED && this.autoDetectCount < 100) { // We skip up to 10 Parse messages before forcing the connection to detect a client (or not). // This is a safety measure against clients that might send a very large number of Parse // messages with client-side statements directly after connecting. This could in worst case @@ -1009,15 +1007,13 @@ public void maybeDetermineWellKnownClient(ParseMessage parseMessage) { && skippedAutoDetectParseMessages.size() < 10) { skippedAutoDetectParseMessages.add(parseMessage); } else { - if (this.wellKnownClient == WellKnownClient.UNSPECIFIED - && getServer().getOptions().shouldAutoDetectClient()) { + if (getServer().getOptions().shouldAutoDetectClient()) { setWellKnownClient( ClientAutoDetector.detectClient(skippedAutoDetectParseMessages, parseMessage)); } maybeSetApplicationName(); skippedAutoDetectParseMessages.clear(); - // Make sure that we only try to detect the client once. - this.hasDeterminedClientUsingQuery = true; + this.autoDetectCount++; } } } @@ -1052,7 +1048,7 @@ List getSkippedAutoDetectParseMessages() { @VisibleForTesting boolean isHasDeterminedClientUsingQuery() { - return this.hasDeterminedClientUsingQuery; + return this.wellKnownClient != WellKnownClient.UNSPECIFIED || this.autoDetectCount >= 10; } /** Status of a {@link ConnectionHandler} */ diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java index 2f80f2fa8f..f2b71f5f02 100644 --- a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java +++ b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java @@ -120,7 +120,7 @@ public static class Builder { Builder() {} - Builder setEnvironment(Map environment) { + public Builder setEnvironment(Map environment) { this.environment = Preconditions.checkNotNull(environment); return this; } diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/PgCatalog.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/PgCatalog.java index 4fae4114ce..3e7058d5eb 100644 --- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/PgCatalog.java +++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/PgCatalog.java @@ -672,7 +672,7 @@ public class PgAttribute implements PgCatalogTable { public static final String PG_ATTRIBUTE_CTE = "pg_attribute as (\n" - + "select '''\"' || table_schema || '\".\"' || table_name || '\"''' as attrelid,\n" + + "select (ABS(spanner.farm_fingerprint(table_schema || '.' || table_name)) % 2147483647) as attrelid,\n" + " column_name as attname,\n" + " case regexp_replace(c.spanner_type, '\\(.*\\)', '')\n" + " when 'boolean' then 16\n" @@ -712,7 +712,7 @@ public class PgAttribute implements PgCatalogTable { + " c.spanner_type\n" + "from information_schema.columns c\n" + "union all\n" - + "select '''\"' || i.table_schema || '\".\"' || i.table_name || '\".\"' || i.index_name || '\"''' as attrelid,\n" + + "select (ABS(spanner.farm_fingerprint(i.table_schema || '.' || i.table_name || '.' || i.index_name)) % 2147483647) as attrelid,\n" + " i.column_name as attname,\n" + " case regexp_replace(c.spanner_type, '\\(.*\\)', '')\n" + " when 'boolean' then 16\n" @@ -770,8 +770,8 @@ public class PgAttrdef implements PgCatalogTable { public static final String PG_ATTRDEF_CTE = "pg_attrdef as (\n" - + "select '''\"' || table_schema || '\".\"' || table_name || '\".\"' || column_name || '\"''' as oid,\n" - + " '''\"' || table_schema || '\".\"' || table_name || '\"''' as adrelid,\n" + + "select (ABS(spanner.farm_fingerprint(table_schema || '.' || table_name || '.' || column_name)) % 2147483647) as oid,\n" + + " (ABS(spanner.farm_fingerprint(table_schema || '.' || table_name)) % 2147483647) as adrelid,\n" + " ordinal_position as adnum,\n" + " coalesce(column_default, generation_expression) as adbin\n" + "from information_schema.columns c\n" @@ -788,7 +788,7 @@ public class PgConstraint implements PgCatalogTable { public static final String PG_CONSTRAINT_CTE = "pg_constraint as (\n" + "select\n" - + " '''\"' || tc.constraint_schema || '\".\"' || tc.constraint_name || '\"''' as oid,\n" + + " (ABS(spanner.farm_fingerprint(tc.constraint_schema || '.' || tc.constraint_name)) % 2147483647) as oid,\n" + " tc.constraint_name as conname, 2200 as connamespace,\n" + " case tc.constraint_type\n" + " when 'PRIMARY KEY' then 'p'\n" @@ -796,9 +796,9 @@ public class PgConstraint implements PgCatalogTable { + " when 'FOREIGN KEY' then 'f'\n" + " else ''\n" + " end as contype, false as condeferrable, false as condeferred, true as convalidated,\n" - + " '''\"' || tc.table_schema || '\".\"' || tc.table_name || '\"''' as conrelid,\n" + + " (ABS(spanner.farm_fingerprint(tc.table_schema || '.' || tc.table_name)) % 2147483647) as conrelid,\n" + " 0::bigint as contypid, '0'::varchar as conindid, '0'::varchar as conparentid,\n" - + " '''\"' || uc.table_schema || '\".\"' || uc.table_name || '\"''' as confrelid,\n" + + " (ABS(spanner.farm_fingerprint(uc.table_schema || '.' || uc.table_name)) % 2147483647) as confrelid,\n" + " case rc.update_rule\n" + " when 'CASCADE' then 'c'\n" + " when 'NO ACTION' then 'a'\n" diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java index 55b0e73158..8f951b87f2 100644 --- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java +++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java @@ -30,6 +30,7 @@ import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata; import com.google.cloud.spanner.pgadapter.statements.BackendConnection.ConnectionState; import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient; +import com.google.cloud.spanner.pgadapter.utils.QueryPartReplacer; import com.google.cloud.spanner.pgadapter.wireprotocol.BindMessage; import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage.ManuallyCreatedToken; import com.google.cloud.spanner.pgadapter.wireprotocol.DescribeMessage; @@ -117,11 +118,25 @@ public void execute() throws Exception { /** Replaces any known unsupported query (e.g. JDBC metadata queries). */ static Tuple replaceKnownUnsupportedQueries( WellKnownClient client, OptionsMetadata options, String sql) { + boolean replaced = false; if ((options.isReplaceJdbcMetadataQueries() || client == WellKnownClient.JDBC) && JdbcMetadataStatementHelper.isPotentialJdbcMetadataStatement(sql)) { - return Tuple.of(Boolean.TRUE, JdbcMetadataStatementHelper.replaceJdbcMetadataStatement(sql)); + sql = JdbcMetadataStatementHelper.replaceJdbcMetadataStatement(sql); + replaced = true; } - return Tuple.of(Boolean.FALSE, sql); + if (client != null) { + for (QueryPartReplacer replacer : client.getQueryPartReplacements()) { + Tuple result = replacer.replace(sql); + if (!result.x().equals(sql)) { + sql = result.x(); + replaced = true; + } + if (result.y() == QueryPartReplacer.ReplacementStatus.STOP) { + break; + } + } + } + return Tuple.of(replaced, sql); } /** diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java index 5ffd5237ac..6e45a53ca9 100644 --- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java +++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java @@ -394,6 +394,34 @@ public ImmutableList getDdlReplacements() { + ")")); } }, + ADBC { + private final Pattern adbcTypeQuery = + Pattern.compile("(?s).*pg_type.*WHERE.*typreceive.*typsend.*", Pattern.CASE_INSENSITIVE); + + @Override + boolean isClient(List orderedParameterKeys, Map parameters) { + return false; + } + + @Override + boolean isClient(List skippedParseMessages, ParseMessage parseMessage) { + return parseMessage.getSql() != null && adbcTypeQuery.matcher(parseMessage.getSql()).find(); + } + + @Override + boolean isClient(List skippedParseMessages, List statements) { + return !statements.isEmpty() && adbcTypeQuery.matcher(statements.get(0).getSql()).find(); + } + + @Override + public ImmutableList getQueryPartReplacements() { + return ImmutableList.of( + RegexQueryPartReplacer.replace( + Pattern.compile("typreceive\\s*!=\\s*0"), "typreceive != '0'::varchar"), + RegexQueryPartReplacer.replace( + Pattern.compile("typsend\\s*!=\\s*0"), "typsend != '0'::varchar")); + } + }, PGX { @Override boolean isClient(List orderedParameterKeys, Map parameters) { diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/Logging.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/Logging.java index 55f307bbac..3aee8c0562 100644 --- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/Logging.java +++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/Logging.java @@ -26,7 +26,7 @@ public enum Action { /** Format a log message by prepending the current thread name and method to the message. */ public static Supplier format(String method, Supplier message) { return () -> - String.format("[%s]: [%s] " + message.get(), Thread.currentThread().getName(), method); + String.format("[%s]: [%s] %s", Thread.currentThread().getName(), method, message.get()); } /** Create a log message with the current thread name, method and action. */ @@ -40,6 +40,6 @@ public static Supplier format(String method, Action action) { public static Supplier format(String method, Action action, Supplier message) { return () -> String.format( - "[%s]: [%s] [%s] " + message.get(), Thread.currentThread().getName(), method, action); + "[%s]: [%s] [%s] %s", Thread.currentThread().getName(), method, action, message.get()); } } diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java index efa5cd84ca..a40f7f1f35 100644 --- a/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java +++ b/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java @@ -652,7 +652,7 @@ public void testMaybeDetermineWellKnownClient_stopsSkippingParseMessagesAfter10M // UNSPECIFIED. connectionHandler.maybeDetermineWellKnownClient(parseMessage); - assertTrue(connectionHandler.isHasDeterminedClientUsingQuery()); + assertFalse(connectionHandler.isHasDeterminedClientUsingQuery()); assertTrue(connectionHandler.getSkippedAutoDetectParseMessages().isEmpty()); } @@ -665,7 +665,10 @@ public void testBuildConnectionUrl() { == null); OptionsMetadata options = - OptionsMetadata.newBuilder().setCredentials(NoCredentials.getInstance()).build(); + OptionsMetadata.newBuilder() + .setCredentials(NoCredentials.getInstance()) + .setEnvironment(ImmutableMap.of()) + .build(); // Check that the dialect is included in the connection URL. This is required to support the // 'autoConfigEmulator' property. assertEquals( @@ -694,11 +697,16 @@ public void testBuildConnectionUrl() { .setProject("test-project") .setInstance("test-instance") .setDatabase("test-database") + .setEnvironment(ImmutableMap.of()) .build(), buildProperties(ImmutableMap.of()), ImmutableMap.of())); // Enable the autoConfigEmulator flag through the options builder. - OptionsMetadata emulatorOptions = OptionsMetadata.newBuilder().autoConfigureEmulator().build(); + OptionsMetadata emulatorOptions = + OptionsMetadata.newBuilder() + .autoConfigureEmulator() + .setEnvironment(ImmutableMap.of()) + .build(); assertEquals( "cloudspanner:/projects/my-project/instances/my-instance/databases/my-database;userAgent=pg-adapter;autoConfigEmulator=true;defaultSequenceKind=bit_reversed_positive;dialect=postgresql", buildConnectionURL( @@ -770,7 +778,10 @@ public void testBuildConnectionUrl() { "true", () -> { OptionsMetadata optionsWithSystemProperty = - OptionsMetadata.newBuilder().setCredentials(NoCredentials.getInstance()).build(); + OptionsMetadata.newBuilder() + .setCredentials(NoCredentials.getInstance()) + .setEnvironment(ImmutableMap.of()) + .build(); assertEquals( "cloudspanner:/projects/my-project/instances/my-instance/databases/my-database;userAgent=pg-adapter;defaultSequenceKind=bit_reversed_positive;useVirtualThreads=true;dialect=postgresql", buildConnectionURL( @@ -784,7 +795,10 @@ public void testBuildConnectionUrl() { "true", () -> { OptionsMetadata optionsWithSystemProperty = - OptionsMetadata.newBuilder().setCredentials(NoCredentials.getInstance()).build(); + OptionsMetadata.newBuilder() + .setCredentials(NoCredentials.getInstance()) + .setEnvironment(ImmutableMap.of()) + .build(); assertEquals( "cloudspanner:/projects/my-project/instances/my-instance/databases/my-database;userAgent=pg-adapter;defaultSequenceKind=bit_reversed_positive;useVirtualGrpcTransportThreads=true;dialect=postgresql", buildConnectionURL( diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadataTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadataTest.java index 4fb55a6b75..f5c3eae393 100644 --- a/src/test/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadataTest.java +++ b/src/test/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadataTest.java @@ -132,12 +132,19 @@ public void testBuildConnectionUrlWithFullPath() { assertEquals( "cloudspanner:/projects/test-project/instances/test-instance/databases/test-database;userAgent=pg-adapter;credentials=credentials.json", - new OptionsMetadata(new String[] {"-c", "credentials.json"}) + new OptionsMetadata( + ImmutableMap.of(), + System.getProperty("os.name"), + DEFAULT_STARTUP_TIMEOUT, + new String[] {"-c", "credentials.json"}) .buildConnectionURL( "projects/test-project/instances/test-instance/databases/test-database")); assertEquals( "cloudspanner:/projects/test-project/instances/test-instance/databases/test-database;userAgent=pg-adapter;credentials=credentials.json", new OptionsMetadata( + ImmutableMap.of(), + System.getProperty("os.name"), + DEFAULT_STARTUP_TIMEOUT, new String[] { "-p", "test-project", "-i", "test-instance", "-c", "credentials.json" }) @@ -171,7 +178,11 @@ public void testBuildConnectionUrlWithDefaultProjectId() { == null); OptionsMetadata useDefaultProjectIdOptions = - new OptionsMetadata(new String[] {"-i", "test-instance", "-c", "credentials.json"}) { + new OptionsMetadata( + ImmutableMap.of(), + System.getProperty("os.name"), + DEFAULT_STARTUP_TIMEOUT, + new String[] {"-i", "test-instance", "-c", "credentials.json"}) { @Override String getDefaultProjectId() { return "custom-test-project"; @@ -181,7 +192,11 @@ String getDefaultProjectId() { "cloudspanner:/projects/custom-test-project/instances/test-instance/databases/test-database;userAgent=pg-adapter;credentials=credentials.json", useDefaultProjectIdOptions.buildConnectionURL("test-database")); OptionsMetadata noProjectIdOptions = - new OptionsMetadata(new String[] {"-i", "test-instance", "-c", "credentials.json"}) { + new OptionsMetadata( + ImmutableMap.of(), + System.getProperty("os.name"), + DEFAULT_STARTUP_TIMEOUT, + new String[] {"-i", "test-instance", "-c", "credentials.json"}) { @Override String getDefaultProjectId() { return null; @@ -202,7 +217,11 @@ public void testBuildConnectionUrlWithDefaultCredentials() { == null); OptionsMetadata useDefaultCredentials = - new OptionsMetadata(new String[] {"-p", "test-project", "-i", "test-instance"}) { + new OptionsMetadata( + ImmutableMap.of(), + System.getProperty("os.name"), + DEFAULT_STARTUP_TIMEOUT, + new String[] {"-p", "test-project", "-i", "test-instance"}) { @Override void tryGetDefaultCredentials() {} }; @@ -210,7 +229,11 @@ void tryGetDefaultCredentials() {} "cloudspanner:/projects/test-project/instances/test-instance/databases/test-database;userAgent=pg-adapter", useDefaultCredentials.buildConnectionURL("test-database")); OptionsMetadata noDefaultCredentialsOptions = - new OptionsMetadata(new String[] {"-p", "test-project", "-i", "test-instance"}) { + new OptionsMetadata( + ImmutableMap.of(), + System.getProperty("os.name"), + DEFAULT_STARTUP_TIMEOUT, + new String[] {"-p", "test-project", "-i", "test-instance"}) { @Override void tryGetDefaultCredentials() throws IOException { throw new IOException("test exception"); diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/python/adbc/AdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/python/adbc/AdbcMockServerTest.java new file mode 100644 index 0000000000..5b18a15dcb --- /dev/null +++ b/src/test/java/com/google/cloud/spanner/pgadapter/python/adbc/AdbcMockServerTest.java @@ -0,0 +1,337 @@ +// Copyright 2026 Google LLC +// +// 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 com.google.cloud.spanner.pgadapter.python.adbc; + +import static com.google.cloud.spanner.pgadapter.python.PythonTestUtil.run; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; + +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.pgadapter.AbstractMockServerTest; +import com.google.cloud.spanner.pgadapter.python.PythonTest; +import com.google.cloud.spanner.pgadapter.python.PythonTestUtil; +import com.google.cloud.spanner.pgadapter.statements.PgCatalog; +import com.google.cloud.spanner.pgadapter.statements.PgCatalog.PgAttribute; +import com.google.cloud.spanner.pgadapter.statements.PgCatalog.PgNamespace; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.ListValue; +import com.google.protobuf.Value; +import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.ResultSet; +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.StructType; +import com.google.spanner.v1.StructType.Field; +import com.google.spanner.v1.Type; +import com.google.spanner.v1.TypeCode; +import java.io.File; +import java.util.List; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; + +@RunWith(Parameterized.class) +@Category(PythonTest.class) +public class AdbcMockServerTest extends AbstractMockServerTest { + static final String DIRECTORY_NAME = "./src/test/python/adbc_driver_postgresql"; + + @Parameter public String host; + + @Parameters(name = "host = {0}") + public static List data() { + return ImmutableList.of(new Object[] {"localhost"}, new Object[] {"/tmp"}); + } + + @BeforeClass + public static void createVirtualEnv() throws Exception { + PythonTestUtil.createVirtualEnv(DIRECTORY_NAME); + } + + String createConnectionString() { + if ("localhost".equals(host)) { + return String.format("postgresql://localhost:%d/d?sslmode=disable", pgServer.getLocalPort()); + } else { + // For unix socket, let's see if ADBC supports it. Standard URI format for unix socket is not + // perfectly standardized. + // Let's try to pass the path as a query parameter or similar if the driver supports it, or + // just use host=... if it accepts it. + // If it fails, we will find out. + return String.format( + "postgresql://localhost:%d/d?sslmode=disable&host=%s", pgServer.getLocalPort(), host); + } + } + + String execute(String method) throws Exception { + return execute(method, createConnectionString()); + } + + static String execute(String method, String connectionString) throws Exception { + File directory = new File(DIRECTORY_NAME); + return run( + new String[] { + directory.getAbsolutePath() + "/venv/bin/python3", + "adbc_tests.py", + method, + connectionString + }, + DIRECTORY_NAME); + } + + private ResultSet createResultSet() { + return ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.INT64).build()) + .setName("") + .build()) + .build()) + .build()) + .addRows( + ListValue.newBuilder() + .addValues(Value.newBuilder().setStringValue(String.valueOf(1)).build()) + .build()) + .build(); + } + + @Before + public void setupStartupQueries() { + mockStartupQueries(); + } + + private void mockStartupQueries() { + String sql = + "with " + + PgAttribute.PG_ATTRIBUTE_CTE + + "\n\n" + + "SELECT\n" + + " attrelid,\n" + + " attname,\n" + + " atttypid\n" + + "FROM\n" + + " pg_attribute\n" + + "ORDER BY\n" + + " attrelid, attnum\n"; + mockSpanner.putStatementResult( + StatementResult.query(Statement.of(sql), createEmptyAttributeResultSet())); + + String typeSql = + "with " + + PgNamespace.PG_NAMESPACE_CTE + + ",\n" + + PgCatalog.PG_TYPE_CTE_EMULATED + + "\n" + + "SELECT oid, typname, typreceive, typbasetype, typrelid, typarray FROM pg_type WHERE (typreceive != 0 OR typsend != 0) AND typtype != 'r' AND typreceive::TEXT != 'array_recv'"; + mockSpanner.putStatementResult( + StatementResult.query(Statement.of(typeSql), createEmptyTypeResultSet())); + } + + private ResultSet createEmptyTypeResultSet() { + return ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.INT64).build()) + .setName("oid") + .build()) + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.STRING).build()) + .setName("typname") + .build()) + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.STRING).build()) + .setName("typreceive") + .build()) + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.INT64).build()) + .setName("typbasetype") + .build()) + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.STRING).build()) + .setName("typrelid") + .build()) + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.INT64).build()) + .setName("typarray") + .build()) + .build()) + .build()) + .build(); + } + + private ResultSet createEmptyAttributeResultSet() { + return ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.STRING).build()) + .setName("attrelid") + .build()) + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.STRING).build()) + .setName("attname") + .build()) + .addFields( + Field.newBuilder() + .setType(Type.newBuilder().setCode(TypeCode.INT64).build()) + .setName("atttypid") + .build()) + .build()) + .build()) + .build(); + } + + @Test + public void testSelect1() throws Exception { + String sql = "SELECT 1"; + + mockSpanner.putStatementResult(StatementResult.query(Statement.of(sql), createResultSet())); + + String actualOutput = execute("select1"); + // Depending on ADBC output, it might be (1,) or something similar. + // Let's verify what it is. For psycopg3 it was (1,). + // Let's assume it is similar. + String expectedOutput = "(b'\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x01',)\n"; + assertEquals(expectedOutput, actualOutput); + + ExecuteSqlRequest request = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).stream() + .filter(r -> r.getSql().equals(sql)) + .findFirst() + .orElse(null); + assertNotNull("No ExecuteSqlRequest found for " + sql, request); + assertEquals(sql, request.getSql()); + } + + @Test + public void testSelectString() throws Exception { + String sql = "SELECT 'foo'"; + + com.google.spanner.v1.ResultSet resultSet = + com.google.spanner.v1.ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("C") + .setType(Type.newBuilder().setCode(TypeCode.STRING).build()) + .build()) + .build()) + .build()) + .addRows( + ListValue.newBuilder() + .addValues(Value.newBuilder().setStringValue("foo").build()) + .build()) + .build(); + + mockSpanner.putStatementResult(StatementResult.query(Statement.of(sql), resultSet)); + + String actualOutput = execute("select_string"); + // We expect it to be a tuple with bytes or string. Let's see. + // Assuming binary protocol, it might be (b'foo',). Or just ('foo',). + // Let's print actual output in the failure message if it fails. + String expectedOutput = "(b'foo',)\n"; // Assuming bytes as a starting point. + assertEquals(expectedOutput, actualOutput); + } + + @Test + public void testSelectBoolean() throws Exception { + String sql = "SELECT true"; + + com.google.spanner.v1.ResultSet resultSet = + com.google.spanner.v1.ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("C") + .setType(Type.newBuilder().setCode(TypeCode.BOOL).build()) + .build()) + .build()) + .build()) + .addRows( + ListValue.newBuilder() + .addValues(Value.newBuilder().setBoolValue(true).build()) + .build()) + .build(); + + mockSpanner.putStatementResult(StatementResult.query(Statement.of(sql), resultSet)); + + String actualOutput = execute("select_boolean"); + // We expect it to be a tuple with bytes or boolean. Let's see. + // Assuming binary protocol, it might be (b'\x01',). Or just (True,). + // Let's print actual output in the failure message if it fails. + String expectedOutput = "(b'\\x01',)\n"; + // Let's guess (True,) first, if it fails we will know the exact format. + assertEquals(expectedOutput, actualOutput); + } + + @Test + public void testSelectTimestamp() throws Exception { + String sql = "SELECT '2020-01-01T00:00:00Z'::timestamp"; + + com.google.spanner.v1.ResultSet resultSet = + com.google.spanner.v1.ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("C") + .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP).build()) + .build()) + .build()) + .build()) + .addRows( + ListValue.newBuilder() + .addValues(Value.newBuilder().setStringValue("2020-01-01T00:00:00Z").build()) + .build()) + .build(); + + mockSpanner.putStatementResult(StatementResult.query(Statement.of(sql), resultSet)); + + String actualOutput = execute("select_timestamp"); + // We expect it to be a tuple with bytes or datetime. Let's see. + // Assuming binary protocol, it might be bytes. + // Let's print actual output in the failure message if it fails. + String expectedOutput = "(b'\\x00\\x02>\\x07\\x86\\xc2`\\x00',)\n"; + assertEquals(expectedOutput, actualOutput); + } +} diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/python/adbc/ITAdbcTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/python/adbc/ITAdbcTest.java new file mode 100644 index 0000000000..3ade107b31 --- /dev/null +++ b/src/test/java/com/google/cloud/spanner/pgadapter/python/adbc/ITAdbcTest.java @@ -0,0 +1,93 @@ +// Copyright 2026 Google LLC +// +// 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 com.google.cloud.spanner.pgadapter.python.adbc; + +import static com.google.cloud.spanner.pgadapter.PgAdapterTestEnv.getOnlyAllTypesDdl; +import static org.junit.Assert.assertEquals; + +import com.google.cloud.spanner.Database; +import com.google.cloud.spanner.pgadapter.IntegrationTest; +import com.google.cloud.spanner.pgadapter.PgAdapterTestEnv; +import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata; +import com.google.cloud.spanner.pgadapter.python.PythonTestUtil; +import java.util.Collections; +import java.util.logging.Handler; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; + +@Category(IntegrationTest.class) +@RunWith(Parameterized.class) +public class ITAdbcTest implements IntegrationTest { + private static final PgAdapterTestEnv testEnv = new PgAdapterTestEnv(); + private static Database database; + + @Parameter public String host; + + @Parameters(name = "host = {0}") + public static Object[] data() { + OptionsMetadata options = new OptionsMetadata(new String[] {"-p p", "-i i"}); + return options.isDomainSocketEnabled() + ? new Object[] {"localhost", "/tmp"} + : new Object[] {"localhost"}; + } + + @BeforeClass + public static void setup() throws Exception { + Logger.getLogger("com.google.cloud.spanner.pgadapter").setLevel(Level.FINEST); + for (Handler handler : Logger.getLogger("").getHandlers()) { + handler.setLevel(Level.FINEST); + } + + // We reuse the virtual environment from AdbcMockServerTest + PythonTestUtil.createVirtualEnv(AdbcMockServerTest.DIRECTORY_NAME); + testEnv.setUp(); + database = testEnv.createDatabase(getDdlStatements()); + testEnv.startPGAdapterServerWithDefaultDatabase(database.getId(), Collections.emptyList()); + } + + @AfterClass + public static void teardown() { + testEnv.stopPGAdapterServer(); + testEnv.cleanUp(); + } + + private static Iterable getDdlStatements() { + return getOnlyAllTypesDdl(); + } + + private String createConnectionString() { + return String.format( + "host=%s port=%d dbname=d sslmode=disable", host, testEnv.getServer().getLocalPort()); + } + + private String execute(String method) throws Exception { + return AdbcMockServerTest.execute(method, createConnectionString()); + } + + @Test + public void testSelect1() throws Exception { + String actualOutput = execute("select1"); + String expectedOutput = "(1,)\n"; + assertEquals(expectedOutput, actualOutput); + } +} diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetectorTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetectorTest.java index ca6613e354..cde62986ef 100644 --- a/src/test/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetectorTest.java +++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetectorTest.java @@ -484,4 +484,27 @@ public void testSetting() { WellKnownClient.DEFAULT_UNSPECIFIED.set(true); } } + + @Test + public void testAdbc() { + // Test detection by query + assertEquals( + WellKnownClient.ADBC, + ClientAutoDetector.detectClient( + ImmutableList.of(), + ImmutableList.of( + Statement.of( + "WITH pg_type AS (SELECT 1) SELECT oid, typname, typreceive, typbasetype, typrelid, typarray FROM pg_type WHERE (typreceive != 0 OR typsend != 0) AND typtype != 'r' AND typreceive::TEXT != 'array_recv'")))); + + // Test replacement + String sql = + "WITH pg_type AS (SELECT 1) SELECT oid, typname, typreceive, typbasetype, typrelid, typarray FROM pg_type WHERE (typreceive != 0 OR typsend != 0) AND typtype != 'r' AND typreceive::TEXT != 'array_recv'"; + String replaced = sql; + for (QueryPartReplacer replacer : WellKnownClient.ADBC.getQueryPartReplacements()) { + replaced = replacer.replace(replaced).x(); + } + assertEquals( + "WITH pg_type AS (SELECT 1) SELECT oid, typname, typreceive, typbasetype, typrelid, typarray FROM pg_type WHERE (typreceive != '0'::varchar OR typsend != '0'::varchar) AND typtype != 'r' AND typreceive::TEXT != 'array_recv'", + replaced); + } } diff --git a/src/test/python/adbc_driver_postgresql/adbc_tests.py b/src/test/python/adbc_driver_postgresql/adbc_tests.py new file mode 100644 index 0000000000..41bdd32039 --- /dev/null +++ b/src/test/python/adbc_driver_postgresql/adbc_tests.py @@ -0,0 +1,75 @@ +""" Copyright 2026 Google LLC + + 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 argparse +import sys +import adbc_driver_postgresql.dbapi as dbapi + + +def select1(conn_string: str): + # ADBC might expect a URI or a set of key-value pairs. + # Let's try to pass the connection string as-is or convert it to a URI if needed. + # If it is a URI (starts with postgresql://), use it. + # Otherwise, assume it is key-value pairs and pass it as a URI if it's just host/port. + + # For now, let's just try to connect with the string as-is. + # If it's "host=localhost port=port...", we might need to convert it to a URI for ADBC if it doesn't support KV pairs in connect(). + # Most ADBC drivers expect a URI. + uri = conn_string + if not uri.startswith("postgresql://"): + # Simple conversion for testing, assuming the format is "host=X port=Y dbname=Z ..." + # This is a bit hacky, but let's see if we can make it work. + # Or we can just pass the URI from Java. + pass + + with dbapi.connect(uri) as conn: + with conn.cursor() as cur: + cur.execute("SELECT 1") + print(cur.fetchone()) + + +def select_string(conn_string: str): + with dbapi.connect(conn_string) as conn: + with conn.cursor() as cur: + cur.execute("SELECT 'foo'") + print(cur.fetchone()) + + +def select_boolean(conn_string: str): + with dbapi.connect(conn_string) as conn: + with conn.cursor() as cur: + cur.execute("SELECT true") + print(cur.fetchone()) + + +def select_timestamp(conn_string: str): + with dbapi.connect(conn_string) as conn: + with conn.cursor() as cur: + cur.execute("SELECT '2020-01-01T00:00:00Z'::timestamp") + print(cur.fetchone()) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run ADBC tests.") + parser.add_argument("method", type=str, help="Test method to run") + parser.add_argument("conn_string", type=str, help="Connection string for PGAdapter") + args = parser.parse_args() + + method = globals().get(args.method) + if method: + method(args.conn_string) + else: + print(f"Unknown method: {args.method}") + sys.exit(1) diff --git a/src/test/python/adbc_driver_postgresql/requirements.txt b/src/test/python/adbc_driver_postgresql/requirements.txt new file mode 100644 index 0000000000..37173def78 --- /dev/null +++ b/src/test/python/adbc_driver_postgresql/requirements.txt @@ -0,0 +1,3 @@ +--extra-index-url https://pypi.org/simple/ +adbc-driver-postgresql +pyarrow