Skip to content

Commit 486ae5d

Browse files
teoincontattoteoincontatto
committed
fix: SGStream migration fail when a TOASTable column has NOT NULL constraint
1 parent 1b1c35f commit 486ae5d

10 files changed

Lines changed: 966 additions & 39 deletions

File tree

‎stackgres-k8s/e2e/spec/stream-to-cluster‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -161,11 +161,11 @@ $(
161161
if is_array_not_supported "$TYPE_CATEGORY" "$TYPE_TYPE" "$TYPE_NAME" "$TYPE_FORMAT_TYPE"
162162
then
163163
cat << TABLE_EOF
164-
CREATE TABLE complex_$TYPE_NAME(i bigint,\$execute\$ || quote_ident('c_$TYPE_NAME') || ' $TYPE_NAME' || \$execute\$, updated boolean DEFAULT false, PRIMARY KEY (i));
164+
CREATE TABLE complex_$TYPE_NAME(i bigint,\$execute\$ || quote_ident('c_$TYPE_NAME') || ' $TYPE_NAME NOT NULL' || \$execute\$, updated boolean DEFAULT false, PRIMARY KEY (i));
165165
TABLE_EOF
166166
else
167167
cat << TABLE_EOF
168-
CREATE TABLE complex_$TYPE_NAME(i bigint,\$execute\$ || quote_ident('c_$TYPE_NAME') || ' $TYPE_NAME' || ',' || quote_ident('c_$TYPE_NAME') || '_array $TYPE_NAME[]' || \$execute\$, updated boolean DEFAULT false, PRIMARY KEY (i));
168+
CREATE TABLE complex_$TYPE_NAME(i bigint,\$execute\$ || quote_ident('c_$TYPE_NAME') || ' $TYPE_NAME NOT NULL' || ',' || quote_ident('c_$TYPE_NAME') || '_array $TYPE_NAME[] NOT NULL' || \$execute\$, updated boolean DEFAULT false, PRIMARY KEY (i));
169169
TABLE_EOF
170170
fi
171171
)

‎stackgres-k8s/src/common/src/main/java/io/stackgres/common/crd/sgstream/StackGresStreamTargetJdbcSinkDebeziumProperties.java‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@ public class StackGresStreamTargetJdbcSinkDebeziumProperties {
2929

3030
private String connectionUrlParameters;
3131

32+
private Boolean detectInsertMode;
33+
34+
private Boolean removePlaceholders;
35+
3236
@JsonProperty("connectionPoolMin_size")
3337
private Integer connectionPoolMinSize;
3438

@@ -109,6 +113,22 @@ public void setConnectionUrlParameters(String connectionUrlParameters) {
109113
this.connectionUrlParameters = connectionUrlParameters;
110114
}
111115

116+
public Boolean getDetectInsertMode() {
117+
return detectInsertMode;
118+
}
119+
120+
public void setDetectInsertMode(Boolean detectInsertMode) {
121+
this.detectInsertMode = detectInsertMode;
122+
}
123+
124+
public Boolean getRemovePlaceholders() {
125+
return removePlaceholders;
126+
}
127+
128+
public void setRemovePlaceholders(Boolean removePlaceholders) {
129+
this.removePlaceholders = removePlaceholders;
130+
}
131+
112132
public Integer getConnectionPoolMinSize() {
113133
return connectionPoolMinSize;
114134
}
@@ -306,10 +326,10 @@ public int hashCode() {
306326
return Objects.hash(batchSize, collectionNameFormat, collectionNamingStrategy,
307327
columnNamingStrategy, connectionPoolAcquireIncrement, connectionPoolMaxSize,
308328
connectionPoolMinSize, connectionPoolTimeout, connectionUrlParameters, databaseTimeZone,
309-
deleteEnabled, dialectPostgresPostgisSchema, dialectSqlserverIdentityInsert,
310-
flushMaxRetries, flushRetryDelayMs, insertMode, primaryKeyFields, primaryKeyMode,
311-
quoteIdentifiers, schemaEvolution, tableNameFormat, tableNamingStrategy, truncateEnabled,
312-
useReductionBuffer, useTimeZone);
329+
deleteEnabled, detectInsertMode, dialectPostgresPostgisSchema,
330+
dialectSqlserverIdentityInsert, flushMaxRetries, flushRetryDelayMs, insertMode,
331+
primaryKeyFields, primaryKeyMode, quoteIdentifiers, removePlaceholders, schemaEvolution,
332+
tableNameFormat, tableNamingStrategy, truncateEnabled, useReductionBuffer, useTimeZone);
313333
}
314334

