driver-adapters: support arrays in pg and neon (#4310)

Fixes most scalar list tests and some other tests that use scalar lists. Fixed tests: ``` < new::regressions::prisma_13097::prisma_13097::group_by_boolean_array < new::regressions::prisma_13097::prisma_13097::group_by_enum_array < queries::batch::select_one_compound::compound_batch::should_only_batch_if_possible_list < queries::filters::field_reference::bytes_filter::bytes_filter::inclusion_filter < queries::filters::field_reference::bytes_filter::bytes_filter::scalar_list_filters < queries::filters::list_filters::json_lists::equality < queries::filters::list_filters::json_lists::has < queries::filters::list_filters::json_lists::has_every < queries::filters::list_filters::json_lists::has_some < queries::filters::list_filters::json_lists::is_empty < writes::data_types::scalar_list::base::basic_types::behave_like_regular_val_for_create_and_update < writes::data_types::scalar_list::base::basic_types::create_mut_return_items_with_empty_lists < writes::data_types::scalar_list::base::basic_types::create_mut_work_with_list_vals < writes::data_types::scalar_list::base::basic_types::set_base < writes::data_types::scalar_list::base::basic_types::update_mut_push_empty_scalar_list < writes::data_types::scalar_list::decimal::decimal::behave_like_regular_val_for_create_and_update < writes::data_types::scalar_list::decimal::decimal::create_mut_return_items_with_empty_lists < writes::data_types::scalar_list::decimal::decimal::create_mut_work_with_list_vals < writes::data_types::scalar_list::decimal::decimal::update_mut_push_empty_scalar_list < writes::data_types::scalar_list::defaults::basic::basic_empty_write < writes::data_types::scalar_list::defaults::basic::basic_write < writes::data_types::scalar_list::defaults::decimal::basic_empty_write < writes::data_types::scalar_list::defaults::json::basic_empty_write < writes::data_types::scalar_list::defaults::json::basic_write < writes::data_types::scalar_list::json::json::behave_like_regular_val_for_create_and_update < writes::data_types::scalar_list::json::json::create_mut_return_items_with_empty_lists < writes::data_types::scalar_list::json::json::create_mut_work_with_list_vals < writes::data_types::scalar_list::json::json::update_mut_push_empty_scalar_list < writes::top_level_mutations::create_list::create_list::create_not_accept_null_in_set ``` Relevant tests that are not fixed by these changes yet and will need to be addressed in future PRs: ``` raw::sql::null_list::null_list::null_scalar_lists writes::data_types::scalar_list::defaults::decimal::basic_write ``` Fixes: prisma/team-orm#374
prisma · Oct 3, 2023 · 3b37c31 · 3b37c31
1 parent 06e0f08
commit 3b37c31
Show file tree

Hide file tree

Showing 7 changed files with 366 additions and 24 deletions.
diff --git a/...gine/connector-test-kit-rs/query-engine-tests/tests/writes/data_types/scalar_list/json.rs b/...gine/connector-test-kit-rs/query-engine-tests/tests/writes/data_types/scalar_list/json.rs
@@ -70,8 +70,9 @@ mod json {
         // TODO: This specific query currently cannot be sent from the JS client.
         // The client _always_ sends an array as plain json and never as an array of json.
         // We're temporarily ignoring it for the JSON protocol because we can't differentiate a list of json values from a json array.
+        // Similarly, this does not currently work with driver adapters.
         // https://github.com/prisma/prisma/issues/18019
-        if runner.protocol().is_graphql() {
+        if runner.protocol().is_graphql() && !runner.is_external_executor() {
             match_connector_result!(
               &runner,
               r#"mutation {
@@ -161,8 +162,9 @@ mod json {
         // TODO: This specific query currently cannot be sent from the JS client.
         // The client _always_ sends an array as plain json and never as an array of json.
         // We're temporarily ignoring it for the JSON protocol because we can't differentiate a list of json values from a json array.
+        // Similarly, this does not currently work with driver adapters.
         // https://github.com/prisma/prisma/issues/18019
-        if runner.protocol().is_graphql() {
+        if runner.protocol().is_graphql() && !runner.is_external_executor() {
             match_connector_result!(
               &runner,
               r#"mutation {

diff --git a/query-engine/connector-test-kit-rs/query-tests-setup/src/runner/mod.rs b/query-engine/connector-test-kit-rs/query-tests-setup/src/runner/mod.rs
@@ -480,4 +480,8 @@ impl Runner {
     pub fn protocol(&self) -> EngineProtocol {
         self.protocol
     }
+
+    pub fn is_external_executor(&self) -> bool {
+        matches!(self.executor, RunnerExecutor::External(_))
+    }
 }
diff --git a/query-engine/driver-adapters/js/adapter-neon/src/conversion.ts b/query-engine/driver-adapters/js/adapter-neon/src/conversion.ts
@@ -3,6 +3,31 @@ import { types } from '@neondatabase/serverless'
 
 const NeonColumnType = types.builtins
 
+/**
+ * PostgreSQL array column types (not defined in NeonColumnType).
+ */
+const ArrayColumnType = {
+  BOOL_ARRAY: 1000,
+  BYTEA_ARRAY: 1001,
+  BPCHAR_ARRAY: 1014,
+  CHAR_ARRAY: 1002,
+  DATE_ARRAY: 1182,
+  FLOAT4_ARRAY: 1021,
+  FLOAT8_ARRAY: 1022,
+  INT2_ARRAY: 1005,
+  INT4_ARRAY: 1007,
+  JSONB_ARRAY: 3807,
+  JSON_ARRAY: 199,
+  MONEY_ARRAY: 791,
+  NUMERIC_ARRAY: 1231,
+  TEXT_ARRAY: 1009,
+  TIMESTAMP_ARRAY: 1115,
+  TIME_ARRAY: 1183,
+  UUID_ARRAY: 2951,
+  VARCHAR_ARRAY: 1015,
+  XML_ARRAY: 143,
+}
+
 /**
  * This is a simplification of quaint's value inference logic. Take a look at quaint's conversion.rs
  * module to see how other attributes of the field packet such as the field length are used to infer
@@ -48,6 +73,40 @@ export function fieldToColumnType(fieldTypeId: number): ColumnType {
       return ColumnTypeEnum.Text
     case NeonColumnType['BYTEA']:
       return ColumnTypeEnum.Bytes
+
+    case ArrayColumnType.INT2_ARRAY:
+    case ArrayColumnType.INT4_ARRAY:
+      return ColumnTypeEnum.Int32Array
+    case ArrayColumnType.FLOAT4_ARRAY:
+      return ColumnTypeEnum.FloatArray
+    case ArrayColumnType.FLOAT8_ARRAY:
+      return ColumnTypeEnum.DoubleArray
+    case ArrayColumnType.NUMERIC_ARRAY:
+    case ArrayColumnType.MONEY_ARRAY:
+      return ColumnTypeEnum.NumericArray
+    case ArrayColumnType.BOOL_ARRAY:
+      return ColumnTypeEnum.BooleanArray
+    case ArrayColumnType.CHAR_ARRAY:
+      return ColumnTypeEnum.CharArray
+    case ArrayColumnType.TEXT_ARRAY:
+    case ArrayColumnType.VARCHAR_ARRAY:
+    case ArrayColumnType.BPCHAR_ARRAY:
+    case ArrayColumnType.XML_ARRAY:
+      return ColumnTypeEnum.TextArray
+    case ArrayColumnType.DATE_ARRAY:
+      return ColumnTypeEnum.DateArray
+    case ArrayColumnType.TIME_ARRAY:
+      return ColumnTypeEnum.TimeArray
+    case ArrayColumnType.TIMESTAMP_ARRAY:
+      return ColumnTypeEnum.DateTimeArray
+    case ArrayColumnType.JSON_ARRAY:
+    case ArrayColumnType.JSONB_ARRAY:
+      return ColumnTypeEnum.JsonArray
+    case ArrayColumnType.BYTEA_ARRAY:
+      return ColumnTypeEnum.BytesArray
+    case ArrayColumnType.UUID_ARRAY:
+      return ColumnTypeEnum.UuidArray
+
     default:
       if (fieldTypeId >= 10000) {
         // Postgres Custom Types
@@ -78,14 +137,20 @@ const parsePgBytes = types.getTypeParser(NeonColumnType.BYTEA) as (_: string) =>
  * Convert bytes to a JSON-encodable representation since we can't
  * currently send a parsed Buffer or ArrayBuffer across JS to Rust
  * boundary.
+ */
+function convertBytes(serializedBytes: string): number[] {
+  const buffer = parsePgBytes(serializedBytes)
+  return encodeBuffer(buffer)
+}
+
+/**
  * TODO:
  * 1. Check if using base64 would be more efficient than this encoding.
  * 2. Consider the possibility of eliminating re-encoding altogether
  *    and passing bytea hex format to the engine if that can be aligned
  *    with other adapter flavours.
  */
-function convertBytes(serializedBytes: string): number[] {
-  const buffer = parsePgBytes(serializedBytes)
+function encodeBuffer(buffer: Buffer) {
   return Array.from(new Uint8Array(buffer))
 }
 
@@ -97,3 +162,20 @@ types.setTypeParser(NeonColumnType.JSONB, convertJson)
 types.setTypeParser(NeonColumnType.JSON, convertJson)
 types.setTypeParser(NeonColumnType.MONEY, money => money.slice(1))
 types.setTypeParser(NeonColumnType.BYTEA, convertBytes)
+
+const parseBytesArray = types.getTypeParser(ArrayColumnType.BYTEA_ARRAY) as (_: string) => Buffer[]
+
+types.setTypeParser(ArrayColumnType.BYTEA_ARRAY, (serializedBytesArray) => {
+  const buffers = parseBytesArray(serializedBytesArray)
+  return buffers.map(encodeBuffer)
+})
+
+const parseTextArray = types.getTypeParser(ArrayColumnType.TEXT_ARRAY) as (_: string) => string[]
+
+types.setTypeParser(ArrayColumnType.TIME_ARRAY, parseTextArray)
+types.setTypeParser(ArrayColumnType.DATE_ARRAY, parseTextArray)
+types.setTypeParser(ArrayColumnType.TIMESTAMP_ARRAY, parseTextArray)
+
+types.setTypeParser(ArrayColumnType.MONEY_ARRAY, (moneyArray) =>
+  parseTextArray(moneyArray).map((money) => money.slice(1)),
+)
diff --git a/query-engine/driver-adapters/js/adapter-pg/src/conversion.ts b/query-engine/driver-adapters/js/adapter-pg/src/conversion.ts
@@ -3,6 +3,31 @@ import { types } from 'pg'
 
 const PgColumnType = types.builtins
 
+/**
+ * PostgreSQL array column types (not defined in PgColumnType).
+ */
+const ArrayColumnType = {
+  BOOL_ARRAY: 1000,
+  BYTEA_ARRAY: 1001,
+  BPCHAR_ARRAY: 1014,
+  CHAR_ARRAY: 1002,
+  DATE_ARRAY: 1182,
+  FLOAT4_ARRAY: 1021,
+  FLOAT8_ARRAY: 1022,
+  INT2_ARRAY: 1005,
+  INT4_ARRAY: 1007,
+  JSONB_ARRAY: 3807,
+  JSON_ARRAY: 199,
+  MONEY_ARRAY: 791,
+  NUMERIC_ARRAY: 1231,
+  TEXT_ARRAY: 1009,
+  TIMESTAMP_ARRAY: 1115,
+  TIME_ARRAY: 1183,
+  UUID_ARRAY: 2951,
+  VARCHAR_ARRAY: 1015,
+  XML_ARRAY: 143,
+}
+
 /**
  * This is a simplification of quaint's value inference logic. Take a look at quaint's conversion.rs
  * module to see how other attributes of the field packet such as the field length are used to infer
@@ -48,6 +73,40 @@ export function fieldToColumnType(fieldTypeId: number): ColumnType {
       return ColumnTypeEnum.Text
     case PgColumnType['BYTEA']:
       return ColumnTypeEnum.Bytes
+
+    case ArrayColumnType.INT2_ARRAY:
+    case ArrayColumnType.INT4_ARRAY:
+      return ColumnTypeEnum.Int32Array
+    case ArrayColumnType.FLOAT4_ARRAY:
+      return ColumnTypeEnum.FloatArray
+    case ArrayColumnType.FLOAT8_ARRAY:
+      return ColumnTypeEnum.DoubleArray
+    case ArrayColumnType.NUMERIC_ARRAY:
+    case ArrayColumnType.MONEY_ARRAY:
+      return ColumnTypeEnum.NumericArray
+    case ArrayColumnType.BOOL_ARRAY:
+      return ColumnTypeEnum.BooleanArray
+    case ArrayColumnType.CHAR_ARRAY:
+      return ColumnTypeEnum.CharArray
+    case ArrayColumnType.TEXT_ARRAY:
+    case ArrayColumnType.VARCHAR_ARRAY:
+    case ArrayColumnType.BPCHAR_ARRAY:
+    case ArrayColumnType.XML_ARRAY:
+      return ColumnTypeEnum.TextArray
+    case ArrayColumnType.DATE_ARRAY:
+      return ColumnTypeEnum.DateArray
+    case ArrayColumnType.TIME_ARRAY:
+      return ColumnTypeEnum.TimeArray
+    case ArrayColumnType.TIMESTAMP_ARRAY:
+      return ColumnTypeEnum.DateTimeArray
+    case ArrayColumnType.JSON_ARRAY:
+    case ArrayColumnType.JSONB_ARRAY:
+      return ColumnTypeEnum.JsonArray
+    case ArrayColumnType.BYTEA_ARRAY:
+      return ColumnTypeEnum.BytesArray
+    case ArrayColumnType.UUID_ARRAY:
+      return ColumnTypeEnum.UuidArray
+
     default:
       if (fieldTypeId >= 10000) {
         // Postgres Custom Types
@@ -78,14 +137,20 @@ const parsePgBytes = types.getTypeParser(PgColumnType.BYTEA) as (_: string) => B
  * Convert bytes to a JSON-encodable representation since we can't
  * currently send a parsed Buffer or ArrayBuffer across JS to Rust
  * boundary.
+ */
+function convertBytes(serializedBytes: string): number[] {
+  const buffer = parsePgBytes(serializedBytes)
+  return encodeBuffer(buffer)
+}
+
+/**
  * TODO:
  * 1. Check if using base64 would be more efficient than this encoding.
  * 2. Consider the possibility of eliminating re-encoding altogether
  *    and passing bytea hex format to the engine if that can be aligned
  *    with other adapter flavours.
  */
-function convertBytes(serializedBytes: string): number[] {
-  const buffer = parsePgBytes(serializedBytes)
+function encodeBuffer(buffer: Buffer) {
   return Array.from(new Uint8Array(buffer))
 }
 
@@ -97,3 +162,20 @@ types.setTypeParser(PgColumnType.JSONB, convertJson)
 types.setTypeParser(PgColumnType.JSON, convertJson)
 types.setTypeParser(PgColumnType.MONEY, money => money.slice(1))
 types.setTypeParser(PgColumnType.BYTEA, convertBytes)
+
+const parseBytesArray = types.getTypeParser(ArrayColumnType.BYTEA_ARRAY) as (_: string) => Buffer[]
+
+types.setTypeParser(ArrayColumnType.BYTEA_ARRAY, (serializedBytesArray) => {
+  const buffers = parseBytesArray(serializedBytesArray)
+  return buffers.map(encodeBuffer)
+})
+
+const parseTextArray = types.getTypeParser(ArrayColumnType.TEXT_ARRAY) as (_: string) => string[]
+
+types.setTypeParser(ArrayColumnType.TIME_ARRAY, parseTextArray)
+types.setTypeParser(ArrayColumnType.DATE_ARRAY, parseTextArray)
+types.setTypeParser(ArrayColumnType.TIMESTAMP_ARRAY, parseTextArray)
+
+types.setTypeParser(ArrayColumnType.MONEY_ARRAY, (moneyArray) =>
+  parseTextArray(moneyArray).map((money) => money.slice(1)),
+)
diff --git a/query-engine/driver-adapters/js/driver-adapter-utils/src/const.ts b/query-engine/driver-adapters/js/driver-adapter-utils/src/const.ts
@@ -3,24 +3,43 @@
 // them via regular dictionaries.
 // See: https://hackmd.io/@dzearing/Sk3xV0cLs
 export const ColumnTypeEnum = {
-  'Int32': 0,
-  'Int64': 1,
-  'Float': 2,
-  'Double': 3,
-  'Numeric': 4,
-  'Boolean': 5,
-  'Char': 6,
-  'Text': 7,
-  'Date': 8,
-  'Time': 9,
-  'DateTime': 10,
-  'Json': 11,
-  'Enum': 12,
-  'Bytes': 13,
-  'Set': 14,
-  'Uuid': 15,
-  // ...
-  'UnknownNumber': 128
+  // Scalars
+  Int32: 0,
+  Int64: 1,
+  Float: 2,
+  Double: 3,
+  Numeric: 4,
+  Boolean: 5,
+  Char: 6,
+  Text: 7,
+  Date: 8,
+  Time: 9,
+  DateTime: 10,
+  Json: 11,
+  Enum: 12,
+  Bytes: 13,
+  Set: 14,
+  Uuid: 15,
+
+  // Arrays
+  Int32Array: 64,
+  Int64Array: 65,
+  FloatArray: 66,
+  DoubleArray: 67,
+  NumericArray: 68,
+  BooleanArray: 69,
+  CharArray: 70,
+  TextArray: 71,
+  DateArray: 72,
+  TimeArray: 73,
+  DateTimeArray: 74,
+  JsonArray: 75,
+  EnumArray: 76,
+  BytesArray: 77,
+  UuidArray: 78,
+
+  // Custom
+  UnknownNumber: 128,
 } as const
 
 // This string value paired with `ColumnType.Json` will be treated as JSON `null`

diff --git a/query-engine/driver-adapters/src/conversion.rs b/query-engine/driver-adapters/src/conversion.rs
@@ -1,4 +1,5 @@
 use napi::bindgen_prelude::{FromNapiValue, ToNapiValue};
+use napi::NapiValue;
 use quaint::ast::Value as QuaintValue;
 use serde::Serialize;
 use serde_json::value::Value as JsonValue;
@@ -9,6 +10,7 @@ pub enum JSArg {
     RawString(String),
     Value(serde_json::Value),
     Buffer(Vec<u8>),
+    Array(Vec<JSArg>),
 }
 
 impl From<JsonValue> for JSArg {
@@ -35,6 +37,23 @@ impl ToNapiValue for JSArg {
             JSArg::Buffer(bytes) => {
                 ToNapiValue::to_napi_value(env, napi::Env::from_raw(env).create_buffer_with_data(bytes)?.into_raw())
             }
+            // While arrays are encodable as JSON generally, their element might not be, or may be
+            // represented in a different way than we need. We use this custom logic for all arrays
+            // to avoid having separate `JsonArray` and `BytesArray` variants in `JSArg` and
+            // avoid complicating the logic in `conv_params`.
+            JSArg::Array(items) => {
+                let env = napi::Env::from_raw(env);
+                let mut array = env.create_array(items.len().try_into().expect("JS array length must fit into u32"))?;
+
+                for (index, item) in items.into_iter().enumerate() {
+                    let js_value = ToNapiValue::to_napi_value(env.raw(), item)?;
+                    // TODO: NapiRaw could be implemented for sys::napi_value directly, there should
+                    // be no need for re-wrapping; submit a patch to napi-rs and simplify here.
+                    array.set(index as u32, napi::JsUnknown::from_raw_unchecked(env.raw(), js_value))?;
+                }
+
+                ToNapiValue::to_napi_value(env.raw(), array)
+            }
         }
     }
 }
@@ -62,6 +81,7 @@ pub fn conv_params(params: &[QuaintValue<'_>]) -> serde_json::Result<Vec<JSArg>>
                 },
                 None => JsonValue::Null.into(),
             },
+            QuaintValue::Array(Some(items)) => JSArg::Array(conv_params(items)?),
             quaint_value => JSArg::from(JsonValue::from(quaint_value.clone())),
         };