Skip to content

Commit

Permalink
Fixed RD-15104: Support arrays in SQL output (#534)
Browse files Browse the repository at this point in the history 
Postgres queries can legitimately return arrays as columns. So far the
SQL compiler service wouldn't support them. Here's a patch to support
them. Added a test to exercise the logic and make sure it all goes to
the output.
  • Loading branch information
@bgaidioz
bgaidioz authored Nov 15, 2024
1 parent ccec888 commit f222e30
Show file tree
Hide file tree
Showing 3 changed files with 230 additions and 0 deletions.
17 changes: 17 additions & 0 deletions sql-compiler/src/main/scala/com/rawlabs/sql/compiler/SqlTypesUtils.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
69 changes: 69 additions & 0 deletions sql-compiler/src/main/scala/com/rawlabs/sql/compiler/writers/TypedResultSetJsonWriter.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import com.rawlabs.compiler.{
RawDoubleType,
RawFloatType,
RawIntType,
RawInterval,
RawIntervalType,
RawIterableType,
RawListType,
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
144 changes: 144 additions & 0 deletions sql-compiler/src/test/scala/com/rawlabs/sql/compiler/TestSqlCompilerServiceAirports.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,150 @@ 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[CAST(1.1 AS REAL), CAST(2.2 AS REAL), CAST(3.3 AS REAL), CAST(4.4 AS REAL)] AS float_array,
| ARRAY[CAST(1.1 AS DOUBLE PRECISION), CAST(2.2 AS DOUBLE PRECISION), CAST(3.3 AS DOUBLE PRECISION), CAST(4.4 AS DOUBLE PRECISION)] AS double_array,
| ARRAY[1.1::numeric, 2.2::numeric, 3.3::numeric, 4.4::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
| ],
| "float_array": [
| 1.1,
| 2.2,
| 3.3,
| 4.4
| ],
| "double_array": [
| 1.1,
| 2.2,
| 3.3,
| 4.4
| ],
| "decimal_array": [
| 1.1,
| 2.2,
| 3.3,
| 4.4
| ],
| "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
Expand Down

0 comments on commit f222e30

Please sign in to comment.