Skip to content

Commit 72847cf

Browse files
shunpingshunping
committed
Add sqlserver read and write to managed io. TransformTranslation and test are also included.
1 parent 94c8f70 commit 72847cf

7 files changed

Lines changed: 414 additions & 2 deletions

File tree

‎model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/external_transforms.proto‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,10 @@ message ManagedTransforms {
8888
"beam:schematransform:org.apache.beam:oracle_read:v1"];
8989
ORACLE_WRITE = 12 [(org.apache.beam.model.pipeline.v1.beam_urn) =
9090
"beam:schematransform:org.apache.beam:oracle_write:v1"];
91+
SQLSERVER_READ = 13 [(org.apache.beam.model.pipeline.v1.beam_urn) =
92+
"beam:schematransform:org.apache.beam:sql_server_read:v1"];
93+
SQLSERVER_WRITE = 14 [(org.apache.beam.model.pipeline.v1.beam_urn) =
94+
"beam:schematransform:org.apache.beam:sql_server_write:v1"];
9195
}
9296
}
9397

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

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,20 +18,30 @@
1818
package org.apache.beam.sdk.io.jdbc.providers;
1919

2020
import static org.apache.beam.sdk.io.jdbc.JdbcUtil.MSSQL;
21+
import static org.apache.beam.sdk.util.construction.BeamUrns.getUrn;
2122

2223
import com.google.auto.service.AutoService;
24+
import java.util.Collections;
25+
import java.util.List;
26+
import org.apache.beam.model.pipeline.v1.ExternalTransforms;
2327
import org.apache.beam.sdk.io.jdbc.JdbcReadSchemaTransformProvider;
28+
import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
2429
import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
2530
import org.checkerframework.checker.initialization.qual.Initialized;
2631
import org.checkerframework.checker.nullness.qual.NonNull;
2732
import org.checkerframework.checker.nullness.qual.UnknownKeyFor;
33+
import org.slf4j.Logger;
34+
import org.slf4j.LoggerFactory;
2835

