[FLINK-39224] Support JSON, JSONB types for PostgreSQL (#189)

Author: och5351

flink-connector-jdbc-postgres/src/main/java/org/apache/flink/connector/jdbc/postgres/database/catalog/PostgresTypeMapper.java

@@ -85,6 +85,8 @@ public class PostgresTypeMapper implements JdbcCatalogTypeMapper {
     private static final String PG_CHARACTER_VARYING = "varchar";
     private static final String PG_CHARACTER_VARYING_ARRAY = "_varchar";
     private static final String PG_UUID = "uuid";
+    private static final String PG_JSON = "json";
+    private static final String PG_JSONB = "jsonb";
 
     @Override
     public DataType mapping(ObjectPath tablePath, ResultSetMetaData metadata, int colIndex)
@@ -158,6 +160,8 @@ protected DataType getMapping(String pgType, int precision, int scale) {
             case PG_CHARACTER_VARYING_ARRAY:
                 return DataTypes.ARRAY(DataTypes.VARCHAR(precision));
             case PG_TEXT:
+            case PG_JSON:
+            case PG_JSONB:
                 return DataTypes.STRING();
             case PG_TEXT_ARRAY:
                 return DataTypes.ARRAY(DataTypes.STRING());

flink-connector-jdbc-postgres/src/test/java/org/apache/flink/connector/jdbc/postgres/database/catalog/PostgresCatalogITCase.java

@@ -23,6 +23,9 @@
 import org.apache.flink.types.Row;
 import org.apache.flink.util.CollectionUtil;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -213,4 +216,37 @@ void testNullUuidTypes() {
                                 .collect());
         assertThat(results).hasToString("[+I[1, null]]");
     }
+
+    @Test
+    void testJsonTypes() throws JsonProcessingException {
+
+        List<Row> results =
+                CollectionUtil.iteratorToList(
+                        tEnv.sqlQuery(String.format("select * from %s", TABLE_JSON_TYPE))
+                                .execute()
+                                .collect());
+
+        ObjectMapper mapper = new ObjectMapper();
+        JsonNode expectedJson =
+                mapper.readTree(
+                        "\"test1\":{\"test1-1\":\"testValue\",\"test1-2\":1,\"test1-3\":[\"test1-3-1\",\"test1-3-2\"]}, 2, \"test2\"");
+
+        assertThat(results).hasToString("[+I[" + expectedJson.toString() + "]]");
+    }
+
+    @Test
+    void testJsonbTypes() throws JsonProcessingException {
+        List<Row> results =
+                CollectionUtil.iteratorToList(
+                        tEnv.sqlQuery(String.format("select * from %s", TABLE_JSONB_TYPE))
+                                .execute()
+                                .collect());
+
+        ObjectMapper mapper = new ObjectMapper();
+        JsonNode expectedJson =
+                mapper.readTree(
+                        "\"test1\":{\"test1-1\":\"testValue\",\"test1-2\":1,\"test1-3\":[\"test1-3-1\",\"test1-3-2\"]}, 2, \"test2\"");
+
+        assertThat(results).hasToString("[+I[" + transformPGJsonb(expectedJson.toString()) + "]]");
+    }
 }

flink-connector-jdbc-postgres/src/test/java/org/apache/flink/connector/jdbc/postgres/database/catalog/PostgresCatalogTest.java

@@ -24,6 +24,7 @@
 import org.apache.flink.table.catalog.exceptions.DatabaseNotExistException;
 import org.apache.flink.table.catalog.exceptions.TableNotExistException;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import org.junit.jupiter.api.Test;
 
 import java.util.Arrays;
@@ -68,6 +69,8 @@ void testListTables() throws DatabaseNotExistException {
                 .isEqualTo(
                         Arrays.asList(
                                 "public.array_table",
+                                "public.json_table",
+                                "public.jsonb_table",
                                 "public.primitive_table",
                                 "public.primitive_table2",
                                 "public.serial_table",
@@ -203,4 +206,19 @@ void testNullUuidDataTypes() throws TableNotExistException {
                         new ObjectPath(PostgresCatalog.DEFAULT_DATABASE, TABLE_UUID_TYPE2));
         assertThat(table.getUnresolvedSchema()).isEqualTo(getNullUuidTable().schema);
     }
+
+    @Test
+    void testJsonDataTypes() throws TableNotExistException, JsonProcessingException {
+        CatalogBaseTable table =
+                catalog.getTable(new ObjectPath(PostgresCatalog.DEFAULT_DATABASE, TABLE_JSON_TYPE));
+        assertThat(table.getUnresolvedSchema()).isEqualTo(getJsonTable().schema);
+    }
+
+    @Test
+    void testJsonbDataTypes() throws TableNotExistException, JsonProcessingException {
+        CatalogBaseTable table =
+                catalog.getTable(
+                        new ObjectPath(PostgresCatalog.DEFAULT_DATABASE, TABLE_JSONB_TYPE));
+        assertThat(table.getUnresolvedSchema()).isEqualTo(getJsonbTable().schema);
+    }
 }

flink-connector-jdbc-postgres/src/test/java/org/apache/flink/connector/jdbc/postgres/database/catalog/PostgresCatalogTestBase.java

@@ -26,6 +26,9 @@
 import org.apache.flink.table.api.Schema;
 import org.apache.flink.table.types.logical.DecimalType;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 
@@ -57,12 +60,14 @@ private static DatabaseMetadata getStaticMetadata() {
     protected static final String TABLE_SERIAL_TYPE = "serial_table";
     protected static final String TABLE_UUID_TYPE = "uuid_table";
     protected static final String TABLE_UUID_TYPE2 = "uuid_table2";
+    protected static final String TABLE_JSON_TYPE = "json_table";
+    protected static final String TABLE_JSONB_TYPE = "jsonb_table";
 
     protected static String baseUrl;
     protected static PostgresCatalog catalog;
 
     @BeforeAll
-    static void init() throws SQLException {
+    static void init() throws SQLException, JsonProcessingException {
         // jdbc:postgresql://localhost:50807/postgres?user=postgres
         String jdbcUrl = getStaticMetadata().getJdbcUrl();
         // jdbc:postgresql://localhost:50807/
@@ -115,6 +120,11 @@ static void init() throws SQLException {
         createTable(
                 PostgresTablePath.fromFlinkTableName(TABLE_UUID_TYPE2),
                 getNullUuidTable().pgSchemaSql);
+        createTable(
+                PostgresTablePath.fromFlinkTableName(TABLE_JSON_TYPE), getJsonTable().pgSchemaSql);
+        createTable(
+                PostgresTablePath.fromFlinkTableName(TABLE_JSONB_TYPE),
+                getJsonbTable().pgSchemaSql);
 
         executeSQL(
                 PostgresCatalog.DEFAULT_DATABASE,
@@ -142,6 +152,14 @@ static void init() throws SQLException {
                 String.format(
                         "insert into %s values (%s);",
                         TABLE_UUID_TYPE2, getNullUuidTable().values));
+        executeSQL(
+                PostgresCatalog.DEFAULT_DATABASE,
+                String.format(
+                        "insert into %s values (%s);", TABLE_JSON_TYPE, getJsonTable().values));
+        executeSQL(
+                PostgresCatalog.DEFAULT_DATABASE,
+                String.format(
+                        "insert into %s values (%s);", TABLE_JSONB_TYPE, getJsonbTable().values));
     }
 
     @AfterAll
@@ -186,6 +204,14 @@ static void afterAll() throws SQLException {
                 PostgresCatalog.DEFAULT_DATABASE,
                 String.format(
                         "DROP TABLE %s ", PostgresTablePath.fromFlinkTableName(TABLE_UUID_TYPE2)));
+        executeSQL(
+                PostgresCatalog.DEFAULT_DATABASE,
+                String.format(
+                        "DROP TABLE %s ", PostgresTablePath.fromFlinkTableName(TABLE_JSON_TYPE)));
+        executeSQL(
+                PostgresCatalog.DEFAULT_DATABASE,
+                String.format(
+                        "DROP TABLE %s ", PostgresTablePath.fromFlinkTableName(TABLE_JSONB_TYPE)));
     }
 
     public static void createTable(PostgresTablePath tablePath, String tableSchemaSql)
@@ -439,4 +465,33 @@ public static TestTable getNullUuidTable() {
                 "id INT, " + "uid_col UUID",
                 "1, NULL");
     }
+
+    static String transformPGJsonb(String json) {
+        return json.replaceAll(":", ": ").replaceAll(",", ", ");
+    }
+
+    static TestTable getJsonTable() throws JsonProcessingException {
+        ObjectMapper mapper = new ObjectMapper();
+        JsonNode jsonNode =
+                mapper.readTree(
+                        "\"test1\":{\"test1-1\":\"testValue\",\"test1-2\":1,\"test1-3\":[\"test1-3-1\",\"test1-3-2\"]}, 2, \"test2\"");
+
+        return new TestTable(
+                Schema.newBuilder().column("json_col", DataTypes.STRING()).build(),
+                "json_col JSON",
+                String.format("'%s'", jsonNode.toString()));
+    }
+
+    static TestTable getJsonbTable() throws JsonProcessingException {
+
+        ObjectMapper mapper = new ObjectMapper();
+        JsonNode jsonNode =
+                mapper.readTree(
+                        "\"test1\":{\"test1-1\":\"testValue\",\"test1-2\":1,\"test1-3\":[\"test1-3-1\",\"test1-3-2\"]}, 2, \"test2\"");
+
+        return new TestTable(
+                Schema.newBuilder().column("jsonb_col", DataTypes.STRING()).build(),
+                "jsonb_col JSONB",
+                String.format("'%s'", jsonNode.toString()));
+    }
 }

flink-connector-jdbc-postgres/src/test/java/org/apache/flink/connector/jdbc/postgres/table/PostgresDynamicTableSourceITCase.java

@@ -25,10 +25,15 @@
 import org.apache.flink.table.api.DataTypes;
 import org.apache.flink.types.Row;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
 import java.math.BigDecimal;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 
 import static org.apache.flink.connector.jdbc.testutils.tables.TableBuilder.dbType;
@@ -51,30 +56,54 @@ protected TableRow createInputTable() {
                 // other fields
                 field("real_col", dbType("REAL"), DataTypes.FLOAT()),
                 field("double_col", dbType("DOUBLE PRECISION"), DataTypes.DOUBLE()),
-                field("time_col", dbType("TIME"), DataTypes.TIME()));
+                field("time_col", dbType("TIME"), DataTypes.TIME()),
+                field("json_col", dbType("JSON"), DataTypes.STRING()),
+                field("jsonb_col", dbType("JSONB"), DataTypes.STRING()));
+    }
+
+    private JsonNode toJson(String json) throws JsonProcessingException {
+        ObjectMapper mapper = new ObjectMapper();
+        return mapper.readTree(json);
+    }
+
+    private String transformPGJsonb(String json) {
+        return json.replaceAll(":", ": ").replaceAll(",", ", ");
     }
 
     protected List<Row> getTestData() {
 
         String uuid1 = "123e4567-e89b-12d3-a456-426614174000";
         String uuid2 = "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11";
+        String json =
+                "\"test1\":{\"test1-1\":\"testValue\",\"test1-2\":1,\"test1-3\":[\"test1-3-1\",\"test1-3-2\"]}, 2, \"test2\"";
+
+        try {
+            JsonNode jsonNode = toJson(json);
 
-        return Arrays.asList(
-                Row.of(
-                        1L,
-                        uuid1,
-                        BigDecimal.valueOf(100.1234),
-                        LocalDateTime.parse("2020-01-01T15:35:00.123456"),
-                        1.175E-37F,
-                        1.79769E308D,
-                        LocalTime.parse("15:35")),
-                Row.of(
-                        2L,
-                        uuid2,
-                        BigDecimal.valueOf(101.1234),
-                        LocalDateTime.parse("2020-01-01T15:36:01.123456"),
-                        -1.175E-37F,
-                        -1.79769E308,
-                        LocalTime.parse("15:36:01")));
+            return Arrays.asList(
+                    Row.of(
+                            1L,
+                            uuid1,
+                            BigDecimal.valueOf(100.1234),
+                            LocalDateTime.parse("2020-01-01T15:35:00.123456"),
+                            1.175E-37F,
+                            1.79769E308D,
+                            LocalTime.parse("15:35"),
+                            jsonNode.toString(),
+                            transformPGJsonb(jsonNode.toString())),
+                    Row.of(
+                            2L,
+                            uuid2,
+                            BigDecimal.valueOf(101.1234),
+                            LocalDateTime.parse("2020-01-01T15:36:01.123456"),
+                            -1.175E-37F,
+                            -1.79769E308,
+                            LocalTime.parse("15:36:01"),
+                            jsonNode.toString(),
+                            transformPGJsonb(jsonNode.toString())));
+        } catch (JsonProcessingException e) {
+            e.printStackTrace();
+        }
+        return Collections.emptyList();
     }
 }