Skip to content

Commit f308690

Browse files
cgardenscgardens
authored
Exposing SSL-only version of Postgres Source (#6362)
1 parent d9adbd3 commit f308690

File tree

22 files changed

+605
-69
lines changed

22 files changed

+605
-69
lines changed

‎airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/base/spec_modification/SpecModifyingSource.java

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
/*
2+
* Copyright (c) 2021 Airbyte, Inc., all rights reserved.
3+
*/
4+
5+
package io.airbyte.integrations.base.spec_modification;
6+
7+
import com.fasterxml.jackson.databind.JsonNode;
8+
import io.airbyte.commons.util.AutoCloseableIterator;
9+
import io.airbyte.integrations.base.Source;
10+
import io.airbyte.protocol.models.AirbyteCatalog;
11+
import io.airbyte.protocol.models.AirbyteConnectionStatus;
12+
import io.airbyte.protocol.models.AirbyteMessage;
13+
import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
14+
import io.airbyte.protocol.models.ConnectorSpecification;
15+
16+
/**
17+
* In some cases we want to prune or mutate the spec for an existing source. The common case is that
18+
* we want to remove features that are not appropriate for some reason. e.g. In cloud, we do not
19+
* want to allow users to send data unencrypted.
20+
*/
21+
public abstract class SpecModifyingSource implements Source {
22+
23+
private final Source source;
24+
25+
public SpecModifyingSource(final Source source) {
26+
this.source = source;
27+
}
28+
29+
public abstract ConnectorSpecification modifySpec(ConnectorSpecification originalSpec) throws Exception;
30+
31+
@Override
32+
public ConnectorSpecification spec() throws Exception {
33+
return modifySpec(source.spec());
34+
}
35+
36+
@Override
37+
public AirbyteConnectionStatus check(final JsonNode config) throws Exception {
38+
return source.check(config);
39+
}
40+
41+
@Override
42+
public AirbyteCatalog discover(final JsonNode config) throws Exception {
43+
return source.discover(config);
44+
}
45+
46+
@Override
47+
public AutoCloseableIterator<AirbyteMessage> read(final JsonNode config, final ConfiguredAirbyteCatalog catalog, final JsonNode state)
48+
throws Exception {
49+
return source.read(config, catalog, state);
50+
}
51+
52+
}

‎airbyte-integrations/connectors/source-clickhouse/src/test-integration/java/io/airbyte/integrations/io/airbyte/integration_tests/sources/ClickHouseJdbcSourceAcceptanceTest.java

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ public String getDriverClass() {
3838
}
3939

4040
@Override
41-
public String createTableQuery(String tableName, String columnClause, String primaryKeyClause) {
41+
public String createTableQuery(final String tableName, final String columnClause, final String primaryKeyClause) {
4242
// ClickHouse requires Engine to be mentioned as part of create table query.
4343
// Refer : https://clickhouse.tech/docs/en/engines/table-engines/ for more information
4444
return String.format("CREATE TABLE %s(%s) %s",
@@ -56,12 +56,12 @@ public void tearDown() throws SQLException {
5656
}
5757

5858
@Override
59-
public String primaryKeyClause(List<String> columns) {
59+
public String primaryKeyClause(final List<String> columns) {
6060
if (columns.isEmpty()) {
6161
return "";
6262
}
6363

64-
StringBuilder clause = new StringBuilder();
64+
final StringBuilder clause = new StringBuilder();
6565
clause.append("(");
6666
for (int i = 0; i < columns.size(); i++) {
6767
clause.append(columns.get(i));
@@ -91,7 +91,7 @@ public void setup() throws Exception {
9191
}
9292

9393
@Override
94-
public AbstractJdbcSource getSource() {
94+
public AbstractJdbcSource getJdbcSource() {
9595
return new ClickHouseSource();
9696
}
9797

‎airbyte-integrations/connectors/source-cockroachdb/src/test/java/io/airbyte/integrations/source/cockroachdb/CockroachDbJdbcSourceAcceptanceTest.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ public boolean supportsSchemas() {
9595
}
9696

9797
@Override
98-
public AbstractJdbcSource getSource() {
98+
public AbstractJdbcSource getJdbcSource() {
9999
return new CockroachDbSource();
100100
}
101101

@@ -360,7 +360,7 @@ void testReadMultipleTables() throws Exception {
360360
@Test
361361
void testReadMultipleTablesIncrementally() throws Exception {
362362
final String tableName2 = TABLE_NAME + 2;
363-
String streamName2 = streamName + 2;
363+
final String streamName2 = streamName + 2;
364364
database.execute(ctx -> {
365365
ctx.createStatement().execute(
366366
createTableQuery(getFullyQualifiedTableName(tableName2), "id INTEGER, name VARCHAR(200)",
@@ -493,7 +493,7 @@ void testDiscoverWithMultipleSchemas() throws Exception {
493493
.withSourceDefinedPrimaryKey(List.of(List.of(COL_ROW_ID))));
494494

495495
// sort streams by name so that we are comparing lists with the same order.
496-
Comparator<AirbyteStream> schemaTableCompare = Comparator
496+
final Comparator<AirbyteStream> schemaTableCompare = Comparator
497497
.comparing(stream -> stream.getNamespace() + "." + stream.getName());
498498
expected.getStreams().sort(schemaTableCompare);
499499
actual.getStreams().sort(schemaTableCompare);

‎airbyte-integrations/connectors/source-db2/src/test/java/io.airbyte.integrations.source.db2/Db2JdbcSourceAcceptanceTest.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ public void setup() throws Exception {
7272
public void clean() throws Exception {
7373
// In Db2 before dropping a schema, all objects that were in that schema must be dropped or moved to
7474
// another schema.
75-
for (String tableName : TEST_TABLES) {
75+
for (final String tableName : TEST_TABLES) {
7676
final String dropTableQuery = String
7777
.format("DROP TABLE IF EXISTS %s.%s", SCHEMA_NAME, tableName);
7878
super.database.execute(connection -> connection.createStatement().execute(dropTableQuery));
@@ -116,7 +116,7 @@ public String getDriverClass() {
116116
}
117117

118118
@Override
119-
public AbstractJdbcSource getSource() {
119+
public AbstractJdbcSource getJdbcSource() {
120120
return new Db2Source();
121121
}
122122

‎airbyte-integrations/connectors/source-jdbc/src/test/java/io/airbyte/integrations/source/jdbc/AbstractJdbcSourceAcceptanceTest.java

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ public boolean supportsSchemas() {
6565
}
6666

6767
@Override
68-
public AbstractJdbcSource getSource() {
68+
public AbstractJdbcSource getJdbcSource() {
6969
return new PostgresTestSource();
7070
}
7171

@@ -95,8 +95,8 @@ public PostgresTestSource() {
9595
}
9696

9797
@Override
98-
public JsonNode toDatabaseConfig(JsonNode config) {
99-
ImmutableMap.Builder<Object, Object> configBuilder = ImmutableMap.builder()
98+
public JsonNode toDatabaseConfig(final JsonNode config) {
99+
final ImmutableMap.Builder<Object, Object> configBuilder = ImmutableMap.builder()
100100
.put("username", config.get("username").asText())
101101
.put("jdbc_url", String.format("jdbc:postgresql://%s:%s/%s",
102102
config.get("host").asText(),
@@ -115,7 +115,7 @@ public Set<String> getExcludedInternalNameSpaces() {
115115
return Set.of("information_schema", "pg_catalog", "pg_internal", "catalog_history");
116116
}
117117

118-
public static void main(String[] args) throws Exception {
118+
public static void main(final String[] args) throws Exception {
119119
final Source source = new PostgresTestSource();
120120
LOGGER.info("starting source: {}", PostgresTestSource.class);
121121
new IntegrationRunner(source).run(args);

‎airbyte-integrations/connectors/source-jdbc/src/testFixtures/java/io/airbyte/integrations/source/jdbc/test/JdbcSourceAcceptanceTest.java

Lines changed: 56 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import io.airbyte.commons.util.MoreIterators;
2323
import io.airbyte.db.Databases;
2424
import io.airbyte.db.jdbc.JdbcDatabase;
25+
import io.airbyte.integrations.base.Source;
2526
import io.airbyte.integrations.source.jdbc.AbstractJdbcSource;
2627
import io.airbyte.integrations.source.jdbc.SourceJdbcUtils;
2728
import io.airbyte.integrations.source.relationaldb.models.DbState;
@@ -50,6 +51,7 @@
5051
import java.util.List;
5152
import java.util.Optional;
5253
import java.util.Set;
54+
import java.util.function.Function;
5355
import java.util.stream.Collectors;
5456
import org.hamcrest.Matchers;
5557
import org.junit.jupiter.api.Test;
@@ -95,7 +97,7 @@ public abstract class JdbcSourceAcceptanceTest {
9597

9698
public JsonNode config;
9799
public JdbcDatabase database;
98-
public AbstractJdbcSource source;
100+
public Source source;
99101
public static String streamName;
100102

101103
/**
@@ -126,21 +128,43 @@ public abstract class JdbcSourceAcceptanceTest {
126128
/**
127129
* An instance of the source that should be tests.
128130
*
131+
* @return abstract jdbc source
132+
*/
133+
public abstract AbstractJdbcSource getJdbcSource();
134+
135+
/**
136+
* In some cases the Source that is being tested may be an AbstractJdbcSource, but because it is
137+
* decorated, Java cannot recognize it as such. In these cases, as a workaround a user can choose to
138+
* override getJdbcSource and have it return null. Then they can override this method with the
139+
* decorated source AND override getToDatabaseConfigFunction with the appropriate
140+
* toDatabaseConfigFunction that is hidden behind the decorator.
141+
*
129142
* @return source
130143
*/
131-
public abstract AbstractJdbcSource getSource();
144+
public Source getSource() {
145+
return getJdbcSource();
146+
}
132147

133-
protected String createTableQuery(String tableName, String columnClause, String primaryKeyClause) {
148+
/**
149+
* See getSource() for when to override this method.
150+
*
151+
* @return a function that maps a source's config to a jdbc config.
152+
*/
153+
public Function<JsonNode, JsonNode> getToDatabaseConfigFunction() {
154+
return getJdbcSource()::toDatabaseConfig;
155+
}
156+
157+
protected String createTableQuery(final String tableName, final String columnClause, final String primaryKeyClause) {
134158
return String.format("CREATE TABLE %s(%s %s %s)",
135159
tableName, columnClause, primaryKeyClause.equals("") ? "" : ",", primaryKeyClause);
136160
}
137161

138-
protected String primaryKeyClause(List<String> columns) {
162+
protected String primaryKeyClause(final List<String> columns) {
139163
if (columns.isEmpty()) {
140164
return "";
141165
}
142166

143-
StringBuilder clause = new StringBuilder();
167+
final StringBuilder clause = new StringBuilder();
144168
clause.append("PRIMARY KEY (");
145169
for (int i = 0; i < columns.size(); i++) {
146170
clause.append(columns.get(i));
@@ -155,7 +179,7 @@ protected String primaryKeyClause(List<String> columns) {
155179
public void setup() throws Exception {
156180
source = getSource();
157181
config = getConfig();
158-
final JsonNode jdbcConfig = source.toDatabaseConfig(config);
182+
final JsonNode jdbcConfig = getToDatabaseConfigFunction().apply(config);
159183

160184
streamName = TABLE_NAME;
161185

@@ -253,7 +277,7 @@ void testCheckFailure() throws Exception {
253277
@Test
254278
void testDiscover() throws Exception {
255279
final AirbyteCatalog actual = filterOutOtherSchemas(source.discover(config));
256-
AirbyteCatalog expected = getCatalog(getDefaultNamespace());
280+
final AirbyteCatalog expected = getCatalog(getDefaultNamespace());
257281
assertEquals(expected.getStreams().size(), actual.getStreams().size());
258282
actual.getStreams().forEach(actualStream -> {
259283
final Optional<AirbyteStream> expectedStream =
@@ -265,7 +289,7 @@ void testDiscover() throws Exception {
265289
});
266290
}
267291

268-
protected AirbyteCatalog filterOutOtherSchemas(AirbyteCatalog catalog) {
292+
protected AirbyteCatalog filterOutOtherSchemas(final AirbyteCatalog catalog) {
269293
if (supportsSchemas()) {
270294
final AirbyteCatalog filteredCatalog = Jsons.clone(catalog);
271295
filteredCatalog.setStreams(filteredCatalog.getStreams()
@@ -312,7 +336,7 @@ void testDiscoverWithMultipleSchemas() throws Exception {
312336
Field.of(COL_NAME, JsonSchemaPrimitive.STRING))
313337
.withSupportedSyncModes(Lists.newArrayList(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL)));
314338
// sort streams by name so that we are comparing lists with the same order.
315-
Comparator<AirbyteStream> schemaTableCompare = Comparator.comparing(stream -> stream.getNamespace() + "." + stream.getName());
339+
final Comparator<AirbyteStream> schemaTableCompare = Comparator.comparing(stream -> stream.getNamespace() + "." + stream.getName());
316340
expected.getStreams().sort(schemaTableCompare);
317341
actual.getStreams().sort(schemaTableCompare);
318342
assertEquals(expected, filterOutOtherSchemas(actual));
@@ -325,7 +349,7 @@ void testReadSuccess() throws Exception {
325349
source.read(config, getConfiguredCatalogWithOneStream(getDefaultNamespace()), null));
326350

327351
setEmittedAtToNull(actualMessages);
328-
List<AirbyteMessage> expectedMessages = getTestMessages();
352+
final List<AirbyteMessage> expectedMessages = getTestMessages();
329353
assertThat(expectedMessages, Matchers.containsInAnyOrder(actualMessages.toArray()));
330354
assertThat(actualMessages, Matchers.containsInAnyOrder(expectedMessages.toArray()));
331355
}
@@ -596,7 +620,7 @@ void testReadOneTableIncrementallyTwice() throws Exception {
596620
@Test
597621
void testReadMultipleTablesIncrementally() throws Exception {
598622
final String tableName2 = TABLE_NAME + 2;
599-
String streamName2 = streamName + 2;
623+
final String streamName2 = streamName + 2;
600624
database.execute(ctx -> {
601625
ctx.createStatement().execute(
602626
createTableQuery(getFullyQualifiedTableName(tableName2), "id INTEGER, name VARCHAR(200)", ""));
@@ -692,34 +716,34 @@ void testReadMultipleTablesIncrementally() throws Exception {
692716

693717
// when initial and final cursor fields are the same.
694718
private void incrementalCursorCheck(
695-
String cursorField,
696-
String initialCursorValue,
697-
String endCursorValue,
698-
List<AirbyteMessage> expectedRecordMessages)
719+
final String cursorField,
720+
final String initialCursorValue,
721+
final String endCursorValue,
722+
final List<AirbyteMessage> expectedRecordMessages)
699723
throws Exception {
700724
incrementalCursorCheck(cursorField, cursorField, initialCursorValue, endCursorValue,
701725
expectedRecordMessages);
702726
}
703727

704728
private void incrementalCursorCheck(
705-
String initialCursorField,
706-
String cursorField,
707-
String initialCursorValue,
708-
String endCursorValue,
709-
List<AirbyteMessage> expectedRecordMessages)
729+
final String initialCursorField,
730+
final String cursorField,
731+
final String initialCursorValue,
732+
final String endCursorValue,
733+
final List<AirbyteMessage> expectedRecordMessages)
710734
throws Exception {
711735
incrementalCursorCheck(initialCursorField, cursorField, initialCursorValue, endCursorValue,
712736
expectedRecordMessages,
713737
getConfiguredCatalogWithOneStream(getDefaultNamespace()).getStreams().get(0));
714738
}
715739

716740
private void incrementalCursorCheck(
717-
String initialCursorField,
718-
String cursorField,
719-
String initialCursorValue,
720-
String endCursorValue,
721-
List<AirbyteMessage> expectedRecordMessages,
722-
ConfiguredAirbyteStream airbyteStream)
741+
final String initialCursorField,
742+
final String cursorField,
743+
final String initialCursorValue,
744+
final String endCursorValue,
745+
final List<AirbyteMessage> expectedRecordMessages,
746+
final ConfiguredAirbyteStream airbyteStream)
723747
throws Exception {
724748
airbyteStream.setSyncMode(SyncMode.INCREMENTAL);
725749
airbyteStream.setCursorField(Lists.newArrayList(cursorField));
@@ -856,13 +880,13 @@ protected ConfiguredAirbyteStream createTableWithSpaces() throws SQLException {
856880
Field.of(COL_LAST_NAME_WITH_SPACE, JsonSchemaPrimitive.STRING));
857881
}
858882

859-
public String getFullyQualifiedTableName(String tableName) {
883+
public String getFullyQualifiedTableName(final String tableName) {
860884
return SourceJdbcUtils.getFullyQualifiedTableName(getDefaultSchemaName(), tableName);
861885
}
862886

863887
public void createSchemas() throws SQLException {
864888
if (supportsSchemas()) {
865-
for (String schemaName : TEST_SCHEMAS) {
889+
for (final String schemaName : TEST_SCHEMAS) {
866890
final String createSchemaQuery = String.format("CREATE SCHEMA %s;", schemaName);
867891
database.execute(connection -> connection.createStatement().execute(createSchemaQuery));
868892
}
@@ -871,15 +895,15 @@ public void createSchemas() throws SQLException {
871895

872896
public void dropSchemas() throws SQLException {
873897
if (supportsSchemas()) {
874-
for (String schemaName : TEST_SCHEMAS) {
898+
for (final String schemaName : TEST_SCHEMAS) {
875899
final String dropSchemaQuery = String
876900
.format(DROP_SCHEMA_QUERY, schemaName);
877901
database.execute(connection -> connection.createStatement().execute(dropSchemaQuery));
878902
}
879903
}
880904
}
881905

882-
private JsonNode convertIdBasedOnDatabase(int idValue) {
906+
private JsonNode convertIdBasedOnDatabase(final int idValue) {
883907
if (getDriverClass().toLowerCase().contains("oracle")) {
884908
return Jsons.jsonNode(BigDecimal.valueOf(idValue));
885909
} else if (getDriverClass().toLowerCase().contains("snowflake")) {
@@ -902,8 +926,8 @@ protected String getDefaultNamespace() {
902926
}
903927
}
904928

905-
protected static void setEmittedAtToNull(Iterable<AirbyteMessage> messages) {
906-
for (AirbyteMessage actualMessage : messages) {
929+
protected static void setEmittedAtToNull(final Iterable<AirbyteMessage> messages) {
930+
for (final AirbyteMessage actualMessage : messages) {
907931
if (actualMessage.getRecord() != null) {
908932
actualMessage.getRecord().setEmittedAt(null);
909933
}

‎airbyte-integrations/connectors/source-mssql/src/test/java/io/airbyte/integrations/source/mssql/MssqlJdbcSourceAcceptanceTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ public JsonNode getConfig() {
7474
}
7575

7676
@Override
77-
public AbstractJdbcSource getSource() {
77+
public AbstractJdbcSource getJdbcSource() {
7878
return new MssqlSource();
7979
}
8080

0 commit comments

Comments
 (0)