Skip to content

Commit 04db40d

Browse files
authored
[Fix] [Connector-V2] Postgres support for multiple primary keys (#8526)
1 parent 6e4ee46 commit 04db40d

File tree

2 files changed

+21
-10
lines changed

2 files changed

+21
-10
lines changed

Diff for: seatunnel-connectors-v2/connector-jdbc/src/main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/catalog/psql/PostgresCreateTableSqlBuilder.java

+18-8
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,11 @@ public String build(TablePath tablePath) {
7373
buildColumnSql(column), fieldIde))
7474
.collect(Collectors.toList());
7575

76+
// add primary key
77+
if (createIndex && primaryKey != null) {
78+
columnSqls.add("\t" + buildPrimaryKeySql());
79+
}
80+
7681
if (createIndex && CollectionUtils.isNotEmpty(constraintKeys)) {
7782
for (ConstraintKey constraintKey : constraintKeys) {
7883
if (StringUtils.isBlank(constraintKey.getConstraintName())
@@ -134,14 +139,6 @@ private String buildColumnSql(Column column) {
134139
if (!column.isNullable()) {
135140
columnSql.append(" NOT NULL");
136141
}
137-
138-
// Add primary key directly after the column if it is a primary key
139-
if (createIndex
140-
&& primaryKey != null
141-
&& primaryKey.getColumnNames().contains(column.getName())) {
142-
columnSql.append(" PRIMARY KEY");
143-
}
144-
145142
return columnSql.toString();
146143
}
147144

@@ -163,6 +160,19 @@ private String buildColumnCommentSql(Column column, String tableName) {
163160
return columnCommentSql.toString();
164161
}
165162

163+
private String buildPrimaryKeySql() {
164+
String constraintName = UUID.randomUUID().toString().replace("-", "");
165+
String primaryKeyColumns =
166+
primaryKey.getColumnNames().stream()
167+
.map(
168+
column ->
169+
String.format(
170+
"\"%s\"",
171+
CatalogUtils.getFieldIde(column, fieldIde)))
172+
.collect(Collectors.joining(","));
173+
return "CONSTRAINT \"" + constraintName + "\" PRIMARY KEY (" + primaryKeyColumns + ")";
174+
}
175+
166176
private String buildUniqueKeySql(ConstraintKey constraintKey) {
167177
String constraintName = UUID.randomUUID().toString().replace("-", "");
168178
String indexColumns =

Diff for: seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/catalog/psql/PostgresCreateTableSqlBuilderTest.java

+3-2
Original file line numberDiff line numberDiff line change
@@ -52,9 +52,10 @@ void build() {
5252
catalogTable.getTableId().toTablePath());
5353
String pattern =
5454
"CREATE TABLE \"test\" \\(\n"
55-
+ "\"id\" int4 NOT NULL PRIMARY KEY,\n"
55+
+ "\"id\" int4 NOT NULL,\n"
5656
+ "\"name\" text NOT NULL,\n"
5757
+ "\"age\" int4 NOT NULL,\n"
58+
+ "\tCONSTRAINT \"([a-zA-Z0-9]+)\" PRIMARY KEY \\(\"id\",\"name\"\\),\n"
5859
+ "\tCONSTRAINT \"([a-zA-Z0-9]+)\" UNIQUE \\(\"name\"\\)\n"
5960
+ "\\);";
6061
Assertions.assertTrue(
@@ -142,7 +143,7 @@ private CatalogTable catalogTable(boolean otherDB) {
142143
TableSchema tableSchema =
143144
TableSchema.builder()
144145
.columns(columns)
145-
.primaryKey(PrimaryKey.of("pk_id", Lists.newArrayList("id")))
146+
.primaryKey(PrimaryKey.of("pk_id_name", Lists.newArrayList("id", "name")))
146147
.constraintKey(
147148
Lists.newArrayList(
148149
ConstraintKey.of(

0 commit comments

Comments
 (0)