Skip to content

Commit 82fa91b

Browse files
jsuhani-2026jsuhani-2026
committed
test: add integration test for PostgreSQL partitioned tables
1 parent 06d6601 commit 82fa91b

4 files changed

Lines changed: 236 additions & 0 deletions

File tree

‎v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PostgreSQLPartitionedTablesIT.java‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
/*
2+
* Copyright (C) 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License"); you may not
5+
* use this file except in compliance with the License. You may obtain a copy of
6+
* the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12+
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13+
* License for the specific language governing permissions and limitations under
14+
* the License.
15+
*/
16+
package com.google.cloud.teleport.v2.templates;
17+
18+
import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult;
19+
20+
import com.google.cloud.spanner.Struct;
21+
import com.google.cloud.teleport.metadata.SkipDirectRunnerTest;
22+
import com.google.cloud.teleport.metadata.TemplateIntegrationTest;
23+
import com.google.common.collect.ImmutableList;
24+
import java.util.HashMap;
25+
import java.util.List;
26+
import java.util.Map;
27+
import org.apache.beam.it.common.PipelineLauncher;
28+
import org.apache.beam.it.common.PipelineOperator;
29+
import org.apache.beam.it.common.utils.ResourceManagerUtils;
30+
import org.apache.beam.it.gcp.spanner.SpannerResourceManager;
31+
import org.apache.beam.it.gcp.spanner.matchers.SpannerAsserts;
32+
import org.apache.beam.it.jdbc.PostgresResourceManager;
33+
import org.junit.AfterClass;
34+
import org.junit.Before;
35+
import org.junit.Test;
36+
import org.junit.experimental.categories.Category;
37+
import org.junit.runner.RunWith;
38+
import org.junit.runners.JUnit4;
39+
40+
/**
41+
* An integration test for {@link SourceDbToSpanner} Flex template which tests a migration from a
42+
* PostgreSQL database containing partitioned tables.
43+
*/
44+
@Category({TemplateIntegrationTest.class, SkipDirectRunnerTest.class})
45+
@TemplateIntegrationTest(SourceDbToSpanner.class)
46+
@RunWith(JUnit4.class)
47+
public class PostgreSQLPartitionedTablesIT extends SourceDbToSpannerITBase {
48+
49+
private static boolean initialized = false;
50+
private static PipelineLauncher.LaunchInfo jobInfo;
51+
52+
public static PostgresResourceManager postgresSQLResourceManager;
53+
public static SpannerResourceManager gsqlSpannerResourceManager;
54+
public static SpannerResourceManager pgDialectSpannerResourceManager;
55+
56+
private static final String POSTGRESQL_DDL_RESOURCE = "PartitionedTablesIT/postgresql-schema.sql";
57+
private static final String SPANNER_GSQL_DDL_RESOURCE =
58+
"PartitionedTablesIT/spanner-gsql-schema.sql";
59+
private static final String SPANNER_PG_DDL_RESOURCE = "PartitionedTablesIT/spanner-pg-schema.sql";
60+
61+
@Before
62+
public void setUp() throws Exception {
63+
synchronized (PostgreSQLPartitionedTablesIT.class) {
64+
if (!initialized) {
65+
postgresSQLResourceManager = setUpPostgreSQLResourceManager();
66+
gsqlSpannerResourceManager = setUpSpannerResourceManager();
67+
pgDialectSpannerResourceManager = setUpPGDialectSpannerResourceManager();
68+
69+
loadSQLFileResource(postgresSQLResourceManager, POSTGRESQL_DDL_RESOURCE);
70+
71+
initialized = true;
72+
}
73+
}
74+
}
75+
76+
@AfterClass
77+
public static void cleanUp() {
78+
ResourceManagerUtils.cleanResources(
79+
postgresSQLResourceManager, gsqlSpannerResourceManager, pgDialectSpannerResourceManager);
80+
}
81+
82+
@Test
83+
public void testPostgreSQLPartitionedTablesGoogleSQLDialect() throws Exception {
84+
createSpannerDDL(gsqlSpannerResourceManager, SPANNER_GSQL_DDL_RESOURCE);
85+
jobInfo =
86+
launchDataflowJob(
87+
getClass().getSimpleName() + "GSQL",
88+
null,
89+
"measurements_range,employees_list,orders_hash",
90+
postgresSQLResourceManager,
91+
gsqlSpannerResourceManager,
92+
new HashMap<>(),
93+
null);
94+
95+
PipelineOperator.Result result = pipelineOperator().waitUntilDone(createConfig(jobInfo));
96+
assertThatResult(result).isLaunchFinished();
97+
98+
verifyData(gsqlSpannerResourceManager);
99+
}
100+
101+
@Test
102+
public void testPostgreSQLPartitionedTablesPostgreSQLDialect() throws Exception {
103+
createSpannerDDL(pgDialectSpannerResourceManager, SPANNER_PG_DDL_RESOURCE);
104+
jobInfo =
105+
launchDataflowJob(
106+
getClass().getSimpleName() + "PG",
107+
null,
108+
"measurements_range,employees_list,orders_hash",
109+
postgresSQLResourceManager,
110+
pgDialectSpannerResourceManager,
111+
new HashMap<>(),
112+
null);
113+
114+
PipelineOperator.Result result = pipelineOperator().waitUntilDone(createConfig(jobInfo));
115+
assertThatResult(result).isLaunchFinished();
116+
117+
verifyData(pgDialectSpannerResourceManager);
118+
}
119+
120+
private void verifyData(SpannerResourceManager spannerResourceManager) {
121+
List<Map<String, Object>> measurementsPostgreSQL =
122+
postgresSQLResourceManager.runSQLQuery(
123+
"SELECT id, city_id, logdate, peaktemp FROM measurements_range");
124+
ImmutableList<Struct> measurementsSpanner =
125+
spannerResourceManager.readTableRecords(
126+
"measurements_range", "id", "city_id", "logdate", "peaktemp");
127+
SpannerAsserts.assertThatStructs(measurementsSpanner)
128+
.hasRecordsUnorderedCaseInsensitiveColumns(measurementsPostgreSQL);
129+
130+
List<Map<String, Object>> employeesPostgreSQL =
131+
postgresSQLResourceManager.runSQLQuery("SELECT id, name, department FROM employees_list");
132+
ImmutableList<Struct> employeesSpanner =
133+
spannerResourceManager.readTableRecords(
134+
"employees_list", "id", "name", "department");
135+
SpannerAsserts.assertThatStructs(employeesSpanner)
136+
.hasRecordsUnorderedCaseInsensitiveColumns(employeesPostgreSQL);
137+
138+
List<Map<String, Object>> ordersPostgreSQL =
139+
postgresSQLResourceManager.runSQLQuery(
140+
"SELECT order_id, customer_id, amount FROM orders_hash");
141+
ImmutableList<Struct> ordersSpanner =
142+
spannerResourceManager.readTableRecords(
143+
"orders_hash", "order_id", "customer_id", "amount");
144+
SpannerAsserts.assertThatStructs(ordersSpanner)
145+
.hasRecordsUnorderedCaseInsensitiveColumns(ordersPostgreSQL);
146+
}
147+
}

