Skip to content

Commit d15e5b8

Browse files
shunpingshunping
committed
Add schema transform translation and test for mysql read and write
1 parent f3e2edc commit d15e5b8

4 files changed

Lines changed: 340 additions & 2 deletions

File tree

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

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
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.sdk.io.jdbc.providers.ReadFromMySqlSchemaTransformProvider.MySqlReadSchemaTransform;
21+
import static org.apache.beam.sdk.io.jdbc.providers.WriteToMySqlSchemaTransformProvider.MySqlWriteSchemaTransform;
22+
import static org.apache.beam.sdk.schemas.transforms.SchemaTransformTranslation.SchemaTransformPayloadTranslator;
23+
24+
import com.google.auto.service.AutoService;
25+
import java.util.Map;
26+
import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
27+
import org.apache.beam.sdk.transforms.PTransform;
28+
import org.apache.beam.sdk.util.construction.PTransformTranslation;
29+
import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar;
30+
import org.apache.beam.sdk.values.Row;
31+
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
32+
33+
public class MySqlSchemaTransformTranslation {
34+
static class MySqlReadSchemaTransformTranslator
35+
extends SchemaTransformPayloadTranslator<MySqlReadSchemaTransform> {
36+
@Override
37+
public SchemaTransformProvider provider() {
38+
return new ReadFromMySqlSchemaTransformProvider();
39+
}
40+
41+
@Override
42+
public Row toConfigRow(MySqlReadSchemaTransform transform) {
43+
return transform.getConfigurationRow();
44+
}
45+
}
46+
47+
@AutoService(TransformPayloadTranslatorRegistrar.class)
48+
public static class ReadRegistrar implements TransformPayloadTranslatorRegistrar {
49+
@Override
50+
@SuppressWarnings({
51+
"rawtypes",
52+
})
53+
public Map<
54+
? extends Class<? extends PTransform>,
55+
? extends PTransformTranslation.TransformPayloadTranslator>
56+
getTransformPayloadTranslators() {
57+
return ImmutableMap
58+
.<Class<? extends PTransform>, PTransformTranslation.TransformPayloadTranslator>builder()
59+
.put(MySqlReadSchemaTransform.class, new MySqlReadSchemaTransformTranslator())
60+
.build();
61+
}
62+
}
63+
64+
static class MySqlWriteSchemaTransformTranslator
65+
extends SchemaTransformPayloadTranslator<MySqlWriteSchemaTransform> {
66+
@Override
67+
public SchemaTransformProvider provider() {
68+
return new WriteToMySqlSchemaTransformProvider();
69+
}
70+
71+
@Override
72+
public Row toConfigRow(MySqlWriteSchemaTransform transform) {
73+
return transform.getConfigurationRow();
74+
}
75+
}
76+
77+
@AutoService(TransformPayloadTranslatorRegistrar.class)
78+
public static class WriteRegistrar implements TransformPayloadTranslatorRegistrar {
79+
@Override
80+
@SuppressWarnings({
81+
"rawtypes",
82+
})
83+
public Map<
84+
? extends Class<? extends PTransform>,
85+
? extends PTransformTranslation.TransformPayloadTranslator>
86+
getTransformPayloadTranslators() {
87+
return ImmutableMap
88+
.<Class<? extends PTransform>, PTransformTranslation.TransformPayloadTranslator>builder()
89+
.put(MySqlWriteSchemaTransform.class, new MySqlWriteSchemaTransformTranslator())
90+
.build();
91+
}
92+
}
93+
}

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,13 @@ protected String jdbcType() {
7070
"The fetchSize option is ignored. It is required to set useCursorFetch=true"
7171
+ " in the JDBC URL when using fetchSize for MySQL");
7272
}
73-
return super.from(configuration);
73+
return new MySqlReadSchemaTransform(configuration);
74+
}
75+
76+
public static class MySqlReadSchemaTransform extends JdbcReadSchemaTransform {
77+
public MySqlReadSchemaTransform(JdbcReadSchemaTransformConfiguration config) {
78+
super(config, MYSQL);
79+
config.validate(MYSQL);
80+
}
7481
}
7582
}

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,13 @@ protected String jdbcType() {
5555
throw new IllegalArgumentException(
5656
String.format("Wrong JDBC type. Expected '%s' but got '%s'", jdbcType(), jdbcType));
5757
}
58-
return super.from(configuration);
58+
return new MySqlWriteSchemaTransform(configuration);
59+
}
60+
61+
public static class MySqlWriteSchemaTransform extends JdbcWriteSchemaTransform {
62+
public MySqlWriteSchemaTransform(JdbcWriteSchemaTransformConfiguration config) {
63+
super(config, MYSQL);
64+
config.validate(MYSQL);
65+
}
5966
}
6067
}

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

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
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.MySqlSchemaTransformTranslation.MySqlReadSchemaTransformTranslator;
22+
import static org.apache.beam.sdk.io.jdbc.providers.MySqlSchemaTransformTranslation.MySqlWriteSchemaTransformTranslator;
23+
import static org.apache.beam.sdk.io.jdbc.providers.ReadFromMySqlSchemaTransformProvider.MySqlReadSchemaTransform;
24+
import static org.apache.beam.sdk.io.jdbc.providers.WriteToMySqlSchemaTransformProvider.MySqlWriteSchemaTransform;
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 MysqlSchemaTransformTranslationTest {
56+
@ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder();
57+
58+
@Rule public transient ExpectedException thrown = ExpectedException.none();
59+
60+
static final WriteToMySqlSchemaTransformProvider WRITE_PROVIDER =
61+
new WriteToMySqlSchemaTransformProvider();
62+
static final ReadFromMySqlSchemaTransformProvider READ_PROVIDER =
63+
new ReadFromMySqlSchemaTransformProvider();
64+
65+
static final Row READ_CONFIG =
66+
Row.withSchema(READ_PROVIDER.configurationSchema())
67+
.withFieldValue("jdbc_url", "jdbc:mysql://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:mysql://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+
MySqlWriteSchemaTransform writeTransform =
101+
(MySqlWriteSchemaTransform) WRITE_PROVIDER.from(WRITE_CONFIG);
102+
103+
MySqlWriteSchemaTransformTranslator translator = new MySqlWriteSchemaTransformTranslator();
104+
Row translatedRow = translator.toConfigRow(writeTransform);
105+
106+
MySqlWriteSchemaTransform writeTransformFromRow =
107+
translator.fromConfigRow(translatedRow, PipelineOptionsFactory.create());
108+
109+
assertEquals(WRITE_CONFIG, writeTransformFromRow.getConfigurationRow());
110+
}
111+
112+
@Test
113+
public void testWriteTransformProtoTranslation()
114+
throws InvalidProtocolBufferException, IOException {
115+
// First build a pipeline
116+
Pipeline p = Pipeline.create();
117+
Schema inputSchema = Schema.builder().addStringField("name").build();
118+
PCollection<Row> input =
119+
p.apply(
120+
Create.of(
121+
Collections.singletonList(
122+
Row.withSchema(inputSchema).addValue("test").build())))
123+
.setRowSchema(inputSchema);
124+
125+
MySqlWriteSchemaTransform writeTransform =
126+
(MySqlWriteSchemaTransform) WRITE_PROVIDER.from(WRITE_CONFIG);
127+
PCollectionRowTuple.of("input", input).apply(writeTransform);
128+
129+
// Then translate the pipeline to a proto and extract MySqlWriteSchemaTransform proto
130+
RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
131+
List<RunnerApi.PTransform> writeTransformProto =
132+
pipelineProto.getComponents().getTransformsMap().values().stream()
133+
.filter(
134+
tr -> {
135+
RunnerApi.FunctionSpec spec = tr.getSpec();
136+
try {
137+
return spec.getUrn().equals(BeamUrns.getUrn(SCHEMA_TRANSFORM))
138+
&& SchemaTransformPayload.parseFrom(spec.getPayload())
139+
.getIdentifier()
140+
.equals(WRITE_PROVIDER.identifier());
141+
} catch (InvalidProtocolBufferException e) {
142+
throw new RuntimeException(e);
143+
}
144+
})
145+
.collect(Collectors.toList());
146+
assertEquals(1, writeTransformProto.size());
147+
RunnerApi.FunctionSpec spec = writeTransformProto.get(0).getSpec();
148+
149+
// Check that the proto contains correct values
150+
SchemaTransformPayload payload = SchemaTransformPayload.parseFrom(spec.getPayload());
151+
Schema schemaFromSpec = SchemaTranslation.schemaFromProto(payload.getConfigurationSchema());
152+
assertEquals(WRITE_PROVIDER.configurationSchema(), schemaFromSpec);
153+
Row rowFromSpec = RowCoder.of(schemaFromSpec).decode(payload.getConfigurationRow().newInput());
154+
155+
assertEquals(WRITE_CONFIG, rowFromSpec);
156+
157+
// Use the information in the proto to recreate the MySqlWriteSchemaTransform
158+
MySqlWriteSchemaTransformTranslator translator = new MySqlWriteSchemaTransformTranslator();
159+
MySqlWriteSchemaTransform writeTransformFromSpec =
160+
translator.fromConfigRow(rowFromSpec, PipelineOptionsFactory.create());
161+
162+
assertEquals(WRITE_CONFIG, writeTransformFromSpec.getConfigurationRow());
163+
}
164+
165+
@Test
166+
public void testReCreateReadTransformFromRow() {
167+
// setting a subset of fields here.
168+
MySqlReadSchemaTransform readTransform =
169+
(MySqlReadSchemaTransform) READ_PROVIDER.from(READ_CONFIG);
170+
171+
MySqlReadSchemaTransformTranslator translator = new MySqlReadSchemaTransformTranslator();
172+
Row row = translator.toConfigRow(readTransform);
173+
174+
MySqlReadSchemaTransform readTransformFromRow =
175+
translator.fromConfigRow(row, PipelineOptionsFactory.create());
176+
177+
assertEquals(READ_CONFIG, readTransformFromRow.getConfigurationRow());
178+
}
179+
180+
@Test
181+
public void testReadTransformProtoTranslation()
182+
throws InvalidProtocolBufferException, IOException {
183+
// First build a pipeline
184+
Pipeline p = Pipeline.create();
185+
186+
MySqlReadSchemaTransform readTransform =
187+
(MySqlReadSchemaTransform) READ_PROVIDER.from(READ_CONFIG);
188+
189+
// Mock inferBeamSchema since it requires database connection.
190+
Schema expectedSchema = Schema.builder().addStringField("name").build();
191+
try (MockedStatic<JdbcIO.ReadRows> mock = Mockito.mockStatic(JdbcIO.ReadRows.class)) {
192+
mock.when(() -> JdbcIO.ReadRows.inferBeamSchema(Mockito.any(), Mockito.any()))
193+
.thenReturn(expectedSchema);
194+
PCollectionRowTuple.empty(p).apply(readTransform);
195+
}
196+
197+
// Then translate the pipeline to a proto and extract MySqlReadSchemaTransform proto
198+
RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
199+
List<RunnerApi.PTransform> readTransformProto =
200+
pipelineProto.getComponents().getTransformsMap().values().stream()
201+
.filter(
202+
tr -> {
203+
RunnerApi.FunctionSpec spec = tr.getSpec();
204+
try {
205+
return spec.getUrn().equals(BeamUrns.getUrn(SCHEMA_TRANSFORM))
206+
&& SchemaTransformPayload.parseFrom(spec.getPayload())
207+
.getIdentifier()
208+
.equals(READ_PROVIDER.identifier());
209+
} catch (InvalidProtocolBufferException e) {
210+
throw new RuntimeException(e);
211+
}
212+
})
213+
.collect(Collectors.toList());
214+
assertEquals(1, readTransformProto.size());
215+
RunnerApi.FunctionSpec spec = readTransformProto.get(0).getSpec();
216+
217+
// Check that the proto contains correct values
218+
SchemaTransformPayload payload = SchemaTransformPayload.parseFrom(spec.getPayload());
219+
Schema schemaFromSpec = SchemaTranslation.schemaFromProto(payload.getConfigurationSchema());
220+
assertEquals(READ_PROVIDER.configurationSchema(), schemaFromSpec);
221+
Row rowFromSpec = RowCoder.of(schemaFromSpec).decode(payload.getConfigurationRow().newInput());
222+
assertEquals(READ_CONFIG, rowFromSpec);
223+
224+
// Use the information in the proto to recreate the MySqlReadSchemaTransform
225+
MySqlReadSchemaTransformTranslator translator = new MySqlReadSchemaTransformTranslator();
226+
MySqlReadSchemaTransform readTransformFromSpec =
227+
translator.fromConfigRow(rowFromSpec, PipelineOptionsFactory.create());
228+
229+
assertEquals(READ_CONFIG, readTransformFromSpec.getConfigurationRow());
230+
}
231+
}

0 commit comments

Comments
 (0)