From 4e59b38430df8b7b5b3402ae9b6b4979ea2c589b Mon Sep 17 00:00:00 2001
From: erezrokah <erezrokah@users.noreply.github.com>
Date: Fri, 11 Aug 2023 19:09:11 +0200
Subject: [PATCH] feat: Add JSON scalar

---
 src/memdb/tables.ts  | 10 ++++--
 src/scalar/json.ts   | 77 ++++++++++++++++++++++++++++++++++++++++++++
 src/scalar/scalar.ts |  6 ++++
 src/schema/meta.ts   |  4 +--
 4 files changed, 93 insertions(+), 4 deletions(-)
 create mode 100644 src/scalar/json.ts

diff --git a/src/memdb/tables.ts b/src/memdb/tables.ts
index 9ea80af..d8cb6e0 100644
--- a/src/memdb/tables.ts
+++ b/src/memdb/tables.ts
@@ -4,6 +4,7 @@ import { createColumn } from '../schema/column.js';
 import { addCQIDsColumns } from '../schema/meta.js';
 import { pathResolver, parentColumnResolver } from '../schema/resolvers.js';
 import { createTable } from '../schema/table.js';
+import { JSONType } from '../types/json.js';
 
 export const createTables = () => {
   const allTables = [
@@ -12,8 +13,8 @@ export const createTables = () => {
       title: 'Table 1',
       description: 'Table 1 description',
       resolver: (clientMeta, parent, stream) => {
-        stream.write({ id: 'id-1' });
-        stream.write({ id: 'id-2' });
+        stream.write({ id: 'id-1', json: '{ "a": 1 }' });
+        stream.write({ id: 'id-2', json: [1, 2, 3] });
         return Promise.resolve();
       },
       columns: [
@@ -21,6 +22,11 @@ export const createTables = () => {
           name: 'id',
           resolver: pathResolver('id'),
         }),
+        createColumn({
+          name: 'json',
+          resolver: pathResolver('json'),
+          type: new JSONType(),
+        }),
       ],
     }),
     createTable({
diff --git a/src/scalar/json.ts b/src/scalar/json.ts
new file mode 100644
index 0000000..5a23504
--- /dev/null
+++ b/src/scalar/json.ts
@@ -0,0 +1,77 @@
+import { Utf8 as ArrowString } from '@apache-arrow/esnext-esm';
+
+import { Scalar } from './scalar.js';
+import { isInvalid, NULL_VALUE } from './util.js';
+
+const validate = (value: string) => {
+  try {
+    JSON.parse(value);
+    return true;
+  } catch {
+    return false;
+  }
+};
+
+class JSONType implements Scalar<string> {
+  private _valid = false;
+  private _value = '';
+
+  public constructor(v?: unknown) {
+    this.value = v;
+    return this;
+  }
+
+  public get dataType() {
+    return new ArrowString();
+  }
+
+  public get valid(): boolean {
+    return this._valid;
+  }
+
+  public get value(): string {
+    return this._value;
+  }
+
+  public set value(value: unknown) {
+    if (isInvalid(value)) {
+      this._valid = false;
+      return;
+    }
+
+    if (typeof value === 'string') {
+      this._value = value;
+      this._valid = validate(value);
+      return;
+    }
+
+    if (value instanceof Uint8Array) {
+      this._value = new TextDecoder().decode(value);
+      this._valid = validate(this._value);
+      return;
+    }
+
+    if (value instanceof JSONType) {
+      this._value = value.value;
+      this._valid = value.valid;
+      return;
+    }
+
+    try {
+      this._value = JSON.stringify(value);
+      this._valid = true;
+    } catch {
+      throw new Error(`Unable to set '${value}' as JSON`);
+    }
+  }
+
+  public toString() {
+    if (this._valid) {
+      return this._value;
+    }
+
+    return NULL_VALUE;
+  }
+}
+
+export { JSONType as JSON };
diff --git a/src/scalar/scalar.ts b/src/scalar/scalar.ts
index 9b4aa8d..903e326 100644
--- a/src/scalar/scalar.ts
+++ b/src/scalar/scalar.ts
@@ -1,5 +1,6 @@
 import { DataType, Precision } from '@apache-arrow/esnext-esm';
 
+import { JSONType } from '../types/json.js';
 import { UUIDType } from '../types/uuid.js';
 
 import { Bool } from './bool.js';
@@ -9,6 +10,7 @@ import { Float64 } from './float64.js';
 import { Int16 } from './int16.js';
 import { Int32 } from './int32.js';
 import { Int64 } from './int64.js';
+import { JSON as JSONScalar } from './json.js';
 import { List } from './list.js';
 import { Text } from './text.js';
 import { Timestamp } from './timestamp.js';
@@ -93,5 +95,9 @@ export const newScalar = (dataType: DataType): Scalar<Stringable> => {
     return new UUID();
   }
 
+  if (dataType instanceof JSONType) {
+    return new JSONScalar();
+  }
+
   return new Text();
 };
diff --git a/src/schema/meta.ts b/src/schema/meta.ts
index c1b0dec..3f5faa3 100644
--- a/src/schema/meta.ts
+++ b/src/schema/meta.ts
@@ -1,4 +1,4 @@
-import { Binary, TimeNanosecond } from '@apache-arrow/esnext-esm';
+import { Utf8, TimeNanosecond } from '@apache-arrow/esnext-esm';
 
 import { UUIDType } from '../types/uuid.js';
 
@@ -45,7 +45,7 @@ export const cqSyncTimeColumn = createColumn({
 });
 export const cqSourceNameColumn = createColumn({
   name: '_cq_source_name',
-  type: new Binary(),
+  type: new Utf8(),
   description: 'Internal CQ row that references the source plugin name data was retrieved',
   ignoreInTests: true,
 });