From ed442961b1783c17d88bccd83caa5860e434652a Mon Sep 17 00:00:00 2001 From: Benjamin Gaidioz Date: Thu, 14 Nov 2024 16:13:52 +0100 Subject: [PATCH] Fixed RD-15104: Support arrays in SQL output --- .../rawlabs/sql/compiler/SqlTypesUtils.scala | 17 +++ .../writers/TypedResultSetJsonWriter.scala | 69 +++++++++ .../TestSqlCompilerServiceAirports.scala | 137 ++++++++++++++++++ 3 files changed, 223 insertions(+) diff --git a/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/SqlTypesUtils.scala b/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/SqlTypesUtils.scala index f47773a87..420913b05 100644 --- a/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/SqlTypesUtils.scala +++ b/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/SqlTypesUtils.scala @@ -145,6 +145,23 @@ object SqlTypesUtils { case Some(t) => Right(t) case None => Left(s"unsupported type: $pgTypeName") } + case java.sql.Types.ARRAY => tipe.typeName match { + case "_int2" => Right(RawListType(RawShortType(true, false), false, false)) + case "_int4" => Right(RawListType(RawIntType(true, false), false, false)) + case "_int8" => Right(RawListType(RawLongType(true, false), false, false)) + case "_float4" => Right(RawListType(RawFloatType(true, false), false, false)) + case "_float8" => Right(RawListType(RawDoubleType(true, false), false, false)) + case "_numeric" => Right(RawListType(RawDecimalType(true, false), false, false)) + case "_bool" => Right(RawListType(RawBoolType(true, false), false, false)) + case "_varchar" => Right(RawListType(RawStringType(true, false), false, false)) + case "_text" => Right(RawListType(RawStringType(true, false), false, false)) + case "_interval" => Right(RawListType(RawIntervalType(true, false), false, false)) + case "_date" => Right(RawListType(RawDateType(true, false), false, false)) + case "_time" => Right(RawListType(RawTimeType(true, false), false, false)) + case "_timestamp" => Right(RawListType(RawTimestampType(true, false), false, false)) + case "_jsonb" | "_json" => Right(RawListType(RawAnyType(), false, false)) + case "_hstore" => Right(RawListType(RawAnyType(), false, false)) + } case _ => jdbcToRawType.get(jdbcType) match { case Some(rawType) => Right(rawType) case None => Left(s"unsupported type: $pgTypeName") diff --git a/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/writers/TypedResultSetJsonWriter.scala b/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/writers/TypedResultSetJsonWriter.scala index 930f59f0d..1241dbfab 100644 --- a/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/writers/TypedResultSetJsonWriter.scala +++ b/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/writers/TypedResultSetJsonWriter.scala @@ -24,6 +24,7 @@ import com.rawlabs.compiler.{ RawDoubleType, RawFloatType, RawIntType, + RawInterval, RawIntervalType, RawIterableType, RawListType, @@ -37,6 +38,7 @@ import com.rawlabs.compiler.{ } import com.rawlabs.sql.compiler.SqlIntervals.{intervalToString, stringToInterval} import com.rawlabs.compiler.utils.RecordFieldsNaming +import org.postgresql.util.{PGInterval, PGobject} import java.io.{IOException, OutputStream} import java.sql.ResultSet @@ -120,6 +122,73 @@ class TypedResultSetJsonWriter(os: OutputStream, maxRows: Option[Long]) { case _: RawDoubleType => gen.writeNumber(v.getDouble(i)) case _: RawDecimalType => gen.writeNumber(v.getBigDecimal(i)) case _: RawStringType => gen.writeString(v.getString(i)) + case RawListType(innerType, _, _) => + val array = v.getArray(i) + if (v.wasNull()) gen.writeNull() + else { + val values = array.getArray.asInstanceOf[Array[AnyRef]] + gen.writeStartArray() + values.foreach { value => + if (value == null) gen.writeNull() + else { + innerType match { + case _: RawBoolType => gen.writeBoolean(value.asInstanceOf[Boolean]) + case _: RawByteType => gen.writeNumber(value.asInstanceOf[Byte].toInt) + case _: RawShortType => gen.writeNumber(value.asInstanceOf[Short].toInt) + case _: RawIntType => gen.writeNumber(value.asInstanceOf[Int]) + case _: RawLongType => gen.writeNumber(value.asInstanceOf[Long]) + case _: RawStringType => gen.writeString(value.asInstanceOf[String]) + case _: RawFloatType => gen.writeNumber(value.asInstanceOf[Float]) + case _: RawDoubleType => gen.writeNumber(value.asInstanceOf[Double]) + case _: RawDecimalType => gen.writeNumber(value.asInstanceOf[java.math.BigDecimal]) + case _: RawIntervalType => + val interval = value.asInstanceOf[PGInterval] + val rawInterval = RawInterval( + interval.getYears, + interval.getMonths, + 0, + interval.getDays, + interval.getHours, + interval.getMinutes, + interval.getWholeSeconds, + interval.getMicroSeconds + ) + gen.writeString(intervalToString(rawInterval)) + case _: RawDateType => + val date = value.asInstanceOf[java.sql.Date].toLocalDate + gen.writeString(dateFormatter.format(date)) + case _: RawTimeType => + val time = value.asInstanceOf[java.sql.Time].toLocalTime + gen.writeString(timeFormatter.format(time)) + case _: RawTimestampType => + val dateTime = value.asInstanceOf[java.sql.Timestamp].toLocalDateTime + gen.writeString(timestampFormatter.format(dateTime)) + case _: RawAnyType => v.getMetaData.getColumnTypeName(i) match { + case "_jsonb" | "_json" => + val data = value.asInstanceOf[String] + val json = mapper.readTree(data) + writeRawJson(json) + case "_hstore" => + val item = value.asInstanceOf[PGobject] + val str = item.getValue + // Parse the hstore string into a map + val hstoreMap = new java.util.HashMap[String, String]() + str + .split(",") + .foreach { pair => + val Array(k, v) = pair.split("=>") + hstoreMap.put(k.strip.replaceAll("\"", ""), v.strip.replaceAll("\"", "")) + } + // Convert hstore to JSON-like structure + val json = mapper.valueToTree[ObjectNode](hstoreMap) + writeRawJson(json) + } + case _ => throw new IOException("unsupported type") + } + } + } + gen.writeEndArray() + } case _: RawAnyType => v.getMetaData.getColumnTypeName(i) match { case "jsonb" | "json" => val data = v.getString(i) diff --git a/sql-compiler/src/test/scala/com/rawlabs/sql/compiler/TestSqlCompilerServiceAirports.scala b/sql-compiler/src/test/scala/com/rawlabs/sql/compiler/TestSqlCompilerServiceAirports.scala index 8001684c7..88247330b 100644 --- a/sql-compiler/src/test/scala/com/rawlabs/sql/compiler/TestSqlCompilerServiceAirports.scala +++ b/sql-compiler/src/test/scala/com/rawlabs/sql/compiler/TestSqlCompilerServiceAirports.scala @@ -125,6 +125,143 @@ class TestSqlCompilerServiceAirports ) } + // ARRAY types + test("""SELECT + | INTERVAL '1 day' AS interval, + | ARRAY[CAST(1 AS SMALLINT), CAST(2 AS SMALLINT), CAST(3 AS SMALLINT), CAST(4 AS SMALLINT)] AS short_array, + | ARRAY[1, 2, 3, 4] AS integer_array, + | ARRAY[1.0, 2.0, 3.0, 4.0] AS double_array, + | ARRAY[1.0::numeric, 2.0::numeric, 3.0::numeric, 4.0::numeric] AS decimal_array, + | ARRAY[true, false, true, false] AS boolean_array, + | ARRAY[DATE '2021-01-01', DATE '2021-01-02', DATE '2021-01-03', DATE '2021-01-04'] AS date_array, + | ARRAY[TIME '12:00:00', TIME '13:00:00', TIME '14:00:00', TIME '15:00:00'] AS time_array, + | ARRAY[TIMESTAMP '2021-01-01 12:00:00', TIMESTAMP '2021-01-02 13:00:00', TIMESTAMP '2021-01-03 14:00:00', TIMESTAMP '2021-01-04 15:00:00'] AS timestamp_array, + | ARRAY[INTERVAL '1 day', INTERVAL '2 days', INTERVAL '3 days', INTERVAL '4 days'] AS interval_array, + | ARRAY['{"a": 2}'::json, '{"b": 3}'::json, '{"c": 4}'::json, '{"d": 5}'::json] AS json_array, + | ARRAY['{"a": 2}'::jsonb, '{"b": 3}'::jsonb, '{"c": 4}'::jsonb, '{"d": 5}'::jsonb] AS jsonb_array, + | ARRAY['"a" => "2", "b" => "3"'::hstore, '"c" => "4", "d" => "5"'::hstore] AS hstore_array, + | ARRAY['apple', 'banana', 'cherry'] AS text_array;""".stripMargin) { t => + val v = compilerService.validate(t.q, asJson()) + assert(v.messages.isEmpty) + val GetProgramDescriptionSuccess(description) = compilerService.getProgramDescription(t.q, asJson()) + val Some(main) = description.maybeRunnable + assert(main.params.contains(Vector.empty)) + val baos = new ByteArrayOutputStream() + assert( + compilerService.execute( + t.q, + asJson(), + None, + baos + ) == ExecutionSuccess(true) + ) + assert( + baos.toString() == + """[ + | { + | "interval": "P1D", + | "short_array": [ + | 1, + | 2, + | 3, + | 4 + | ], + | "integer_array": [ + | 1, + | 2, + | 3, + | 4 + | ], + | "double_array": [ + | 1.0, + | 2.0, + | 3.0, + | 4.0 + | ], + | "decimal_array": [ + | 1.0, + | 2.0, + | 3.0, + | 4.0 + | ], + | "boolean_array": [ + | true, + | false, + | true, + | false + | ], + | "date_array": [ + | "2021-01-01", + | "2021-01-02", + | "2021-01-03", + | "2021-01-04" + | ], + | "time_array": [ + | "12:00:00.000", + | "13:00:00.000", + | "14:00:00.000", + | "15:00:00.000" + | ], + | "timestamp_array": [ + | "2021-01-01T12:00:00.000", + | "2021-01-02T13:00:00.000", + | "2021-01-03T14:00:00.000", + | "2021-01-04T15:00:00.000" + | ], + | "interval_array": [ + | "P1D", + | "P2D", + | "P3D", + | "P4D" + | ], + | "json_array": [ + | { + | "a": 2.0 + | }, + | { + | "b": 3.0 + | }, + | { + | "c": 4.0 + | }, + | { + | "d": 5.0 + | } + | ], + | "jsonb_array": [ + | { + | "a": 2.0 + | }, + | { + | "b": 3.0 + | }, + | { + | "c": 4.0 + | }, + | { + | "d": 5.0 + | } + | ], + | "hstore_array": [ + | { + | "a": "2", + | "b": "3" + | }, + | { + | "c": "4", + | "d": "5" + | } + | ], + | "text_array": [ + | "apple", + | "banana", + | "cherry" + | ] + | } + |]""".stripMargin.replaceAll("\\s+", "") + ) + } + // To be sure our offset checks aren't fooled by internal postgres parameters called $1, $2, ..., $10 (with several digits) test("""-- @type n integer |SELECT :n + 1