315335
@Override
@@ -332,6 +352,7 @@ public boolean equals(Object obj) {
332352
&& Objects.equals(connectionUrlParameters, other.connectionUrlParameters)
333353
&& Objects.equals(databaseTimeZone, other.databaseTimeZone)
334354
&& Objects.equals(deleteEnabled, other.deleteEnabled)
355+
&& Objects.equals(detectInsertMode, other.detectInsertMode)
335356
&& Objects.equals(dialectPostgresPostgisSchema, other.dialectPostgresPostgisSchema)
336357
&& Objects.equals(dialectSqlserverIdentityInsert, other.dialectSqlserverIdentityInsert)
337358
&& Objects.equals(flushMaxRetries, other.flushMaxRetries)
@@ -340,6 +361,7 @@ public boolean equals(Object obj) {
340361
&& Objects.equals(primaryKeyFields, other.primaryKeyFields)
341362
&& Objects.equals(primaryKeyMode, other.primaryKeyMode)
342363
&& Objects.equals(quoteIdentifiers, other.quoteIdentifiers)
364+
&& Objects.equals(removePlaceholders, other.removePlaceholders)
343365
&& Objects.equals(schemaEvolution, other.schemaEvolution)
344366
&& Objects.equals(tableNameFormat, other.tableNameFormat)
345367
&& Objects.equals(tableNamingStrategy, other.tableNamingStrategy)

‎stackgres-k8s/src/common/src/main/resources/crds/SGStream.yaml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -968,6 +968,14 @@ spec:
968968
type: string
969969
description: |
970970
Paremeters that are set in the JDBC connection URL. See https://jdbc.postgresql.org/documentation/use/
971+
detectInsertMode:
972+
type: boolean
973+
description: |
974+
Default `true`. Parameter `insertMode` is ignored and the inser mode is detected from the record hints.
975+
removePlaceholders:
976+
type: boolean
977+
description: |
978+
Default `true`. When `true` the placeholders are removed from the records.
971979
connectionPoolMin_size:
972980
type: integer
973981
description: |

‎stackgres-k8s/src/restapi/src/main/java/io/stackgres/apiweb/dto/stream/StreamTargetJdbcSinkDebeziumProperties.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@ public class StreamTargetJdbcSinkDebeziumProperties {
1818

1919
private String connectionUrlParameters;
2020

21+
private Boolean detectInsertMode;
22+
23+
private Boolean removePlaceholders;
24+
2125
@JsonProperty("connectionPoolMin_size")
2226
private Integer connectionPoolMinSize;
2327

@@ -79,6 +83,22 @@ public void setConnectionUrlParameters(String connectionUrlParameters) {
7983
this.connectionUrlParameters = connectionUrlParameters;
8084
}
8185

86+
public Boolean getDetectInsertMode() {
87+
return detectInsertMode;
88+
}
89+
90+
public void setDetectInsertMode(Boolean detectInsertMode) {
91+
this.detectInsertMode = detectInsertMode;
92+
}
93+
94+
public Boolean getRemovePlaceholders() {
95+
return removePlaceholders;
96+
}
97+
98+
public void setRemovePlaceholders(Boolean removePlaceholders) {
99+
this.removePlaceholders = removePlaceholders;
100+
}
101+
82102
public Integer getConnectionPoolMinSize() {
83103
return connectionPoolMinSize;
84104
}

‎stackgres-k8s/src/stream/src/main/java/io/stackgres/stream/jobs/target/migration/EnhancedJdbcChangeEventSink.java‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
/*
2+
* Copyright (C) 2019 OnGres, Inc.
3+
* SPDX-License-Identifier: AGPL-3.0-or-later
4+
*/
5+
6+
package io.stackgres.stream.jobs.target.migration;
7+
8+
import java.sql.SQLException;
9+
import java.util.List;
10+
import java.util.Objects;
11+
12+
import io.debezium.connector.jdbc.JdbcSinkConnectorConfig;
13+
import io.debezium.connector.jdbc.JdbcSinkRecord;
14+
import io.debezium.connector.jdbc.dialect.DatabaseDialect;
15+
import io.debezium.connector.jdbc.relational.TableDescriptor;
16+
import io.debezium.metadata.CollectionId;
17+
import io.debezium.util.Stopwatch;
18+
import io.stackgres.stream.jobs.target.migration.jdbc.JdbcChangeEventSink;
19+
import org.hibernate.StatelessSession;
20+
import org.slf4j.Logger;
21+
import org.slf4j.LoggerFactory;
22+
23+
public class EnhancedJdbcChangeEventSink extends JdbcChangeEventSink {
24+
25+
private static final Logger LOGGER = LoggerFactory.getLogger(JdbcChangeEventSink.class);
26+
27+
private final DatabaseDialect dialect;
28+
private final EnhancedRecordWriter recordWriter;
29+
30+
public EnhancedJdbcChangeEventSink(
31+
JdbcSinkConnectorConfig config,
32+
StatelessSession session,
33+
DatabaseDialect dialect,
34+
EnhancedRecordWriter recordWriter) {
35+
super(config, session, dialect, recordWriter);
36+
this.dialect = dialect;
37+
this.recordWriter = recordWriter;
38+
}
39+
40+
@Override
41+
protected void flushBuffer(
42+
CollectionId collectionId,
43+
List<JdbcSinkRecord> toFlush,
44+
TableDescriptor table) throws SQLException {
45+
if (!recordWriter.isDetectInsertMode()) {
46+
super.flushBuffer(collectionId, toFlush, table);
47+
return;
48+
}
49+
Stopwatch flushBufferStopwatch = Stopwatch.reusable();
50+
Stopwatch tableChangesStopwatch = Stopwatch.reusable();
51+
if (!toFlush.isEmpty()) {
52+
LOGGER.debug("Flushing records in JDBC Writer for table: {}", collectionId.name());
53+
tableChangesStopwatch.start();
54+
tableChangesStopwatch.stop();
55+
final int size = toFlush.size();
56+
int currentIndex = 0;
57+
String currentSqlStatement = getSqlStatementWithHints(table, toFlush.get(0));
58+
boolean wasSnapshot = recordWriter.isSnapshot(toFlush.get(0));
59+
for (int index = 1; index <= size; index++) {
60+
final String sqlStatement;
61+
if (index < size) {
62+
final var currentToFlush = toFlush.get(index);
63+
final boolean isSnapshot = recordWriter.isSnapshot(currentToFlush);
64+
if (wasSnapshot && isSnapshot) {
65+
sqlStatement = currentSqlStatement;
66+
} else {
67+
sqlStatement = getSqlStatementWithHints(table, currentToFlush);
68+
}
69+
wasSnapshot = isSnapshot;
70+
if (Objects.equals(sqlStatement, currentSqlStatement)) {
71+
continue;
72+
}
73+
} else {
74+
sqlStatement = null;
75+
}
76+
flushBufferStopwatch.start();
77+
recordWriter.write(toFlush.subList(currentIndex, index), currentSqlStatement);
78+
flushBufferStopwatch.stop();
79+
currentIndex = index;
80+
currentSqlStatement = sqlStatement;
81+
LOGGER.trace("[PERF] Flush buffer execution time {}", flushBufferStopwatch.durations());
82+
LOGGER.trace("[PERF] Table changes execution time {}", tableChangesStopwatch.durations());
83+
}
84+
}
85+
}
86+
87+
private String getSqlStatementWithHints(
88+
TableDescriptor table,
89+
JdbcSinkRecord record) {
90+
if (!record.isDelete()) {
91+
if (recordWriter.isInsert(record)) {
92+
return dialect.getInsertStatement(table, record);
93+
} else {
94+
return dialect.getUpdateStatement(table, record);
95+
}
96+
} else {
97+
return dialect.getDeleteStatement(table, record);
98+
}
99+
}
100+
101+
}

‎stackgres-k8s/src/stream/src/main/java/io/stackgres/stream/jobs/target/migration/EnhancedRecordWriter.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright (C) 2019 OnGres, Inc.
3+
* SPDX-License-Identifier: AGPL-3.0-or-later
4+
*/
5+
6+
package io.stackgres.stream.jobs.target.migration;
7+
8+
import io.debezium.connector.jdbc.JdbcKafkaSinkRecord;
9+
import io.debezium.connector.jdbc.JdbcSinkConnectorConfig;
10+
import io.debezium.connector.jdbc.JdbcSinkRecord;
11+
import io.debezium.connector.jdbc.QueryBinder;
12+
import io.debezium.connector.jdbc.QueryBinderResolver;
13+
import io.debezium.connector.jdbc.dialect.DatabaseDialect;
14+
import io.stackgres.stream.jobs.target.migration.jdbc.RecordWriter;
15+
import org.hibernate.SharedSessionContract;
16+
17+
public class EnhancedRecordWriter extends RecordWriter {
18+
19+
private final boolean detectInsertMode;
20+
21+
public EnhancedRecordWriter(
22+
SharedSessionContract session,
23+
QueryBinderResolver queryBinderResolver,
24+
JdbcSinkConnectorConfig config,
25+
DatabaseDialect dialect,
26+
boolean detectInsertMode) {
27+
super(session, queryBinderResolver, config, dialect);
28+
this.detectInsertMode = detectInsertMode;
29+
}
30+
31+
public boolean isDetectInsertMode() {
32+
return detectInsertMode;
33+
}
34+
35+
@Override
36+
protected void bindValues(JdbcSinkRecord record, QueryBinder queryBinder) {
37+
if (detectInsertMode) {
38+
super.bindValues(record, queryBinder);
39+
}
40+
if (record.isDelete()) {
41+
bindKeyValuesToQuery(record, queryBinder, 1);
42+
return;
43+
}
44+
45+
if (isInsert(record)) {
46+
int index = bindKeyValuesToQuery(record, queryBinder, 1);
47+
bindNonKeyValuesToQuery(record, queryBinder, index);
48+
} else {
49+
int index = bindNonKeyValuesToQuery(record, queryBinder, 1);
50+
bindKeyValuesToQuery(record, queryBinder, index);
51+
}
52+
}
53+
54+
public boolean isSnapshot(JdbcSinkRecord jdbcSinkRecord) {
55+
return jdbcSinkRecord instanceof JdbcKafkaSinkRecord kafkaSinkRecord
56+
&& kafkaSinkRecord.getOriginalKafkaRecord().headers()
57+
.lastWithName(SgClusterStreamMigrationHandler.JdbcHandler.SNAPSHOT_HEADER_KEY) != null;
58+
}
59+
60+
public boolean isInsert(JdbcSinkRecord jdbcSinkRecord) {
61+
return jdbcSinkRecord instanceof JdbcKafkaSinkRecord kafkaSinkRecord
62+
&& kafkaSinkRecord.getOriginalKafkaRecord().headers()
63+
.lastWithName(SgClusterStreamMigrationHandler.JdbcHandler.INSERT_HEADER_KEY) != null;
64+
}
65+
66+
}

0 commit comments

Comments
 (0)