Added support for PG Geometry (fixes drizzle-team#2516), PG Vector (f…

…ixes drizzle-team#2424)
Sukairo-02 · Jul 3, 2024 · 8900997 · 8900997
1 parent 7db1e69
commit 8900997
Show file tree

Hide file tree

Showing 6 changed files with 197 additions and 19 deletions.
diff --git a/drizzle-typebox/src/index.ts b/drizzle-typebox/src/index.ts
@@ -11,6 +11,7 @@ import type {
 	TOptional,
 	TSchema,
 	TString,
+	TTuple,
 	TUnion,
 } from '@sinclair/typebox';
 import { Type } from '@sinclair/typebox';
@@ -81,8 +82,14 @@ type MaybeOptional<
 
 type GetTypeboxType<TColumn extends Column> = TColumn['_']['dataType'] extends infer TDataType
 	? TDataType extends 'custom' ? TAny
-	: TDataType extends 'json' ? Json
-	: TDataType extends 'array' ? TArray<
+	: TDataType extends 'json' ? TColumn['_']['columnType'] extends 'PgGeometryObject' ? TObject<{
+				x: TNumber;
+				y: TNumber;
+			}>
+		: Json
+	: TDataType extends 'array' ? TColumn['_']['columnType'] extends 'PgVector' ? TArray<TNumber>
+		: TColumn['_']['columnType'] extends 'PgGeometry' ? TTuple<[TNumber, TNumber]>
+		: TArray<
 			GetTypeboxType<
 				Assume<
 					TColumn['_'],
@@ -296,11 +303,15 @@ function mapColumnToSchema(column: Column): TSchema {
 	} else if (column.dataType === 'custom') {
 		type = Type.Any();
 	} else if (column.dataType === 'json') {
-		type = jsonSchema;
+		type = column.columnType === 'PgGeometryObject' ? Type.Object({ x: Type.Number(), y: Type.Number() }) : jsonSchema;
 	} else if (column.dataType === 'array') {
-		type = Type.Array(
-			mapColumnToSchema((column as PgArray<any, any>).baseColumn),
-		);
+		type = column.columnType === 'PgVector'
+			? Type.Array(Type.Number())
+			: column.columnType === 'PgGeometry'
+			? Type.Tuple([Type.Number(), Type.Number()])
+			: Type.Array(
+				mapColumnToSchema((column as PgArray<any, any>).baseColumn),
+			);
 	} else if (column.dataType === 'number') {
 		type = Type.Number();
 	} else if (column.dataType === 'bigint') {

diff --git a/drizzle-typebox/tests/pg.test.ts b/drizzle-typebox/tests/pg.test.ts
@@ -1,7 +1,19 @@
 import { Type } from '@sinclair/typebox';
 import { Value } from '@sinclair/typebox/value';
 import test from 'ava';
-import { char, date, integer, pgEnum, pgTable, serial, text, timestamp, varchar } from 'drizzle-orm/pg-core';
+import {
+	char,
+	date,
+	geometry,
+	integer,
+	pgEnum,
+	pgTable,
+	serial,
+	text,
+	timestamp,
+	varchar,
+	vector,
+} from 'drizzle-orm/pg-core';
 import { createInsertSchema, createSelectSchema, Nullable } from '../src';
 import { expectSchemaShape } from './utils';
 
@@ -23,6 +35,13 @@ const users = pgTable('users', {
 		.default('user'),
 	profession: varchar('profession', { length: 20 }).notNull(),
 	initials: char('initials', { length: 2 }).notNull(),
+	vector: vector('vector', { dimensions: 2 }),
+	geoXy: geometry('geometry_xy', {
+		mode: 'xy',
+	}),
+	geoTuple: geometry('geometry_tuple', {
+		mode: 'tuple',
+	}),
 });
 
 const testUser = {
@@ -39,6 +58,12 @@ const testUser = {
 	roleText2: 'admin',
 	profession: 'Software Engineer',
 	initials: 'JD',
+	vector: [1, 2],
+	geoXy: {
+		x: 10,
+		y: 20.3,
+	},
+	geoTuple: [10, 20.3],
 };
 
 test('users insert valid user', (t) => {
@@ -112,6 +137,12 @@ test('users insert schema', (t) => {
 		),
 		profession: Type.String({ maxLength: 20, minLength: 1 }),
 		initials: Type.String({ maxLength: 2, minLength: 1 }),
+		vector: Type.Optional(Nullable(Type.Array(Type.Number()))),
+		geoXy: Type.Optional(Nullable(Type.Object({
+			x: Type.Number(),
+			y: Type.Number(),
+		}))),
+		geoTuple: Type.Optional(Nullable(Type.Tuple([Type.Number(), Type.Number()]))),
 	});
 
 	expectSchemaShape(t, expected).from(actual);
@@ -136,6 +167,12 @@ test('users insert schema w/ defaults', (t) => {
 		),
 		profession: Type.String({ maxLength: 20, minLength: 1 }),
 		initials: Type.String({ maxLength: 2, minLength: 1 }),
+		vector: Type.Optional(Nullable(Type.Array(Type.Number()))),
+		geoXy: Type.Optional(Nullable(Type.Object({
+			x: Type.Number(),
+			y: Type.Number(),
+		}))),
+		geoTuple: Type.Optional(Nullable(Type.Tuple([Type.Number(), Type.Number()]))),
 	});
 
 	expectSchemaShape(t, expected).from(actual);
@@ -170,6 +207,12 @@ test('users select schema', (t) => {
 		roleText2: Type.Union([Type.Literal('admin'), Type.Literal('user')]),
 		profession: Type.String({ maxLength: 20, minLength: 1 }),
 		initials: Type.String({ maxLength: 2, minLength: 1 }),
+		vector: Nullable(Type.Array(Type.Number())),
+		geoXy: Nullable(Type.Object({
+			x: Type.Number(),
+			y: Type.Number(),
+		})),
+		geoTuple: Nullable(Type.Tuple([Type.Number(), Type.Number()])),
 	});
 
 	expectSchemaShape(t, expected).from(actual);
@@ -192,6 +235,12 @@ test('users select schema w/ defaults', (t) => {
 		roleText2: Type.Union([Type.Literal('admin'), Type.Literal('user')]),
 		profession: Type.String({ maxLength: 20, minLength: 1 }),
 		initials: Type.String({ maxLength: 2, minLength: 1 }),
+		vector: Nullable(Type.Array(Type.Number())),
+		geoXy: Nullable(Type.Object({
+			x: Type.Number(),
+			y: Type.Number(),
+		})),
+		geoTuple: Nullable(Type.Tuple([Type.Number(), Type.Number()])),
 	});
 
 	expectSchemaShape(t, expected).from(actual);

diff --git a/drizzle-valibot/src/index.ts b/drizzle-valibot/src/index.ts
@@ -40,6 +40,8 @@ import {
 	record,
 	string,
 	type StringSchema,
+	tuple,
+	TupleSchema,
 	union,
 	uuid,
 } from 'valibot';
@@ -79,9 +81,14 @@ type MaybeOptional<
 
 type GetValibotType<TColumn extends Column> = TColumn['_']['dataType'] extends infer TDataType
 	? TDataType extends 'custom' ? AnySchema
-	: TDataType extends 'json' ? Json
+	: TDataType extends 'json' ? TColumn['_']['columnType'] extends 'PgGeometryObject' ? ObjectSchema<{
+				x: NumberSchema;
+				y: NumberSchema;
+			}>
+		: Json
 	: TDataType extends 'array'
-		? TColumn['_']['baseColumn'] extends Column ? ArraySchema<GetValibotType<TColumn['_']['baseColumn']>>
+		? TColumn['_']['columnType'] extends 'PgGeometry' ? TupleSchema<[NumberSchema, NumberSchema]>
+		: TColumn['_']['columnType'] extends 'PgVector' ? ArraySchema<NumberSchema>
 		: never
 	: TColumn extends { enumValues: [string, ...string[]] }
 		? Equal<TColumn['enumValues'], [string, ...string[]]> extends true ? StringSchema
@@ -285,11 +292,13 @@ function mapColumnToSchema(column: Column): BaseSchema<any, any> {
 	} else if (column.dataType === 'custom') {
 		type = any();
 	} else if (column.dataType === 'json') {
-		type = jsonSchema;
+		type = column.columnType === 'PgGeometryObject' ? object({ x: number(), y: number() }) : jsonSchema;
 	} else if (column.dataType === 'array') {
-		type = array(
-			mapColumnToSchema((column as PgArray<any, any>).baseColumn),
-		);
+		type = column.columnType === 'PgVector'
+			? array(number())
+			: column.columnType === 'PgGeometry'
+			? tuple([number(), number()])
+			: array(mapColumnToSchema((column as PgArray<any, any>).baseColumn));
 	} else if (column.dataType === 'number') {
 		type = number();
 	} else if (column.dataType === 'bigint') {

diff --git a/drizzle-valibot/tests/pg.test.ts b/drizzle-valibot/tests/pg.test.ts
@@ -1,5 +1,17 @@
 import test from 'ava';
-import { char, date, integer, pgEnum, pgTable, serial, text, timestamp, varchar } from 'drizzle-orm/pg-core';
+import {
+	char,
+	date,
+	geometry,
+	integer,
+	pgEnum,
+	pgTable,
+	serial,
+	text,
+	timestamp,
+	varchar,
+	vector,
+} from 'drizzle-orm/pg-core';
 import {
 	array,
 	date as valiDate,
@@ -14,6 +26,7 @@ import {
 	parse,
 	picklist,
 	string,
+	tuple,
 } from 'valibot';
 import { createInsertSchema, createSelectSchema } from '../src';
 import { expectSchemaShape } from './utils';
@@ -36,6 +49,13 @@ const users = pgTable('users', {
 		.default('user'),
 	profession: varchar('profession', { length: 20 }).notNull(),
 	initials: char('initials', { length: 2 }).notNull(),
+	vector: vector('vector', { dimensions: 2 }),
+	geoXy: geometry('geometry_xy', {
+		mode: 'xy',
+	}),
+	geoTuple: geometry('geometry_tuple', {
+		mode: 'tuple',
+	}),
 });
 
 const testUser = {
@@ -52,6 +72,12 @@ const testUser = {
 	roleText2: 'admin' as const,
 	profession: 'Software Engineer',
 	initials: 'JD',
+	vector: [1, 2],
+	geoXy: {
+		x: 10,
+		y: 20.3,
+	},
+	geoTuple: [10, 20.3],
 };
 
 test('users insert valid user', (t) => {
@@ -116,6 +142,12 @@ test('users insert schema', (t) => {
 		roleText2: optional(picklist(['admin', 'user'])),
 		profession: string([maxLength(20), minLength(1)]),
 		initials: string([maxLength(2), minLength(1)]),
+		vector: optional(nullable(array(number()))),
+		geoXy: optional(nullable(object({
+			x: number(),
+			y: number(),
+		}))),
+		geoTuple: optional(nullable(tuple([number(), number()]))),
 	});
 
 	expectSchemaShape(t, expected).from(actual);
@@ -138,6 +170,12 @@ test('users insert schema w/ defaults', (t) => {
 		roleText2: optional(picklist(['admin', 'user'])),
 		profession: string([maxLength(20), minLength(1)]),
 		initials: string([maxLength(2), minLength(1)]),
+		vector: optional(nullable(array(number()))),
+		geoXy: optional(nullable(object({
+			x: number(),
+			y: number(),
+		}))),
+		geoTuple: optional(nullable(tuple([number(), number()]))),
 	});
 
 	expectSchemaShape(t, expected).from(actual);
@@ -164,6 +202,12 @@ test('users select schema', (t) => {
 		roleText2: picklist(['admin', 'user']),
 		profession: string([maxLength(20), minLength(1)]),
 		initials: string([maxLength(2), minLength(1)]),
+		vector: nullable(array(number())),
+		geoXy: nullable(object({
+			x: number(),
+			y: number(),
+		})),
+		geoTuple: nullable(tuple([number(), number()])),
 	});
 
 	expectSchemaShape(t, expected).from(actual);
@@ -186,6 +230,12 @@ test('users select schema w/ defaults', (t) => {
 		roleText2: picklist(['admin', 'user']),
 		profession: string([maxLength(20), minLength(1)]),
 		initials: string([maxLength(2), minLength(1)]),
+		vector: nullable(array(number())),
+		geoXy: nullable(object({
+			x: number(),
+			y: number(),
+		})),
+		geoTuple: nullable(tuple([number(), number()])),
 	});
 
 	expectSchemaShape(t, expected).from(actual);

diff --git a/drizzle-zod/src/index.ts b/drizzle-zod/src/index.ts
@@ -42,8 +42,14 @@ type MaybeOptional<
 
 type GetZodType<TColumn extends Column> = TColumn['_']['dataType'] extends infer TDataType
 	? TDataType extends 'custom' ? z.ZodAny
-	: TDataType extends 'json' ? z.ZodType<Json>
-	: TDataType extends 'array' ? z.ZodArray<GetZodType<Assume<TColumn['_'], { baseColumn: Column }>['baseColumn']>>
+	: TDataType extends 'json' ? TColumn['_']['columnType'] extends 'PgGeometryObject' ? z.ZodObject<{
+				x: z.ZodNumber;
+				y: z.ZodNumber;
+			}, 'strip'>
+		: z.ZodType<Json>
+	: TDataType extends 'array' ? TColumn['_']['columnType'] extends 'PgVector' ? z.ZodArray<z.ZodNumber>
+		: TColumn['_']['columnType'] extends 'PgGeometry' ? z.ZodTuple<[z.ZodNumber, z.ZodNumber]>
+		: z.ZodArray<GetZodType<Assume<TColumn['_'], { baseColumn: Column }>['baseColumn']>>
 	: TColumn extends { enumValues: [string, ...string[]] }
 		? Equal<TColumn['enumValues'], [string, ...string[]]> extends true ? z.ZodString : z.ZodEnum<TColumn['enumValues']>
 	: TDataType extends 'bigint' ? z.ZodBigInt
@@ -202,9 +208,13 @@ function mapColumnToSchema(column: Column): z.ZodTypeAny {
 	} else if (column.dataType === 'custom') {
 		type = z.any();
 	} else if (column.dataType === 'json') {
-		type = jsonSchema;
+		type = column.columnType === 'PgGeometryObject' ? z.object({ x: z.number(), y: z.number() }) : jsonSchema;
 	} else if (column.dataType === 'array') {
-		type = z.array(mapColumnToSchema((column as PgArray<any, any>).baseColumn));
+		type = column.columnType === 'PgVector'
+			? z.array(z.number())
+			: column.columnType === 'PgGeometry'
+			? z.tuple([z.number(), z.number()])
+			: z.array(mapColumnToSchema((column as PgArray<any, any>).baseColumn));
 	} else if (column.dataType === 'number') {
 		type = z.number();
 	} else if (column.dataType === 'bigint') {