diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PostgreSQLPartitionedTablesIT.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PostgreSQLPartitionedTablesIT.java new file mode 100644 index 0000000000..e1340a6f0f --- /dev/null +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PostgreSQLPartitionedTablesIT.java @@ -0,0 +1,145 @@ +/* + * Copyright (C) 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.teleport.v2.templates; + +import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult; + +import com.google.cloud.spanner.Struct; +import com.google.cloud.teleport.metadata.SkipDirectRunnerTest; +import com.google.cloud.teleport.metadata.TemplateIntegrationTest; +import com.google.common.collect.ImmutableList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.it.common.PipelineLauncher; +import org.apache.beam.it.common.PipelineOperator; +import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.gcp.spanner.SpannerResourceManager; +import org.apache.beam.it.gcp.spanner.matchers.SpannerAsserts; +import org.apache.beam.it.jdbc.PostgresResourceManager; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * An integration test for {@link SourceDbToSpanner} Flex template which tests a migration from a + * PostgreSQL database containing partitioned tables. + */ +@Category({TemplateIntegrationTest.class, SkipDirectRunnerTest.class}) +@TemplateIntegrationTest(SourceDbToSpanner.class) +@RunWith(JUnit4.class) +public class PostgreSQLPartitionedTablesIT extends SourceDbToSpannerITBase { + + private static boolean initialized = false; + private static PipelineLauncher.LaunchInfo jobInfo; + + public static PostgresResourceManager postgresSQLResourceManager; + public static SpannerResourceManager gsqlSpannerResourceManager; + public static SpannerResourceManager pgDialectSpannerResourceManager; + + private static final String POSTGRESQL_DDL_RESOURCE = "PartitionedTablesIT/postgresql-schema.sql"; + private static final String SPANNER_GSQL_DDL_RESOURCE = + "PartitionedTablesIT/spanner-gsql-schema.sql"; + private static final String SPANNER_PG_DDL_RESOURCE = "PartitionedTablesIT/spanner-pg-schema.sql"; + + @Before + public void setUp() throws Exception { + synchronized (PostgreSQLPartitionedTablesIT.class) { + if (!initialized) { + postgresSQLResourceManager = setUpPostgreSQLResourceManager(); + gsqlSpannerResourceManager = setUpSpannerResourceManager(); + pgDialectSpannerResourceManager = setUpPGDialectSpannerResourceManager(); + + loadSQLFileResource(postgresSQLResourceManager, POSTGRESQL_DDL_RESOURCE); + + initialized = true; + } + } + } + + @AfterClass + public static void cleanUp() { + ResourceManagerUtils.cleanResources( + postgresSQLResourceManager, gsqlSpannerResourceManager, pgDialectSpannerResourceManager); + } + + @Test + public void testPostgreSQLPartitionedTablesGoogleSQLDialect() throws Exception { + createSpannerDDL(gsqlSpannerResourceManager, SPANNER_GSQL_DDL_RESOURCE); + jobInfo = + launchDataflowJob( + getClass().getSimpleName() + "GSQL", + null, + "measurements_range,employees_list,orders_hash", + postgresSQLResourceManager, + gsqlSpannerResourceManager, + new HashMap<>(), + null); + + PipelineOperator.Result result = pipelineOperator().waitUntilDone(createConfig(jobInfo)); + assertThatResult(result).isLaunchFinished(); + + verifyData(gsqlSpannerResourceManager); + } + + @Test + public void testPostgreSQLPartitionedTablesPostgreSQLDialect() throws Exception { + createSpannerDDL(pgDialectSpannerResourceManager, SPANNER_PG_DDL_RESOURCE); + jobInfo = + launchDataflowJob( + getClass().getSimpleName() + "PG", + null, + "measurements_range,employees_list,orders_hash", + postgresSQLResourceManager, + pgDialectSpannerResourceManager, + new HashMap<>(), + null); + + PipelineOperator.Result result = pipelineOperator().waitUntilDone(createConfig(jobInfo)); + assertThatResult(result).isLaunchFinished(); + + verifyData(pgDialectSpannerResourceManager); + } + + private void verifyData(SpannerResourceManager spannerResourceManager) { + List> measurementsPostgreSQL = + postgresSQLResourceManager.runSQLQuery( + "SELECT id, city_id, logdate, peaktemp FROM measurements_range"); + ImmutableList measurementsSpanner = + spannerResourceManager.readTableRecords( + "measurements_range", "id", "city_id", "logdate", "peaktemp"); + SpannerAsserts.assertThatStructs(measurementsSpanner) + .hasRecordsUnorderedCaseInsensitiveColumns(measurementsPostgreSQL); + + List> employeesPostgreSQL = + postgresSQLResourceManager.runSQLQuery("SELECT id, name, department FROM employees_list"); + ImmutableList employeesSpanner = + spannerResourceManager.readTableRecords("employees_list", "id", "name", "department"); + SpannerAsserts.assertThatStructs(employeesSpanner) + .hasRecordsUnorderedCaseInsensitiveColumns(employeesPostgreSQL); + + List> ordersPostgreSQL = + postgresSQLResourceManager.runSQLQuery( + "SELECT order_id, customer_id, amount FROM orders_hash"); + ImmutableList ordersSpanner = + spannerResourceManager.readTableRecords("orders_hash", "order_id", "customer_id", "amount"); + SpannerAsserts.assertThatStructs(ordersSpanner) + .hasRecordsUnorderedCaseInsensitiveColumns(ordersPostgreSQL); + } +} diff --git a/v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/postgresql-schema.sql b/v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/postgresql-schema.sql new file mode 100644 index 0000000000..8910266a36 --- /dev/null +++ b/v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/postgresql-schema.sql @@ -0,0 +1,50 @@ +CREATE TABLE measurements_range ( + id INT, + city_id INT NOT NULL, + logdate DATE NOT NULL, + peaktemp INT, + PRIMARY KEY (id, logdate) +) PARTITION BY RANGE (logdate); + +CREATE TABLE measurements_range_y2006m02 PARTITION OF measurements_range + FOR VALUES FROM ('2006-02-01') TO ('2006-03-01'); + +CREATE TABLE measurements_range_y2006m03 PARTITION OF measurements_range + FOR VALUES FROM ('2006-03-01') TO ('2006-04-01'); + +CREATE TABLE measurements_range_default PARTITION OF measurements_range DEFAULT; + +INSERT INTO measurements_range (id, city_id, logdate, peaktemp) VALUES (1, 1, '2006-02-15', 33); +INSERT INTO measurements_range (id, city_id, logdate, peaktemp) VALUES (2, 2, '2006-03-15', 35); +INSERT INTO measurements_range (id, city_id, logdate, peaktemp) VALUES (3, 3, '2006-05-10', 40); +INSERT INTO measurements_range_y2006m02 (id, city_id, logdate, peaktemp) VALUES (4, 4, '2006-02-20', 30); + +CREATE TABLE employees_list ( + id INT, + name VARCHAR(50), + department VARCHAR(50), + PRIMARY KEY (id, department) +) PARTITION BY LIST (department); + +CREATE TABLE employees_list_engineering PARTITION OF employees_list FOR VALUES IN ('Engineering'); +CREATE TABLE employees_list_sales PARTITION OF employees_list FOR VALUES IN ('Sales'); +CREATE TABLE employees_list_default PARTITION OF employees_list DEFAULT; + +INSERT INTO employees_list (id, name, department) VALUES (1, 'Alice', 'Engineering'); +INSERT INTO employees_list (id, name, department) VALUES (2, 'Bob', 'Sales'); +INSERT INTO employees_list (id, name, department) VALUES (3, 'Charlie', 'Marketing'); + +CREATE TABLE orders_hash ( + order_id INT, + customer_id INT, + amount INT, + PRIMARY KEY (order_id, customer_id) +) PARTITION BY HASH (customer_id); + +CREATE TABLE orders_hash_p0 PARTITION OF orders_hash FOR VALUES WITH (MODULUS 3, REMAINDER 0); +CREATE TABLE orders_hash_p1 PARTITION OF orders_hash FOR VALUES WITH (MODULUS 3, REMAINDER 1); +CREATE TABLE orders_hash_p2 PARTITION OF orders_hash FOR VALUES WITH (MODULUS 3, REMAINDER 2); + +INSERT INTO orders_hash (order_id, customer_id, amount) VALUES (1, 101, 500); +INSERT INTO orders_hash (order_id, customer_id, amount) VALUES (2, 102, 600); +INSERT INTO orders_hash (order_id, customer_id, amount) VALUES (3, 103, 700); diff --git a/v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/spanner-gsql-schema.sql b/v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/spanner-gsql-schema.sql new file mode 100644 index 0000000000..a5cc022181 --- /dev/null +++ b/v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/spanner-gsql-schema.sql @@ -0,0 +1,18 @@ +CREATE TABLE measurements_range ( + id INT64, + city_id INT64, + logdate DATE, + peaktemp INT64 +) PRIMARY KEY(id, logdate); + +CREATE TABLE employees_list ( + id INT64, + name STRING(50), + department STRING(50) +) PRIMARY KEY(id, department); + +CREATE TABLE orders_hash ( + order_id INT64, + customer_id INT64, + amount INT64 +) PRIMARY KEY(order_id, customer_id); diff --git a/v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/spanner-pg-schema.sql b/v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/spanner-pg-schema.sql new file mode 100644 index 0000000000..b9d40c4f15 --- /dev/null +++ b/v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/spanner-pg-schema.sql @@ -0,0 +1,21 @@ +CREATE TABLE measurements_range ( + id bigint, + city_id bigint, + logdate date, + peaktemp bigint, + PRIMARY KEY (id, logdate) +); + +CREATE TABLE employees_list ( + id bigint, + name character varying(50), + department character varying(50), + PRIMARY KEY (id, department) +); + +CREATE TABLE orders_hash ( + order_id bigint, + customer_id bigint, + amount bigint, + PRIMARY KEY (order_id, customer_id) +);