test: add integration test for PostgreSQL declarative partitioned tables

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<Map<String, Object>> measurementsPostgreSQL =
+        postgresSQLResourceManager.runSQLQuery(
+            "SELECT id, city_id, logdate, peaktemp FROM measurements_range");
+    ImmutableList<Struct> measurementsSpanner =
+        spannerResourceManager.readTableRecords(
+            "measurements_range", "id", "city_id", "logdate", "peaktemp");
+    SpannerAsserts.assertThatStructs(measurementsSpanner)
+        .hasRecordsUnorderedCaseInsensitiveColumns(measurementsPostgreSQL);
+
+    List<Map<String, Object>> employeesPostgreSQL =
+        postgresSQLResourceManager.runSQLQuery("SELECT id, name, department FROM employees_list");
+    ImmutableList<Struct> employeesSpanner =
+        spannerResourceManager.readTableRecords("employees_list", "id", "name", "department");
+    SpannerAsserts.assertThatStructs(employeesSpanner)
+        .hasRecordsUnorderedCaseInsensitiveColumns(employeesPostgreSQL);
+
+    List<Map<String, Object>> ordersPostgreSQL =
+        postgresSQLResourceManager.runSQLQuery(
+            "SELECT order_id, customer_id, amount FROM orders_hash");
+    ImmutableList<Struct> ordersSpanner =
+        spannerResourceManager.readTableRecords("orders_hash", "order_id", "customer_id", "amount");
+    SpannerAsserts.assertThatStructs(ordersSpanner)
+        .hasRecordsUnorderedCaseInsensitiveColumns(ordersPostgreSQL);
+  }
+}

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);

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);

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)
+);