Skip to content

Commit 14b097b

Browse files
shunpingshunping
committed
Add test for postgres schema transform translation.
1 parent 5c7b817 commit 14b097b

4 files changed

Lines changed: 238 additions & 2 deletions

File tree

‎sdks/java/io/jdbc/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ dependencies {
4343
testImplementation project(path: ":sdks:java:managed")
4444
testImplementation project(path: ":sdks:java:testing:test-utils")
4545
testImplementation library.java.junit
46+
testImplementation library.java.mockito_inline
4647
testImplementation library.java.slf4j_api
4748
testImplementation library.java.postgres
4849

‎sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/providers/ReadFromPostgresSchemaTransformProvider.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,12 +81,13 @@ protected String jdbcType() {
8181
.setConnectionInitSql(Collections.emptyList())
8282
.setDisableAutoCommit(true)
8383
.build();
84-
return super.from(configuration);
84+
return new PostgresReadSchemaTransform(configuration);
8585
}
8686

8787
public static class PostgresReadSchemaTransform extends JdbcReadSchemaTransform {
8888
public PostgresReadSchemaTransform(JdbcReadSchemaTransformConfiguration config) {
8989
super(config, POSTGRES);
90+
config.validate(POSTGRES);
9091
}
9192
}
9293
}

‎sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/providers/WriteToPostgresSchemaTransformProvider.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,12 +71,13 @@ protected String jdbcType() {
7171

7272
// Override "connectionInitSql" for postgres
7373
configuration = configuration.toBuilder().setConnectionInitSql(Collections.emptyList()).build();
74-
return super.from(configuration);
74+
return new PostgresWriteSchemaTransform(configuration);
7575
}
7676

7777
public static class PostgresWriteSchemaTransform extends JdbcWriteSchemaTransform {
7878
public PostgresWriteSchemaTransform(JdbcWriteSchemaTransformConfiguration config) {
7979
super(config, POSTGRES);
80+
config.validate(POSTGRES);
8081
}
8182
}
8283
}

