[Fix] [Connector-V2] Postgres support for multiple primary keys

apache · Jan 16, 2025 · 865a231 · 865a231
1 parent dbeeffb
commit 865a231
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 10 deletions.
diff --git a/...pache/seatunnel/connectors/seatunnel/jdbc/catalog/psql/PostgresCreateTableSqlBuilder.java b/...pache/seatunnel/connectors/seatunnel/jdbc/catalog/psql/PostgresCreateTableSqlBuilder.java
@@ -73,6 +73,11 @@ public String build(TablePath tablePath) {
                                                 buildColumnSql(column), fieldIde))
                         .collect(Collectors.toList());
 
+        // add primary key
+        if (createIndex && primaryKey != null) {
+            columnSqls.add("\t" + buildPrimaryKeySql());
+        }
+
         if (createIndex && CollectionUtils.isNotEmpty(constraintKeys)) {
             for (ConstraintKey constraintKey : constraintKeys) {
                 if (StringUtils.isBlank(constraintKey.getConstraintName())
@@ -134,14 +139,6 @@ private String buildColumnSql(Column column) {
         if (!column.isNullable()) {
             columnSql.append(" NOT NULL");
         }
-
-        // Add primary key directly after the column if it is a primary key
-        if (createIndex
-                && primaryKey != null
-                && primaryKey.getColumnNames().contains(column.getName())) {
-            columnSql.append(" PRIMARY KEY");
-        }
-
         return columnSql.toString();
     }
 
@@ -163,6 +160,19 @@ private String buildColumnCommentSql(Column column, String tableName) {
         return columnCommentSql.toString();
     }
 
+    private String buildPrimaryKeySql() {
+        String constraintName = UUID.randomUUID().toString().replace("-", "");
+        String primaryKeyColumns =
+                primaryKey.getColumnNames().stream()
+                        .map(
+                                column ->
+                                        String.format(
+                                                "\"%s\"",
+                                                CatalogUtils.getFieldIde(column, fieldIde)))
+                        .collect(Collectors.joining(","));
+        return "CONSTRAINT \"" + constraintName + "\" PRIMARY KEY (" + primaryKeyColumns + ")";
+    }
+
     private String buildUniqueKeySql(ConstraintKey constraintKey) {
         String constraintName = UUID.randomUUID().toString().replace("-", "");
         String indexColumns =

diff --git a/...e/seatunnel/connectors/seatunnel/jdbc/catalog/psql/PostgresCreateTableSqlBuilderTest.java b/...e/seatunnel/connectors/seatunnel/jdbc/catalog/psql/PostgresCreateTableSqlBuilderTest.java
@@ -52,9 +52,10 @@ void build() {
                                             catalogTable.getTableId().toTablePath());
                             String pattern =
                                     "CREATE TABLE \"test\" \\(\n"
-                                            + "\"id\" int4 NOT NULL PRIMARY KEY,\n"
+                                            + "\"id\" int4 NOT NULL,\n"
                                             + "\"name\" text NOT NULL,\n"
                                             + "\"age\" int4 NOT NULL,\n"
+                                            + "\tCONSTRAINT \"([a-zA-Z0-9]+)\" PRIMARY KEY \\(\"id\",\"name\"\\),\n"
                                             + "\tCONSTRAINT \"([a-zA-Z0-9]+)\" UNIQUE \\(\"name\"\\)\n"
                                             + "\\);";
                             Assertions.assertTrue(
@@ -142,7 +143,7 @@ private CatalogTable catalogTable(boolean otherDB) {
         TableSchema tableSchema =
                 TableSchema.builder()
                         .columns(columns)
-                        .primaryKey(PrimaryKey.of("pk_id", Lists.newArrayList("id")))
+                        .primaryKey(PrimaryKey.of("pk_id_name", Lists.newArrayList("id", "name")))
                         .constraintKey(
                                 Lists.newArrayList(
                                         ConstraintKey.of(