From 24b396f6449b73665b96ffe52b3eba46f19d4fc4 Mon Sep 17 00:00:00 2001 From: Mairbek Khadikov Date: Wed, 18 Oct 2017 15:18:55 -0700 Subject: [PATCH 1/3] Reading spanner schema transform --- .../sdk/io/gcp/spanner/ReadSpannerSchema.java | 94 ++++++++++ .../sdk/io/gcp/spanner/SpannerSchema.java | 166 ++++++++++++++++++ 2 files changed, 260 insertions(+) create mode 100644 sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java create mode 100644 sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java new file mode 100644 index 000000000000..ad2fe11ada2e --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.sdk.io.gcp.spanner; + +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.ReadOnlyTransaction; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.Statement; +import org.apache.beam.sdk.transforms.DoFn; + +/** + * This transform reads Cloud Spanner 'information_schema.*' tables to build the + * {@link SpannerSchema}. + */ +class ReadSpannerSchema extends DoFn { + + private final SpannerConfig config; + + private transient SpannerAccessor spannerAccessor; + + public ReadSpannerSchema(SpannerConfig config) { + this.config = config; + } + + @Setup + public void setup() throws Exception { + spannerAccessor = config.connectToSpanner(); + } + + @Teardown + public void teardown() throws Exception { + spannerAccessor.close(); + } + + @ProcessElement + public void processElement(ProcessContext c) throws Exception { + SpannerSchema.Builder builder = SpannerSchema.builder(); + DatabaseClient databaseClient = spannerAccessor.getDatabaseClient(); + try (ReadOnlyTransaction tx = + databaseClient.readOnlyTransaction()) { + ResultSet resultSet = readTableInfo(tx); + + while (resultSet.next()) { + String tableName = resultSet.getString(0); + String columnName = resultSet.getString(1); + String type = resultSet.getString(2); + + builder.addColumn(tableName, columnName, type); + } + + resultSet = readPrimaryKeyInfo(tx); + while (resultSet.next()) { + String tableName = resultSet.getString(0); + String columnName = resultSet.getString(1); + String ordering = resultSet.getString(2); + + builder.addKeyPart(tableName, columnName, ordering.toUpperCase().equals("DESC")); + } + } + c.output(builder.build()); + } + + private ResultSet readTableInfo(ReadOnlyTransaction tx) { + return tx.executeQuery(Statement.of( + "SELECT c.table_name, c.column_name, c.spanner_type" + + " FROM information_schema.columns as c" + + " WHERE where c.table_catalog = '' AND c.table_schema = ''" + + " ORDER BY c.table_name, c.ordinal_position")); + } + + private ResultSet readPrimaryKeyInfo(ReadOnlyTransaction tx) { + return tx.executeQuery(Statement + .of("SELECT t.table_name, t.column_name, t.column_ordering" + + " FROM information_schema.index_columns AS t " + + " WHERE t.index_name = 'PRIMARY_KEY' AND t.table_catalog = ''" + + " AND t.table_schema = ''" + + " ORDER BY t.table_name, t.ordinal_position")); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java new file mode 100644 index 000000000000..8ba7d25252e3 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java @@ -0,0 +1,166 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.sdk.io.gcp.spanner; + +import com.google.auto.value.AutoValue; +import com.google.cloud.spanner.Type; +import com.google.common.base.Preconditions; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Encapsulates Cloud Spanner Schema. + */ +class SpannerSchema implements Serializable { + private final List tables; + private final Map tableIndex; + private final Map> columns; + private final Map> keyParts; + + public static Builder builder() { + return new Builder(); + } + + /** + * Builder for {@link SpannerSchema}. + */ + static class Builder { + private final Map> columns = new HashMap<>(); + private final Map> keyParts = new HashMap<>(); + + public void addColumn(String table, String name, String type) { + addColumn(table, Column.create(name.toLowerCase(), type)); + } + + private void addColumn(String table, Column column) { + List list = columns.get(table); + if (list == null) { + list = new ArrayList<>(); + columns.put(table.toLowerCase(), list); + } + list.add(column); + } + + public void addKeyPart(String table, String column, boolean desc) { + List list = keyParts.get(table); + if (list == null) { + list = new ArrayList<>(); + keyParts.put(table.toLowerCase(), list); + } + list.add(KeyPart.create(column.toLowerCase(), desc)); + } + + public SpannerSchema build() { + return new SpannerSchema(columns, keyParts); + } + } + + private SpannerSchema(Map> columns, Map> keyParts) { + this.columns = columns; + this.keyParts = keyParts; + tables = new ArrayList<>(columns.keySet()); + tableIndex = new HashMap<>(tables.size()); + Collections.sort(tables); + for (int i = 0; i < tables.size(); i++) { + tableIndex.put(tables.get(i), i); + } + } + + public int getTableIndex(String tableName) { + Integer result = tableIndex.get(tableName); + Preconditions.checkArgument(result != null, "Table %s not found", tableName); + return result; + } + + public String getTableName(int index) { + Preconditions.checkArgument(index < tables.size(), "Invalid table index %d", index); + return tables.get(index); + } + + public List getColumns(String table) { + return columns.get(table); + } + + public List getKeyParts(String table) { + return keyParts.get(table); + } + + @AutoValue + abstract static class KeyPart implements Serializable { + static KeyPart create(String field, boolean desc) { + return new AutoValue_SpannerSchema_KeyPart(field, desc); + } + + abstract String getField(); + + abstract boolean isDesc(); + } + + @AutoValue + abstract static class Column implements Serializable { + + static Column create(String name, Type type) { + return new AutoValue_SpannerSchema_Column(name, type); + } + + static Column create(String name, String spannerType) { + return create(name, parseSpannerType(spannerType)); + } + + public abstract String getName(); + + public abstract Type getType(); + + private static Type parseSpannerType(String spannerType) { + spannerType = spannerType.toUpperCase(); + if (spannerType.equals("BOOL")) { + return Type.bool(); + } + if (spannerType.equals("INT64")) { + return Type.int64(); + } + if (spannerType.equals("FLOAT64")) { + return Type.float64(); + } + if (spannerType.startsWith("STRING")) { + return Type.string(); + } + if (spannerType.startsWith("BYTES")) { + return Type.bytes(); + } + if (spannerType.equals("TIMESTAMP")) { + return Type.timestamp(); + } + if (spannerType.equals("DATE")) { + return Type.date(); + } + + if (spannerType.startsWith("ARRAY")) { + // Substring "ARRAY" + String spannerArrayType = spannerType.substring(6, spannerType.length() - 1); + Type itemType = parseSpannerType(spannerArrayType); + return Type.array(itemType); + } + throw new IllegalArgumentException("Unknown spanner type " + spannerType); + } + } +} From 6f0be30fdf56f20579f09698ad8e626c3a99454b Mon Sep 17 00:00:00 2001 From: Mairbek Khadikov Date: Thu, 19 Oct 2017 17:17:59 -0700 Subject: [PATCH 2/3] Tests --- .../sdk/io/gcp/spanner/SpannerSchema.java | 9 +- .../io/gcp/spanner/ReadSpannerSchemaTest.java | 117 ++++++++++++++++++ .../sdk/io/gcp/spanner/SpannerSchemaTest.java | 44 +++++++ 3 files changed, 167 insertions(+), 3 deletions(-) create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchemaTest.java create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java index 8ba7d25252e3..d395f5a8d93d 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java @@ -47,26 +47,29 @@ static class Builder { private final Map> columns = new HashMap<>(); private final Map> keyParts = new HashMap<>(); - public void addColumn(String table, String name, String type) { + public Builder addColumn(String table, String name, String type) { addColumn(table, Column.create(name.toLowerCase(), type)); + return this; } - private void addColumn(String table, Column column) { + private Builder addColumn(String table, Column column) { List list = columns.get(table); if (list == null) { list = new ArrayList<>(); columns.put(table.toLowerCase(), list); } list.add(column); + return this; } - public void addKeyPart(String table, String column, boolean desc) { + public Builder addKeyPart(String table, String column, boolean desc) { List list = keyParts.get(table); if (list == null) { list = new ArrayList<>(); keyParts.put(table.toLowerCase(), list); } list.add(KeyPart.create(column.toLowerCase(), desc)); + return this; } public SpannerSchema build() { diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchemaTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchemaTest.java new file mode 100644 index 000000000000..621ff50b3064 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchemaTest.java @@ -0,0 +1,117 @@ +package org.apache.beam.sdk.io.gcp.spanner; + +import static org.hamcrest.Matchers.contains; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThat; +import static org.mockito.Matchers.argThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.cloud.spanner.ReadOnlyTransaction; +import com.google.cloud.spanner.ResultSets; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.Struct; +import com.google.cloud.spanner.Type; +import com.google.cloud.spanner.Value; +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.transforms.DoFnTester; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.mockito.ArgumentMatcher; + +/** + * A test of {@link ReadSpannerSchemaTest}. + */ +public class ReadSpannerSchemaTest { + + @Rule + public final transient ExpectedException thrown = ExpectedException.none(); + + private FakeServiceFactory serviceFactory; + private ReadOnlyTransaction mockTx; + + + private static Struct columnMetadata(String tableName, String columnName, String type) { + return Struct.newBuilder().add("table_name", Value.string(tableName)) + .add("column_name", Value.string(columnName)).add("spanner_type", Value.string(type)) + .build(); + } + + private static Struct pkMetadata(String tableName, String columnName, String ordering) { + return Struct.newBuilder().add("table_name", Value.string(tableName)) + .add("column_name", Value.string(columnName)).add("column_ordering", Value.string(ordering)) + .build(); + } + + private void prepareColumnMetadata(ReadOnlyTransaction tx, List rows) { + Type type = Type.struct(Type.StructField.of("table_name", Type.string()), + Type.StructField.of("column_name", Type.string()), + Type.StructField.of("spanner_type", Type.string())); + when(tx.executeQuery(argThat(new ArgumentMatcher() { + + @Override public boolean matches(Object argument) { + if (!(argument instanceof Statement)) { + return false; + } + Statement st = (Statement) argument; + return st.getSql().contains("information_schema.columns"); + } + }))).thenReturn(ResultSets.forRows(type, rows)); + } + + private void preparePkMetadata(ReadOnlyTransaction tx, List rows) { + Type type = Type.struct(Type.StructField.of("table_name", Type.string()), + Type.StructField.of("column_name", Type.string()), + Type.StructField.of("column_ordering", Type.string())); + when(tx.executeQuery(argThat(new ArgumentMatcher() { + + @Override public boolean matches(Object argument) { + if (!(argument instanceof Statement)) { + return false; + } + Statement st = (Statement) argument; + return st.getSql().contains("information_schema.index_columns"); + } + }))).thenReturn(ResultSets.forRows(type, rows)); + } + + @Before + @SuppressWarnings("unchecked") + public void setUp() throws Exception { + serviceFactory = new FakeServiceFactory(); + mockTx = mock(ReadOnlyTransaction.class); + } + + @Test + public void simple() throws Exception { + // Simplest schema: a table with int64 key + ReadOnlyTransaction tx = mock(ReadOnlyTransaction.class); + when(serviceFactory.mockDatabaseClient().readOnlyTransaction()).thenReturn(tx); + + preparePkMetadata(tx, Arrays.asList(pkMetadata("test", "key", "ASC"))); + prepareColumnMetadata(tx, Arrays.asList(columnMetadata("test", "key", "INT64"))); + + SpannerConfig config = SpannerConfig.create().withProjectId("test-project") + .withInstanceId("test-instance").withDatabaseId("test-database") + .withServiceFactory(serviceFactory); + + DoFnTester tester = DoFnTester.of(new ReadSpannerSchema(config)); + List schemas = tester.processBundle(Arrays.asList((Void) null)); + + assertEquals(1, schemas.size()); + + SpannerSchema schema = schemas.get(0); + + assertEquals(0, schema.getTableIndex("test")); + + SpannerSchema.Column column = SpannerSchema.Column.create("key", Type.int64()); + SpannerSchema.KeyPart keyPart = SpannerSchema.KeyPart.create("key", false); + + assertThat(schema.getColumns("test"), contains(column)); + assertThat(schema.getKeyParts("test"), contains(keyPart)); + } + +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java new file mode 100644 index 000000000000..11d07f9fb62e --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java @@ -0,0 +1,44 @@ +package org.apache.beam.sdk.io.gcp.spanner; + +import static org.junit.Assert.assertEquals; + +import org.junit.Test; + +/** + * A test of {@link SpannerSchema}. + */ +public class SpannerSchemaTest { + + @Test + public void testSingleTable() throws Exception { + SpannerSchema schema = SpannerSchema.builder() + .addColumn("test", "pk", "STRING(48)") + .addKeyPart("test", "pk", false) + .addColumn("test", "maxKey", "STRING(MAX)").build(); + + assertEquals(0, schema.getTableIndex("test")); + assertEquals(2, schema.getColumns("test").size()); + assertEquals(1, schema.getKeyParts("test").size()); + } + + @Test + public void testTwoTables() throws Exception { + SpannerSchema schema = SpannerSchema.builder() + .addColumn("test", "pk", "STRING(48)") + .addKeyPart("test", "pk", false) + .addColumn("test", "maxKey", "STRING(MAX)") + + .addColumn("other", "pk", "INT64") + .addKeyPart("other", "pk", true) + .addColumn("other", "maxKey", "STRING(MAX)") + + .build(); + + assertEquals(1, schema.getTableIndex("test")); + assertEquals(2, schema.getColumns("test").size()); + assertEquals(1, schema.getKeyParts("test").size()); + + assertEquals(2, schema.getColumns("other").size()); + assertEquals(1, schema.getKeyParts("other").size()); + } +} From 5223774ea74497167662e0b32cfa4bd1c77b5c60 Mon Sep 17 00:00:00 2001 From: Mairbek Khadikov Date: Fri, 27 Oct 2017 12:22:59 -0700 Subject: [PATCH 3/3] Addressing comments --- .../sdk/io/gcp/spanner/ReadSpannerSchema.java | 2 +- .../sdk/io/gcp/spanner/SpannerSchema.java | 43 +++++-------------- .../io/gcp/spanner/ReadSpannerSchemaTest.java | 2 +- .../sdk/io/gcp/spanner/SpannerSchemaTest.java | 5 ++- 4 files changed, 16 insertions(+), 36 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java index ad2fe11ada2e..e2ade6885fbd 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java @@ -24,7 +24,7 @@ import org.apache.beam.sdk.transforms.DoFn; /** - * This transform reads Cloud Spanner 'information_schema.*' tables to build the + * This {@link DoFn} reads Cloud Spanner 'information_schema.*' tables to build the * {@link SpannerSchema}. */ class ReadSpannerSchema extends DoFn { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java index d395f5a8d93d..931267ce4d93 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java @@ -20,6 +20,7 @@ import com.google.auto.value.AutoValue; import com.google.cloud.spanner.Type; import com.google.common.base.Preconditions; +import com.google.common.collect.ArrayListMultimap; import java.io.Serializable; import java.util.ArrayList; import java.util.Collections; @@ -32,9 +33,8 @@ */ class SpannerSchema implements Serializable { private final List tables; - private final Map tableIndex; - private final Map> columns; - private final Map> keyParts; + private final ArrayListMultimap columns; + private final ArrayListMultimap keyParts; public static Builder builder() { return new Builder(); @@ -44,8 +44,8 @@ public static Builder builder() { * Builder for {@link SpannerSchema}. */ static class Builder { - private final Map> columns = new HashMap<>(); - private final Map> keyParts = new HashMap<>(); + private final ArrayListMultimap columns = ArrayListMultimap.create(); + private final ArrayListMultimap keyParts = ArrayListMultimap.create(); public Builder addColumn(String table, String name, String type) { addColumn(table, Column.create(name.toLowerCase(), type)); @@ -53,22 +53,12 @@ public Builder addColumn(String table, String name, String type) { } private Builder addColumn(String table, Column column) { - List list = columns.get(table); - if (list == null) { - list = new ArrayList<>(); - columns.put(table.toLowerCase(), list); - } - list.add(column); + columns.put(table.toLowerCase(), column); return this; } public Builder addKeyPart(String table, String column, boolean desc) { - List list = keyParts.get(table); - if (list == null) { - list = new ArrayList<>(); - keyParts.put(table.toLowerCase(), list); - } - list.add(KeyPart.create(column.toLowerCase(), desc)); + keyParts.put(table, KeyPart.create(column.toLowerCase(), desc)); return this; } @@ -77,26 +67,15 @@ public SpannerSchema build() { } } - private SpannerSchema(Map> columns, Map> keyParts) { + private SpannerSchema(ArrayListMultimap columns, + ArrayListMultimap keyParts) { this.columns = columns; this.keyParts = keyParts; tables = new ArrayList<>(columns.keySet()); - tableIndex = new HashMap<>(tables.size()); - Collections.sort(tables); - for (int i = 0; i < tables.size(); i++) { - tableIndex.put(tables.get(i), i); - } - } - - public int getTableIndex(String tableName) { - Integer result = tableIndex.get(tableName); - Preconditions.checkArgument(result != null, "Table %s not found", tableName); - return result; } - public String getTableName(int index) { - Preconditions.checkArgument(index < tables.size(), "Invalid table index %d", index); - return tables.get(index); + public List getTables() { + return tables; } public List getColumns(String table) { diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchemaTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchemaTest.java index 621ff50b3064..dc0487c837fb 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchemaTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchemaTest.java @@ -105,7 +105,7 @@ public void simple() throws Exception { SpannerSchema schema = schemas.get(0); - assertEquals(0, schema.getTableIndex("test")); + assertEquals(1, schema.getTables().size()); SpannerSchema.Column column = SpannerSchema.Column.create("key", Type.int64()); SpannerSchema.KeyPart keyPart = SpannerSchema.KeyPart.create("key", false); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java index 11d07f9fb62e..799b000d87be 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java @@ -2,6 +2,7 @@ import static org.junit.Assert.assertEquals; +import java.util.Arrays; import org.junit.Test; /** @@ -16,7 +17,7 @@ public void testSingleTable() throws Exception { .addKeyPart("test", "pk", false) .addColumn("test", "maxKey", "STRING(MAX)").build(); - assertEquals(0, schema.getTableIndex("test")); + assertEquals(1, schema.getTables().size()); assertEquals(2, schema.getColumns("test").size()); assertEquals(1, schema.getKeyParts("test").size()); } @@ -34,7 +35,7 @@ public void testTwoTables() throws Exception { .build(); - assertEquals(1, schema.getTableIndex("test")); + assertEquals(2, schema.getTables().size()); assertEquals(2, schema.getColumns("test").size()); assertEquals(1, schema.getKeyParts("test").size());