‎v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/postgresql-schema.sql‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
CREATE TABLE measurements_range (
2+
id INT,
3+
city_id INT NOT NULL,
4+
logdate DATE NOT NULL,
5+
peaktemp INT,
6+
PRIMARY KEY (id, logdate)
7+
) PARTITION BY RANGE (logdate);
8+
9+
CREATE TABLE measurements_range_y2006m02 PARTITION OF measurements_range
10+
FOR VALUES FROM ('2006-02-01') TO ('2006-03-01');
11+
12+
CREATE TABLE measurements_range_y2006m03 PARTITION OF measurements_range
13+
FOR VALUES FROM ('2006-03-01') TO ('2006-04-01');
14+
15+
CREATE TABLE measurements_range_default PARTITION OF measurements_range DEFAULT;
16+
17+
INSERT INTO measurements_range (id, city_id, logdate, peaktemp) VALUES (1, 1, '2006-02-15', 33);
18+
INSERT INTO measurements_range (id, city_id, logdate, peaktemp) VALUES (2, 2, '2006-03-15', 35);
19+
INSERT INTO measurements_range (id, city_id, logdate, peaktemp) VALUES (3, 3, '2006-05-10', 40);
20+
INSERT INTO measurements_range_y2006m02 (id, city_id, logdate, peaktemp) VALUES (4, 4, '2006-02-20', 30);
21+
22+
CREATE TABLE employees_list (
23+
id INT,
24+
name VARCHAR(50),
25+
department VARCHAR(50),
26+
PRIMARY KEY (id, department)
27+
) PARTITION BY LIST (department);
28+
29+
CREATE TABLE employees_list_engineering PARTITION OF employees_list FOR VALUES IN ('Engineering');
30+
CREATE TABLE employees_list_sales PARTITION OF employees_list FOR VALUES IN ('Sales');
31+
CREATE TABLE employees_list_default PARTITION OF employees_list DEFAULT;
32+
33+
INSERT INTO employees_list (id, name, department) VALUES (1, 'Alice', 'Engineering');
34+
INSERT INTO employees_list (id, name, department) VALUES (2, 'Bob', 'Sales');
35+
INSERT INTO employees_list (id, name, department) VALUES (3, 'Charlie', 'Marketing');
36+
37+
CREATE TABLE orders_hash (
38+
order_id INT,
39+
customer_id INT,
40+
amount INT,
41+
PRIMARY KEY (order_id, customer_id)
42+
) PARTITION BY HASH (customer_id);
43+
44+
CREATE TABLE orders_hash_p0 PARTITION OF orders_hash FOR VALUES WITH (MODULUS 3, REMAINDER 0);
45+
CREATE TABLE orders_hash_p1 PARTITION OF orders_hash FOR VALUES WITH (MODULUS 3, REMAINDER 1);
46+
CREATE TABLE orders_hash_p2 PARTITION OF orders_hash FOR VALUES WITH (MODULUS 3, REMAINDER 2);
47+
48+
INSERT INTO orders_hash (order_id, customer_id, amount) VALUES (1, 101, 500);
49+
INSERT INTO orders_hash (order_id, customer_id, amount) VALUES (2, 102, 600);
50+
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‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
CREATE TABLE measurements_range (
2+
id INT64,
3+
city_id INT64,
4+
logdate DATE,
5+
peaktemp INT64
6+
) PRIMARY KEY(id, logdate);
7+
8+
CREATE TABLE employees_list (
9+
id INT64,
10+
name STRING(50),
11+
department STRING(50)
12+
) PRIMARY KEY(id, department);
13+
14+
CREATE TABLE orders_hash (
15+
order_id INT64,
16+
customer_id INT64,
17+
amount INT64
18+
) PRIMARY KEY(order_id, customer_id);

‎v2/sourcedb-to-spanner/src/test/resources/PartitionedTablesIT/spanner-pg-schema.sql‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
CREATE TABLE measurements_range (
2+
id bigint,
3+
city_id bigint,
4+
logdate date,
5+
peaktemp bigint,
6+
PRIMARY KEY (id, logdate)
7+
);
8+
9+
CREATE TABLE employees_list (
10+
id bigint,
11+
name character varying(50),
12+
department character varying(50),
13+
PRIMARY KEY (id, department)
14+
);
15+
16+
CREATE TABLE orders_hash (
17+
order_id bigint,
18+
customer_id bigint,
19+
amount bigint,
20+
PRIMARY KEY (order_id, customer_id)
21+
);

0 commit comments

Comments
 (0)