From b62b3f11172c8a30f38b0bb0ef441f625a850a28 Mon Sep 17 00:00:00 2001 From: Robin Simonklein Date: Sat, 28 Sep 2024 09:00:37 +0200 Subject: [PATCH] [pg] Add implementation of geography type for postgis exension --- drizzle-orm/src/pg-core/columns/index.ts | 1 + .../columns/postgis_extension/geography.ts | 116 ++++++++++++++++++ .../tests/extensions/postgis/pg.test.ts | 42 ++++--- .../tests/extensions/postgis/postgres.test.ts | 42 ++++--- 4 files changed, 169 insertions(+), 32 deletions(-) create mode 100644 drizzle-orm/src/pg-core/columns/postgis_extension/geography.ts diff --git a/drizzle-orm/src/pg-core/columns/index.ts b/drizzle-orm/src/pg-core/columns/index.ts index 881f53e33..8a57abff5 100644 --- a/drizzle-orm/src/pg-core/columns/index.ts +++ b/drizzle-orm/src/pg-core/columns/index.ts @@ -19,6 +19,7 @@ export * from './macaddr8.ts'; export * from './numeric.ts'; export * from './point.ts'; export * from './postgis_extension/geometry.ts'; +export * from './postgis_extension/geography.ts'; export * from './real.ts'; export * from './serial.ts'; export * from './smallint.ts'; diff --git a/drizzle-orm/src/pg-core/columns/postgis_extension/geography.ts b/drizzle-orm/src/pg-core/columns/postgis_extension/geography.ts new file mode 100644 index 000000000..46356f99a --- /dev/null +++ b/drizzle-orm/src/pg-core/columns/postgis_extension/geography.ts @@ -0,0 +1,116 @@ +import type { ColumnBuilderBaseConfig, ColumnBuilderRuntimeConfig, MakeColumnConfig } from '~/column-builder.ts'; +import type { ColumnBaseConfig } from '~/column.ts'; +import { entityKind } from '~/entity.ts'; +import type { AnyPgTable } from '~/pg-core/table.ts'; + +import type { Equal } from '~/utils.ts'; +import { PgColumn, PgColumnBuilder } from '../common.ts'; +import { parseEWKB } from './utils.ts'; + +export type PgGeographyBuilderInitial = PgGeographyBuilder<{ + name: TName; + dataType: 'array'; + columnType: 'PgGeography'; + data: [number, number]; + driverParam: string; + enumValues: undefined; + generated: undefined; +}>; + +export class PgGeographyBuilder> extends PgColumnBuilder { + static readonly [entityKind]: string = 'PgGeographyBuilder'; + + constructor(name: T['name']) { + super(name, 'array', 'PgGeography'); + } + + /** @internal */ + override build( + table: AnyPgTable<{ name: TTableName }>, + ): PgGeography> { + return new PgGeography>( + table, + this.config as ColumnBuilderRuntimeConfig, + ); + } +} + +export class PgGeography> extends PgColumn { + static readonly [entityKind]: string = 'PgGeography'; + + getSQLType(): string { + return 'geography(point)'; + } + + override mapFromDriverValue(value: string): [number, number] { + return parseEWKB(value); + } + + override mapToDriverValue(value: [number, number]): string { + return `point(${value[0]} ${value[1]})`; + } +} + +export type PgGeographyObjectBuilderInitial = PgGeographyObjectBuilder<{ + name: TName; + dataType: 'json'; + columnType: 'PgGeographyObject'; + data: { lon: number; lat: number }; + driverParam: string; + enumValues: undefined; + generated: undefined; +}>; + +export class PgGeographyObjectBuilder> + extends PgColumnBuilder +{ + static readonly [entityKind]: string = 'PgGeographyObjectBuilder'; + + constructor(name: T['name']) { + super(name, 'json', 'PgGeographyObject'); + } + + /** @internal */ + override build( + table: AnyPgTable<{ name: TTableName }>, + ): PgGeographyObject> { + return new PgGeographyObject>( + table, + this.config as ColumnBuilderRuntimeConfig, + ); + } +} + +export class PgGeographyObject> extends PgColumn { + static readonly [entityKind]: string = 'PgGeographyObject'; + + getSQLType(): string { + return 'geography(point)'; + } + + override mapFromDriverValue(value: string): { lon: number; lat: number } { + const parsed = parseEWKB(value); + return { lon: parsed[0], lat: parsed[1] }; + } + + override mapToDriverValue(value: { lon: number; lat: number }): string { + return `point(${value.lon} ${value.lat})`; + } +} + +interface PgGeographyConfig { + mode?: T; + type?: 'point' | (string & {}); +} + +export function geography( + name: TName, + config?: PgGeographyConfig, +): Equal extends true ? PgGeographyObjectBuilderInitial + : PgGeographyBuilderInitial; +export function geography(name: string, config?: PgGeographyConfig) { + if (!config?.mode || config.mode === 'tuple') { + return new PgGeographyBuilder(name); + } + return new PgGeographyObjectBuilder(name); +} diff --git a/integration-tests/tests/extensions/postgis/pg.test.ts b/integration-tests/tests/extensions/postgis/pg.test.ts index 45d6f19b3..e832adc99 100644 --- a/integration-tests/tests/extensions/postgis/pg.test.ts +++ b/integration-tests/tests/extensions/postgis/pg.test.ts @@ -2,7 +2,7 @@ import Docker from 'dockerode'; import { sql } from 'drizzle-orm'; import type { NodePgDatabase } from 'drizzle-orm/node-postgres'; import { drizzle } from 'drizzle-orm/node-postgres'; -import { bigserial, geometry, line, pgTable, point } from 'drizzle-orm/pg-core'; +import { bigserial, geometry, geography, line, pgTable, point } from 'drizzle-orm/pg-core'; import getPort from 'get-port'; import pg from 'pg'; import { v4 as uuid } from 'uuid'; @@ -85,9 +85,11 @@ const items = pgTable('items', { pointObj: point('point_xy', { mode: 'xy' }), line: line('line'), lineObj: line('line_abc', { mode: 'abc' }), - geo: geometry('geo', { type: 'point' }), - geoObj: geometry('geo_obj', { type: 'point', mode: 'xy' }), - geoSrid: geometry('geo_options', { type: 'point', mode: 'xy', srid: 4000 }), + geometry: geometry('geometry', { type: 'point' }), + geometryObj: geometry('geometry_obj', { type: 'point', mode: 'xy' }), + geometrySrid: geometry('geometry_options', { type: 'point', mode: 'xy', srid: 4000 }), + geography: geometry('geography', { type: 'point' }), + geographyObj: geography('geography_obj', { type: 'point', mode: 'json' }), }); beforeEach(async () => { @@ -99,9 +101,11 @@ beforeEach(async () => { "point_xy" point, "line" line, "line_abc" line, - "geo" geometry(point), - "geo_obj" geometry(point), - "geo_options" geometry(point,4000) + "geometry" geometry(point), + "geometry_obj" geometry(point), + "geometry_options" geometry(point,4000), + "geography" geography(point), + "geography_obj" geography(point) ); `); }); @@ -112,9 +116,11 @@ test('insert + select', async () => { pointObj: { x: 1, y: 2 }, line: [1, 2, 3], lineObj: { a: 1, b: 2, c: 3 }, - geo: [1, 2], - geoObj: { x: 1, y: 2 }, - geoSrid: { x: 1, y: 2 }, + geometry: [1, 2], + geometryObj: { x: 1, y: 2 }, + geometrySrid: { x: 1, y: 2 }, + geography: [1, 2], + geographyObj: { lon: 1, lat: 2 }, }]).returning(); const response = await db.select().from(items); @@ -125,9 +131,11 @@ test('insert + select', async () => { pointObj: { x: 1, y: 2 }, line: [1, 2, 3], lineObj: { a: 1, b: 2, c: 3 }, - geo: [1, 2], - geoObj: { x: 1, y: 2 }, - geoSrid: { x: 1, y: 2 }, + geometry: [1, 2], + geometryObj: { x: 1, y: 2 }, + geometrySrid: { x: 1, y: 2 }, + geography: [1, 2], + geographyObj: { lon: 1, lat: 2 }, }]); expect(response).toStrictEqual([{ @@ -136,8 +144,10 @@ test('insert + select', async () => { pointObj: { x: 1, y: 2 }, line: [1, 2, 3], lineObj: { a: 1, b: 2, c: 3 }, - geo: [1, 2], - geoObj: { x: 1, y: 2 }, - geoSrid: { x: 1, y: 2 }, + geometry: [1, 2], + geometryObj: { x: 1, y: 2 }, + geometrySrid: { x: 1, y: 2 }, + geography: [1, 2], + geographyObj: { lon: 1, lat: 2 }, }]); }); diff --git a/integration-tests/tests/extensions/postgis/postgres.test.ts b/integration-tests/tests/extensions/postgis/postgres.test.ts index 071ea2532..9fe27bbdd 100644 --- a/integration-tests/tests/extensions/postgis/postgres.test.ts +++ b/integration-tests/tests/extensions/postgis/postgres.test.ts @@ -1,6 +1,6 @@ import Docker from 'dockerode'; import { sql } from 'drizzle-orm'; -import { bigserial, geometry, line, pgTable, point } from 'drizzle-orm/pg-core'; +import {bigserial, geometry, geography, line, pgTable, point} from 'drizzle-orm/pg-core'; import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js'; import getPort from 'get-port'; import postgres, { type Sql } from 'postgres'; @@ -87,9 +87,11 @@ const items = pgTable('items', { pointObj: point('point_xy', { mode: 'xy' }), line: line('line'), lineObj: line('line_abc', { mode: 'abc' }), - geo: geometry('geo', { type: 'point' }), - geoObj: geometry('geo_obj', { type: 'point', mode: 'xy' }), - geoSrid: geometry('geo_options', { type: 'point', mode: 'xy', srid: 4000 }), + geometry: geometry('geometry', { type: 'point' }), + geometryObj: geometry('geometry_obj', { type: 'point', mode: 'xy' }), + geometrySrid: geometry('geometry_options', { type: 'point', mode: 'xy', srid: 4000 }), + geography: geometry('geography', { type: 'point' }), + geographyObj: geography('geography_obj', { type: 'point', mode: 'json' }), }); beforeEach(async () => { @@ -101,9 +103,11 @@ beforeEach(async () => { "point_xy" point, "line" line, "line_abc" line, - "geo" geometry(point), - "geo_obj" geometry(point), - "geo_options" geometry(point,4000) + "geometry" geometry(point), + "geometry_obj" geometry(point), + "geometry_options" geometry(point,4000), + "geography" geography(point), + "geography_obj" geography(point) ); `); }); @@ -114,9 +118,11 @@ test('insert + select', async () => { pointObj: { x: 1, y: 2 }, line: [1, 2, 3], lineObj: { a: 1, b: 2, c: 3 }, - geo: [1, 2], - geoObj: { x: 1, y: 2 }, - geoSrid: { x: 1, y: 2 }, + geometry: [1, 2], + geometryObj: { x: 1, y: 2 }, + geometrySrid: { x: 1, y: 2 }, + geography: [1, 2], + geographyObj: { lon: 1, lat: 2 }, }]).returning(); const response = await db.select().from(items); @@ -127,9 +133,11 @@ test('insert + select', async () => { pointObj: { x: 1, y: 2 }, line: [1, 2, 3], lineObj: { a: 1, b: 2, c: 3 }, - geo: [1, 2], - geoObj: { x: 1, y: 2 }, - geoSrid: { x: 1, y: 2 }, + geometry: [1, 2], + geometryObj: { x: 1, y: 2 }, + geometrySrid: { x: 1, y: 2 }, + geography: [1, 2], + geographyObj: { lon: 1, lat: 2 }, }]); expect(response).toStrictEqual([{ @@ -138,8 +146,10 @@ test('insert + select', async () => { pointObj: { x: 1, y: 2 }, line: [1, 2, 3], lineObj: { a: 1, b: 2, c: 3 }, - geo: [1, 2], - geoObj: { x: 1, y: 2 }, - geoSrid: { x: 1, y: 2 }, + geometry: [1, 2], + geometryObj: { x: 1, y: 2 }, + geometrySrid: { x: 1, y: 2 }, + geography: [1, 2], + geographyObj: { lon: 1, lat: 2 }, }]); });