2936
@AutoService(SchemaTransformProvider.class)
3037
public class ReadFromSqlServerSchemaTransformProvider extends JdbcReadSchemaTransformProvider {
3138

39+
private static final Logger LOG =
40+
LoggerFactory.getLogger(ReadFromSqlServerSchemaTransformProvider.class);
41+
3242
@Override
3343
public @UnknownKeyFor @NonNull @Initialized String identifier() {
34-
return "beam:schematransform:org.apache.beam:sql_server_read:v1";
44+
return getUrn(ExternalTransforms.ManagedTransforms.Urns.SQLSERVER_READ);
3545
}
3646

3747
@Override
@@ -43,4 +53,31 @@ public String description() {
4353
protected String jdbcType() {
4454
return MSSQL;
4555
}
56+
57+
@Override
58+
public @UnknownKeyFor @NonNull @Initialized SchemaTransform from(
59+
JdbcReadSchemaTransformConfiguration configuration) {
60+
String jdbcType = configuration.getJdbcType();
61+
if (jdbcType != null && !jdbcType.equals(jdbcType())) {
62+
throw new IllegalArgumentException(
63+
String.format("Wrong JDBC type. Expected '%s' but got '%s'", jdbcType(), jdbcType));
64+
}
65+
66+
List<@org.checkerframework.checker.nullness.qual.Nullable String> connectionInitSql =
67+
configuration.getConnectionInitSql();
68+
if (connectionInitSql != null && !connectionInitSql.isEmpty()) {
69+
LOG.warn("SQL Server does not support connectionInitSql, ignoring.");
70+
}
71+
72+
// Override "connectionInitSql" for sqlserver
73+
configuration = configuration.toBuilder().setConnectionInitSql(Collections.emptyList()).build();
74+
return new SqlServerReadSchemaTransform(configuration);
75+
}
76+
77+
public static class SqlServerReadSchemaTransform extends JdbcReadSchemaTransform {
78+
public SqlServerReadSchemaTransform(JdbcReadSchemaTransformConfiguration config) {
79+
super(config, MSSQL);
80+
config.validate(MSSQL);
81+
}
82+
}
4683
}

‎sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/providers/SqlServerSchemaTransformTranslation.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.ReadFromSqlServerSchemaTransformProvider.SqlServerReadSchemaTransform;
21+
import static org.apache.beam.sdk.io.jdbc.providers.WriteToSqlServerSchemaTransformProvider.SqlServerWriteSchemaTransform;
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 SqlServerSchemaTransformTranslation {
34+
static class SqlServerReadSchemaTransformTranslator
35+
extends SchemaTransformPayloadTranslator<SqlServerReadSchemaTransform> {
36+
@Override
37+
public SchemaTransformProvider provider() {
38+
return new ReadFromSqlServerSchemaTransformProvider();
39+
}
40+
41+
@Override
42+
public Row toConfigRow(SqlServerReadSchemaTransform 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(SqlServerReadSchemaTransform.class, new SqlServerReadSchemaTransformTranslator())
60+
.build();
61+
}
62+
}
63+
64+
static class SqlServerWriteSchemaTransformTranslator
65+
extends SchemaTransformPayloadTranslator<SqlServerWriteSchemaTransform> {
66+
@Override
67+
public SchemaTransformProvider provider() {
68+
return new WriteToSqlServerSchemaTransformProvider();
69+
}
70+
71+
@Override
72+
public Row toConfigRow(SqlServerWriteSchemaTransform 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(SqlServerWriteSchemaTransform.class, new SqlServerWriteSchemaTransformTranslator())
90+
.build();
91+
}
92+
}
93+
}

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

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,20 +18,30 @@
1818
package org.apache.beam.sdk.io.jdbc.providers;
1919

2020
import static org.apache.beam.sdk.io.jdbc.JdbcUtil.MSSQL;
21+
import static org.apache.beam.sdk.util.construction.BeamUrns.getUrn;
2122

2223
import com.google.auto.service.AutoService;
24+
import java.util.Collections;
25+
import java.util.List;
26+
import org.apache.beam.model.pipeline.v1.ExternalTransforms;
2327
import org.apache.beam.sdk.io.jdbc.JdbcWriteSchemaTransformProvider;
28+
import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
2429
import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
2530
import org.checkerframework.checker.initialization.qual.Initialized;
2631
import org.checkerframework.checker.nullness.qual.NonNull;
2732
import org.checkerframework.checker.nullness.qual.UnknownKeyFor;
33+
import org.slf4j.Logger;
34+
import org.slf4j.LoggerFactory;
2835

2936
@AutoService(SchemaTransformProvider.class)
3037
public class WriteToSqlServerSchemaTransformProvider extends JdbcWriteSchemaTransformProvider {
3138

39+
private static final Logger LOG =
40+
LoggerFactory.getLogger(WriteToSqlServerSchemaTransformProvider.class);
41+
3242
@Override
3343
public @UnknownKeyFor @NonNull @Initialized String identifier() {
34-
return "beam:schematransform:org.apache.beam:sql_server_write:v1";
44+
return getUrn(ExternalTransforms.ManagedTransforms.Urns.SQLSERVER_WRITE);
3545
}
3646

3747
@Override
@@ -43,4 +53,31 @@ public String description() {
4353
protected String jdbcType() {
4454
return MSSQL;
4555
}
56+
57+
@Override
58+
public @UnknownKeyFor @NonNull @Initialized SchemaTransform from(
59+
JdbcWriteSchemaTransformConfiguration configuration) {
60+
String jdbcType = configuration.getJdbcType();
61+
if (jdbcType != null && !jdbcType.equals(jdbcType())) {
62+
throw new IllegalArgumentException(
63+
String.format("Wrong JDBC type. Expected '%s' but got '%s'", jdbcType(), jdbcType));
64+
}
65+
66+
List<@org.checkerframework.checker.nullness.qual.Nullable String> connectionInitSql =
67+
configuration.getConnectionInitSql();
68+
if (connectionInitSql != null && !connectionInitSql.isEmpty()) {
69+
LOG.warn("SQL Server does not support connectionInitSql, ignoring.");
70+
}
71+
72+
// Override "connectionInitSql" for sqlserver
73+
configuration = configuration.toBuilder().setConnectionInitSql(Collections.emptyList()).build();
74+
return new SqlServerWriteSchemaTransform(configuration);
75+
}
76+
77+
public static class SqlServerWriteSchemaTransform extends JdbcWriteSchemaTransform {
78+
public SqlServerWriteSchemaTransform(JdbcWriteSchemaTransformConfiguration config) {
79+
super(config, MSSQL);
80+
config.validate(MSSQL);
81+
}
82+
}
4683
}

0 commit comments

Comments
 (0)