Propagate array item types from SQL to Avro schema

spotify · Nov 26, 2024 · 1dc7af2 · 1dc7af2
1 parent 9c72eac
commit 1dc7af2
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 4 deletions.
diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroSchema.java b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroSchema.java
@@ -74,6 +74,8 @@ public static Schema createSchemaByReadingOneRow(
     try (Statement statement = connection.createStatement()) {
       final ResultSet resultSet = statement.executeQuery(queryBuilderArgs.sqlQueryWithLimitOne());
 
+      resultSet.next();
+
       final Schema schema =
           createAvroSchema(
               resultSet,
@@ -107,7 +109,7 @@ public static Schema createAvroSchema(
             .prop("tableName", tableName)
             .prop("connectionUrl", connectionUrl)
             .fields();
-    return createAvroFields(meta, builder, useLogicalTypes).endRecord();
+    return createAvroFields(resultSet, builder, useLogicalTypes).endRecord();
   }
 
   static String getDatabaseTableName(final ResultSetMetaData meta) throws SQLException {
@@ -123,11 +125,13 @@ static String getDatabaseTableName(final ResultSetMetaData meta) throws SQLExcep
   }
 
   private static SchemaBuilder.FieldAssembler<Schema> createAvroFields(
-      final ResultSetMetaData meta,
-      final SchemaBuilder.FieldAssembler<Schema> builder,
+      final ResultSet resultSet,
+        final SchemaBuilder.FieldAssembler<Schema> builder,
       final boolean useLogicalTypes)
       throws SQLException {
 
+    ResultSetMetaData meta = resultSet.getMetaData();
+
     for (int i = 1; i <= meta.getColumnCount(); i++) {
 
       final String columnName;
@@ -153,9 +157,15 @@ private static SchemaBuilder.FieldAssembler<Schema> createAvroFields(
               SchemaBuilder.UnionAccumulator<SchemaBuilder.NullDefault<Schema>>>
           fieldSchemaBuilder = field.type().unionOf().nullBuilder().endNull().and();
 
+      Integer arrayItemType = null;
+      if (columnType == ARRAY) {
+        arrayItemType = resultSet.getArray(i).getBaseType();
+      }
+
       final SchemaBuilder.UnionAccumulator<SchemaBuilder.NullDefault<Schema>> schemaFieldAssembler =
           setAvroColumnType(
               columnType,
+              arrayItemType,
               meta.getPrecision(i),
               columnClassName,
               useLogicalTypes,
@@ -181,6 +191,7 @@ private static SchemaBuilder.FieldAssembler<Schema> createAvroFields(
   private static SchemaBuilder.UnionAccumulator<SchemaBuilder.NullDefault<Schema>>
       setAvroColumnType(
           final int columnType,
+          final Integer arrayItemType,
           final int precision,
           final String columnClassName,
           final boolean useLogicalTypes,
@@ -225,10 +236,12 @@ private static SchemaBuilder.FieldAssembler<Schema> createAvroFields(
         } else {
           return field.bytesType();
         }
+      case ARRAY:
+        return setAvroColumnType(arrayItemType, null, precision, columnClassName,
+            useLogicalTypes, field.array().items());
       case BINARY:
       case VARBINARY:
       case LONGVARBINARY:
-      case ARRAY:
       case BLOB:
         return field.bytesType();
       case DOUBLE:

diff --git a/e2e/e2e.sh b/e2e/e2e.sh
@@ -1,12 +1,15 @@
 #!/usr/bin/env bash
 
+echo $1
+
 # fail on error
 set -o errexit
 set -o nounset
 set -o pipefail
 
 readonly SCRIPT_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null && pwd)"
 readonly PROJECT_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." >/dev/null && pwd)"
+readonly MOUNT_OUTPUT=0
 
 # This file contatins psql views with complex types to validate and troubleshoot dbeam
 
@@ -65,9 +68,14 @@ pack() {
 }
 
 run_docker_dbeam() {
+  OUTPUT_MOUNT_EXP=""
+  if [ "$MOUNT_OUTPUT" -eq 1 ]; then
+    OUTPUT_MOUNT_EXP=--mount="type=bind,source=$SCRIPT_PATH,target=$SCRIPT_PATH"
+  fi
   time docker run --interactive --rm \
     --net="$DOCKER_NETWORK" \
     --mount="type=bind,source=$PROJECT_PATH/dbeam-core/target,target=/dbeam" \
+    $OUTPUT_MOUNT_EXP \
     --memory=1G \
     --entrypoint=/usr/bin/java \
     "$JAVA_DOCKER_IMAGE" \