[Pg] Handle geometry types

drizzle-kit/tests/postgres/pg-tables.test.ts

@@ -203,8 +203,10 @@ test('add table #7', async () => {
 test('add table #8: geometry types', async () => {
 	const to = {
 		users: pgTable('users', {
-			geom: geometry('geom', { type: 'point' }).notNull(),
+			geom: geometry('geom', { type: 'Point' }).notNull(),
 			geom1: geometry('geom1').notNull(),
+			geom2: geometry('geom2', { mode: 'xy', type: 'Point', srid: 4326 }).notNull(),
+			geom3: geometry('geom3', { type: 'MultiPolygon', srid: 3857 }).notNull(),
 		}),
 	};
 
@@ -214,7 +216,7 @@ test('add table #8: geometry types', async () => {
 	// const { sqlStatements: pst } = await push({ db, to });
 
 	const st0 = [
-		`CREATE TABLE "users" (\n\t"geom" geometry(point) NOT NULL,\n\t"geom1" geometry(point) NOT NULL\n);\n`,
+		`CREATE TABLE "users" (\n\t"geom" geometry(Point) NOT NULL,\n\t"geom1" geometry(Point) NOT NULL,\n\t"geom2" geometry(Point,4326) NOT NULL,\n\t"geom3" geometry(MultiPolygon,3857) NOT NULL\n);\n`,
 	];
 	expect(st).toStrictEqual(st0);
 	// expect(pst).toStrictEqual(st0);

drizzle-orm/src/pg-core/columns/postgis_extension/ewkb.ts

@@ -0,0 +1,47 @@
+function hexToBytes(hex: string): Uint8Array {
+	const bytes: number[] = [];
+	for (let c = 0; c < hex.length; c += 2) {
+		bytes.push(Number.parseInt(hex.slice(c, c + 2), 16));
+	}
+	return new Uint8Array(bytes);
+}
+
+function bytesToFloat64(bytes: Uint8Array, offset: number): number {
+	const buffer = new ArrayBuffer(8);
+	const view = new DataView(buffer);
+	for (let i = 0; i < 8; i++) {
+		view.setUint8(i, bytes[offset + i]!);
+	}
+	return view.getFloat64(0, true);
+}
+
+export function parseEWKB(hex: string): { srid: number | undefined; point: [number, number] } {
+	const bytes = hexToBytes(hex);
+
+	let offset = 0;
+
+	// Byte order: 1 is little-endian, 0 is big-endian
+	const byteOrder = bytes[offset];
+	offset += 1;
+
+	const view = new DataView(bytes.buffer);
+	const geomType = view.getUint32(offset, byteOrder === 1);
+	offset += 4;
+
+	let srid: number | undefined;
+	if (geomType & 0x20000000) { // SRID flag
+		srid = view.getUint32(offset, byteOrder === 1);
+		offset += 4;
+	}
+
+	if ((geomType & 0xFFFF) === 1) {
+		const x = bytesToFloat64(bytes, offset);
+		offset += 8;
+		const y = bytesToFloat64(bytes, offset);
+		offset += 8;
+
+		return { srid, point: [x, y] };
+	}
+
+	throw new Error('Unsupported geometry type');
+}

drizzle-orm/src/pg-core/columns/postgis_extension/geometry.ts

@@ -3,39 +3,43 @@ import { entityKind } from '~/entity.ts';
 import type { PgTable } from '~/pg-core/table.ts';
 import { type Equal, getColumnNameAndConfig } from '~/utils.ts';
 import { PgColumn, PgColumnBuilder } from '../common.ts';
-import { parseEWKB } from './utils.ts';
+import { parseEWKB } from './ewkb.ts';
+import type { PgGeometryConfig } from './types.ts';
+import { getGeometrySQLType, mapGeometryToDriverValue } from './utils.ts';
 
 export class PgGeometryBuilder extends PgColumnBuilder<{
 	dataType: 'array geometry';
 	data: [number, number];
 	driverParam: string;
-}, { srid: number | undefined }> {
+}, PgGeometryConfig> {
 	static override readonly [entityKind]: string = 'PgGeometryBuilder';
 
-	constructor(name: string, srid?: number) {
+	static readonly defaultConfig: PgGeometryConfig = { type: 'Point' };
+
+	constructor(name: string, config: PgGeometryConfig = PgGeometryBuilder.defaultConfig) {
 		super(name, 'array geometry', 'PgGeometry');
-		this.config.srid = srid;
+		this.config.type = config.type;
+		this.config.srid = config.srid;
 	}
 
 	/** @internal */
 	override build(table: PgTable<any>) {
 		return new PgGeometry(
 			table,
-			this.config as any,
+			this.config,
 		);
 	}
 }
 
-export class PgGeometry<T extends ColumnBaseConfig<'array geometry'>>
-	extends PgColumn<T, { srid: number | undefined }>
-{
+export class PgGeometry<T extends ColumnBaseConfig<'array geometry'>> extends PgColumn<T, PgGeometryConfig> {
 	static override readonly [entityKind]: string = 'PgGeometry';
 
 	readonly srid = this.config.srid;
+	readonly type = this.config.type;
 	readonly mode = 'tuple';
 
 	getSQLType(): string {
-		return `geometry(point${this.srid === undefined ? '' : `,${this.srid}`})`;
+		return getGeometrySQLType(this.config);
 	}
 
 	override mapFromDriverValue(value: string | [number, number]): [number, number] {
@@ -45,20 +49,22 @@ export class PgGeometry<T extends ColumnBaseConfig<'array geometry'>>
 	}
 
 	override mapToDriverValue(value: [number, number]): string {
-		return `point(${value[0]} ${value[1]})`;
+		return mapGeometryToDriverValue(value, this.config);
 	}
 }
 
 export class PgGeometryObjectBuilder extends PgColumnBuilder<{
 	dataType: 'object geometry';
 	data: { x: number; y: number };
 	driverParam: string;
-}, { srid?: number }> {
+}, PgGeometryConfig> {
 	static override readonly [entityKind]: string = 'PgGeometryObjectBuilder';
 
-	constructor(name: string, srid: number | undefined) {
+	constructor(name: string, config: PgGeometryConfig) {
 		super(name, 'object geometry', 'PgGeometryObject');
-		this.config.srid = srid;
+		this.config.type = config.type;
+		this.config.srid = config.srid;
+		this.config.mode = config.mode;
 	}
 
 	/** @internal */
@@ -70,16 +76,15 @@ export class PgGeometryObjectBuilder extends PgColumnBuilder<{
 	}
 }
 
-export class PgGeometryObject<T extends ColumnBaseConfig<'object geometry'>>
-	extends PgColumn<T, { srid: number | undefined }>
-{
+export class PgGeometryObject<T extends ColumnBaseConfig<'object geometry'>> extends PgColumn<T, PgGeometryConfig> {
 	static override readonly [entityKind]: string = 'PgGeometryObject';
 
 	readonly srid = this.config.srid;
+	readonly type = this.config.type;
 	readonly mode = 'object';
 
 	getSQLType(): string {
-		return `geometry(point${this.srid === undefined ? '' : `,${this.srid}`})`;
+		return getGeometrySQLType(this.config);
 	}
 
 	override mapFromDriverValue(value: string): { x: number; y: number } {
@@ -88,16 +93,10 @@ export class PgGeometryObject<T extends ColumnBaseConfig<'object geometry'>>
 	}
 
 	override mapToDriverValue(value: { x: number; y: number }): string {
-		return `point(${value.x} ${value.y})`;
+		return mapGeometryToDriverValue([value.x, value.y], this.config);
 	}
 }
 
-export interface PgGeometryConfig<T extends 'tuple' | 'xy' = 'tuple' | 'xy'> {
-	mode?: T;
-	type?: 'point' | (string & {});
-	srid?: number;
-}
-
 export function geometry<TMode extends PgGeometryConfig['mode'] & {}>(
 	config?: PgGeometryConfig<TMode>,
 ): Equal<TMode, 'xy'> extends true ? PgGeometryObjectBuilder : PgGeometryBuilder;
@@ -108,7 +107,7 @@ export function geometry<TMode extends PgGeometryConfig['mode'] & {}>(
 export function geometry(a?: string | PgGeometryConfig, b?: PgGeometryConfig) {
 	const { name, config } = getColumnNameAndConfig<PgGeometryConfig>(a, b);
 	if (!config?.mode || config.mode === 'tuple') {
-		return new PgGeometryBuilder(name, config?.srid);
+		return new PgGeometryBuilder(name, config);
 	}
-	return new PgGeometryObjectBuilder(name, config?.srid);
+	return new PgGeometryObjectBuilder(name, config);
 }

drizzle-orm/src/pg-core/columns/postgis_extension/types.ts

@@ -0,0 +1,22 @@
+export type PgGeometryMode = 'tuple' | 'xy';
+
+export type PgGeometryType =
+	| 'Geometry'
+	| 'Point'
+	| 'LineString'
+	| 'Polygon'
+	| 'MultiPoint'
+	| 'MultiLineString'
+	| 'MultiPolygon'
+	| 'GeometryCollection';
+
+export type PgGeometryTypeAnyCase = PgGeometryType | Lowercase<PgGeometryType>;
+
+export interface PgGeometryConfig<
+	T extends PgGeometryMode = PgGeometryMode,
+	G extends PgGeometryType = PgGeometryType,
+> {
+	mode?: T;
+	type?: G | (string & {});
+	srid?: number;
+}

drizzle-orm/src/pg-core/columns/postgis_extension/utils.ts

@@ -1,47 +1,18 @@
-function hexToBytes(hex: string): Uint8Array {
-	const bytes: number[] = [];
-	for (let c = 0; c < hex.length; c += 2) {
-		bytes.push(Number.parseInt(hex.slice(c, c + 2), 16));
-	}
-	return new Uint8Array(bytes);
-}
-
-function bytesToFloat64(bytes: Uint8Array, offset: number): number {
-	const buffer = new ArrayBuffer(8);
-	const view = new DataView(buffer);
-	for (let i = 0; i < 8; i++) {
-		view.setUint8(i, bytes[offset + i]!);
-	}
-	return view.getFloat64(0, true);
-}
+import type { PgGeometryConfig } from './types';
 
-export function parseEWKB(hex: string): { srid: number | undefined; point: [number, number] } {
-	const bytes = hexToBytes(hex);
+export function mapGeometryToDriverValue(value: [number, number], config: PgGeometryConfig): string {
+	let wkt = `point(${value[0]} ${value[1]})`;
 
-	let offset = 0;
-
-	// Byte order: 1 is little-endian, 0 is big-endian
-	const byteOrder = bytes[offset];
-	offset += 1;
-
-	const view = new DataView(bytes.buffer);
-	const geomType = view.getUint32(offset, byteOrder === 1);
-	offset += 4;
-
-	let srid: number | undefined;
-	if (geomType & 0x20000000) { // SRID flag
-		srid = view.getUint32(offset, byteOrder === 1);
-		offset += 4;
+	if (config.srid) {
+		wkt = `SRID=${config.srid};${wkt}`;
 	}
 
-	if ((geomType & 0xFFFF) === 1) {
-		const x = bytesToFloat64(bytes, offset);
-		offset += 8;
-		const y = bytesToFloat64(bytes, offset);
-		offset += 8;
+	return wkt;
+}
 
-		return { srid, point: [x, y] };
-	}
+export function getGeometrySQLType(config: PgGeometryConfig): string {
+	const type = config.type ?? 'Point';
+	const srid = config.srid;
 
-	throw new Error('Unsupported geometry type');
+	return `geometry(${type}${srid ? `,${srid}` : ''})`;
 }