‎sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/providers/PostgresSchemaTransformTranslationTest.java‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.sdk.io.jdbc.providers;
19+
20+
import static org.apache.beam.model.pipeline.v1.ExternalTransforms.ExpansionMethods.Enum.SCHEMA_TRANSFORM;
21+
import static org.apache.beam.sdk.io.jdbc.providers.PostgresSchemaTransformTranslation.PostgresReadSchemaTransformTranslator;
22+
import static org.apache.beam.sdk.io.jdbc.providers.PostgresSchemaTransformTranslation.PostgresWriteSchemaTransformTranslator;
23+
import static org.apache.beam.sdk.io.jdbc.providers.ReadFromPostgresSchemaTransformProvider.PostgresReadSchemaTransform;
24+
import static org.apache.beam.sdk.io.jdbc.providers.WriteToPostgresSchemaTransformProvider.PostgresWriteSchemaTransform;
25+
import static org.junit.Assert.assertEquals;
26+
27+
import java.io.IOException;
28+
import java.util.Collections;
29+
import java.util.List;
30+
import java.util.stream.Collectors;
31+
import org.apache.beam.model.pipeline.v1.ExternalTransforms.SchemaTransformPayload;
32+
import org.apache.beam.model.pipeline.v1.RunnerApi;
33+
import org.apache.beam.sdk.Pipeline;
34+
import org.apache.beam.sdk.coders.RowCoder;
35+
import org.apache.beam.sdk.io.jdbc.JdbcIO;
36+
import org.apache.beam.sdk.options.PipelineOptionsFactory;
37+
import org.apache.beam.sdk.schemas.Schema;
38+
import org.apache.beam.sdk.schemas.SchemaTranslation;
39+
import org.apache.beam.sdk.transforms.Create;
40+
import org.apache.beam.sdk.util.construction.BeamUrns;
41+
import org.apache.beam.sdk.util.construction.PipelineTranslation;
42+
import org.apache.beam.sdk.values.PCollection;
43+
import org.apache.beam.sdk.values.PCollectionRowTuple;
44+
import org.apache.beam.sdk.values.Row;
45+
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException;
46+
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
47+
import org.junit.ClassRule;
48+
import org.junit.Rule;
49+
import org.junit.Test;
50+
import org.junit.rules.ExpectedException;
51+
import org.junit.rules.TemporaryFolder;
52+
import org.mockito.MockedStatic;
53+
import org.mockito.Mockito;
54+
55+
public class PostgresSchemaTransformTranslationTest {
56+
@ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder();
57+
58+
@Rule public transient ExpectedException thrown = ExpectedException.none();
59+
60+
static final WriteToPostgresSchemaTransformProvider WRITE_PROVIDER =
61+
new WriteToPostgresSchemaTransformProvider();
62+
static final ReadFromPostgresSchemaTransformProvider READ_PROVIDER =
63+
new ReadFromPostgresSchemaTransformProvider();
64+
65+
static final Row READ_CONFIG =
66+
Row.withSchema(READ_PROVIDER.configurationSchema())
67+
.withFieldValue("jdbc_url", "jdbc:postgresql://host:port/database")
68+
.withFieldValue("location", "test_table")
69+
.withFieldValue("connection_properties", "some_property")
70+
.withFieldValue("connection_init_sql", ImmutableList.<String>builder().build())
71+
.withFieldValue("driver_class_name", null)
72+
.withFieldValue("driver_jars", null)
73+
.withFieldValue("disable_auto_commit", true)
74+
.withFieldValue("fetch_size", 10)
75+
.withFieldValue("num_partitions", 5)
76+
.withFieldValue("output_parallelization", true)
77+
.withFieldValue("partition_column", "col")
78+
.withFieldValue("read_query", null)
79+
.withFieldValue("username", "my_user")
80+
.withFieldValue("password", "my_pass")
81+
.build();
82+
83+
static final Row WRITE_CONFIG =
84+
Row.withSchema(WRITE_PROVIDER.configurationSchema())
85+
.withFieldValue("jdbc_url", "jdbc:postgresql://host:port/database")
86+
.withFieldValue("location", "test_table")
87+
.withFieldValue("autosharding", true)
88+
.withFieldValue("connection_init_sql", ImmutableList.<String>builder().build())
89+
.withFieldValue("connection_properties", "some_property")
90+
.withFieldValue("driver_class_name", null)
91+
.withFieldValue("driver_jars", null)
92+
.withFieldValue("batch_size", 100L)
93+
.withFieldValue("username", "my_user")
94+
.withFieldValue("password", "my_pass")
95+
.withFieldValue("write_statement", null)
96+
.build();
97+
98+
@Test
99+
public void testRecreateWriteTransformFromRow() {
100+
PostgresWriteSchemaTransform writeTransform =
101+
(PostgresWriteSchemaTransform) WRITE_PROVIDER.from(WRITE_CONFIG);
102+
103+
PostgresWriteSchemaTransformTranslator translator =
104+
new PostgresWriteSchemaTransformTranslator();
105+
Row translatedRow = translator.toConfigRow(writeTransform);
106+
107+
PostgresWriteSchemaTransform writeTransformFromRow =
108+
translator.fromConfigRow(translatedRow, PipelineOptionsFactory.create());
109+
110+
assertEquals(WRITE_CONFIG, writeTransformFromRow.getConfigurationRow());
111+
}
112+
113+
@Test
114+
public void testWriteTransformProtoTranslation()
115+
throws InvalidProtocolBufferException, IOException {
116+
// First build a pipeline
117+
Pipeline p = Pipeline.create();
118+
Schema inputSchema = Schema.builder().addStringField("name").build();
119+
PCollection<Row> input =
120+
p.apply(
121+
Create.of(
122+
Collections.singletonList(
123+
Row.withSchema(inputSchema).addValue("test").build())))
124+
.setRowSchema(inputSchema);
125+
126+
PostgresWriteSchemaTransform writeTransform =
127+
(PostgresWriteSchemaTransform) WRITE_PROVIDER.from(WRITE_CONFIG);
128+
PCollectionRowTuple.of("input", input).apply(writeTransform);
129+
130+
// Then translate the pipeline to a proto and extract PostgresWriteSchemaTransform proto
131+
RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
132+
List<RunnerApi.PTransform> writeTransformProto =
133+
pipelineProto.getComponents().getTransformsMap().values().stream()
134+
.filter(
135+
tr -> {
136+
RunnerApi.FunctionSpec spec = tr.getSpec();
137+
try {
138+
return spec.getUrn().equals(BeamUrns.getUrn(SCHEMA_TRANSFORM))
139+
&& SchemaTransformPayload.parseFrom(spec.getPayload())
140+
.getIdentifier()
141+
.equals(WRITE_PROVIDER.identifier());
142+
} catch (InvalidProtocolBufferException e) {
143+
throw new RuntimeException(e);
144+
}
145+
})
146+
.collect(Collectors.toList());
147+
assertEquals(1, writeTransformProto.size());
148+
RunnerApi.FunctionSpec spec = writeTransformProto.get(0).getSpec();
149+
150+
// Check that the proto contains correct values
151+
SchemaTransformPayload payload = SchemaTransformPayload.parseFrom(spec.getPayload());
152+
Schema schemaFromSpec = SchemaTranslation.schemaFromProto(payload.getConfigurationSchema());
153+
assertEquals(WRITE_PROVIDER.configurationSchema(), schemaFromSpec);
154+
Row rowFromSpec = RowCoder.of(schemaFromSpec).decode(payload.getConfigurationRow().newInput());
155+
156+
assertEquals(WRITE_CONFIG, rowFromSpec);
157+
158+
// Use the information in the proto to recreate the PostgresWriteSchemaTransform
159+
PostgresWriteSchemaTransformTranslator translator =
160+
new PostgresWriteSchemaTransformTranslator();
161+
PostgresWriteSchemaTransform writeTransformFromSpec =
162+
translator.fromConfigRow(rowFromSpec, PipelineOptionsFactory.create());
163+
164+
assertEquals(WRITE_CONFIG, writeTransformFromSpec.getConfigurationRow());
165+
}
166+
167+
@Test
168+
public void testReCreateReadTransformFromRow() {
169+
// setting a subset of fields here.
170+
PostgresReadSchemaTransform readTransform =
171+
(PostgresReadSchemaTransform) READ_PROVIDER.from(READ_CONFIG);
172+
173+
PostgresReadSchemaTransformTranslator translator = new PostgresReadSchemaTransformTranslator();
174+
Row row = translator.toConfigRow(readTransform);
175+
176+
PostgresReadSchemaTransform readTransformFromRow =
177+
translator.fromConfigRow(row, PipelineOptionsFactory.create());
178+
179+
assertEquals(READ_CONFIG, readTransformFromRow.getConfigurationRow());
180+
}
181+
182+
@Test
183+
public void testReadTransformProtoTranslation()
184+
throws InvalidProtocolBufferException, IOException {
185+
// First build a pipeline
186+
Pipeline p = Pipeline.create();
187+
188+
PostgresReadSchemaTransform readTransform =
189+
(PostgresReadSchemaTransform) READ_PROVIDER.from(READ_CONFIG);
190+
191+
// Mock inferBeamSchema since it requires database connection.
192+
Schema expectedSchema = Schema.builder().addStringField("name").build();
193+
try (MockedStatic<JdbcIO.ReadRows> mock = Mockito.mockStatic(JdbcIO.ReadRows.class)) {
194+
mock.when(() -> JdbcIO.ReadRows.inferBeamSchema(Mockito.any(), Mockito.any()))
195+
.thenReturn(expectedSchema);
196+
PCollectionRowTuple.empty(p).apply(readTransform);
197+
}
198+
199+
// Then translate the pipeline to a proto and extract PostgresReadSchemaTransform proto
200+
RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
201+
List<RunnerApi.PTransform> readTransformProto =
202+
pipelineProto.getComponents().getTransformsMap().values().stream()
203+
.filter(
204+
tr -> {
205+
RunnerApi.FunctionSpec spec = tr.getSpec();
206+
try {
207+
return spec.getUrn().equals(BeamUrns.getUrn(SCHEMA_TRANSFORM))
208+
&& SchemaTransformPayload.parseFrom(spec.getPayload())
209+
.getIdentifier()
210+
.equals(READ_PROVIDER.identifier());
211+
} catch (InvalidProtocolBufferException e) {
212+
throw new RuntimeException(e);
213+
}
214+
})
215+
.collect(Collectors.toList());
216+
assertEquals(1, readTransformProto.size());
217+
RunnerApi.FunctionSpec spec = readTransformProto.get(0).getSpec();
218+
219+
// Check that the proto contains correct values
220+
SchemaTransformPayload payload = SchemaTransformPayload.parseFrom(spec.getPayload());
221+
Schema schemaFromSpec = SchemaTranslation.schemaFromProto(payload.getConfigurationSchema());
222+
assertEquals(READ_PROVIDER.configurationSchema(), schemaFromSpec);
223+
Row rowFromSpec = RowCoder.of(schemaFromSpec).decode(payload.getConfigurationRow().newInput());
224+
assertEquals(READ_CONFIG, rowFromSpec);
225+
226+
// Use the information in the proto to recreate the PostgresReadSchemaTransform
227+
PostgresReadSchemaTransformTranslator translator = new PostgresReadSchemaTransformTranslator();
228+
PostgresReadSchemaTransform readTransformFromSpec =
229+
translator.fromConfigRow(rowFromSpec, PipelineOptionsFactory.create());
230+
231+
assertEquals(READ_CONFIG, readTransformFromSpec.getConfigurationRow());
232+
}
233+
}

0 commit comments

Comments
 (0)