Use tagged types for DB Geometry types

arttuka · Jan 4, 2025 · 0612627 · 0612627
1 parent de3f907
commit 0612627
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 36 deletions.
diff --git a/resources/sql/init-db.sql b/resources/sql/init-db.sql
@@ -27,9 +27,9 @@ LANGUAGE plpgsql
 IMMUTABLE
 RETURNS NULL ON NULL INPUT;
 
-CREATE OR REPLACE FUNCTION AsJSON(g geometry) RETURNS text
+CREATE OR REPLACE FUNCTION AsJSON(g geometry) RETURNS jsonb
 AS $$
-  SELECT ST_AsGeoJSON(ST_Transform(g, 4326), 6)
+  SELECT ST_AsGeoJSON(ST_Transform(g, 4326), 6)::jsonb
 $$
 LANGUAGE SQL
 IMMUTABLE

diff --git a/src/server/db/db.ts b/src/server/db/db.ts
@@ -14,8 +14,9 @@ import config from '../config'
 import { Lane, Route, RouteType, Waypoint } from '../../common/types'
 import { Database, PgrDijkstra } from './types'
 import {
-  aggBuilder,
+  arrayAgg,
   asJSON,
+  distance,
   extendKysely,
   hydrate,
   makeLine,
@@ -107,9 +108,7 @@ const insertEndpoints = async (
                 .as('point'),
             ])
             .$call(whereNullOrGreater('depth', depth))
-            .orderBy(({ eb, ref }) =>
-              eb('lane.geom', '<->', ref('origin.geom'))
-            )
+            .orderBy((eb) => distance(eb, 'lane.geom', 'origin.geom'))
             .limit(sql.lit(1))
             .as('l')
         )
@@ -153,14 +152,11 @@ const insertExtraLanes = async (tx: Transaction<Database>): Promise<void> => {
     .with('lane_endpoints', (db) =>
       db
         .selectFrom('endpoint')
-        .select(({ fn }) => {
-          const arrayAgg = aggBuilder(fn, 'array_agg')
-          return [
-            'lane as laneid',
-            arrayAgg<number[]>(['vertex'], 'seq').as('vertex_ids'),
-            arrayAgg<string[]>(['point'], 'seq').as('points'),
-          ]
-        })
+        .select(({ fn }) => [
+          'lane as laneid',
+          arrayAgg(fn, ['vertex'], 'seq').as('vertex_ids'),
+          arrayAgg(fn, ['point'], 'seq').as('points'),
+        ])
         .where('type', '!=', sql.lit('viadirect' as const))
         .where('vertex', '<', sql.lit(0))
         .groupBy(['lane'])

diff --git a/src/server/db/types.ts b/src/server/db/types.ts
@@ -13,6 +13,11 @@ import {
 import { joinList, getSqlType } from './util'
 import { WaypointType } from '../../common/types'
 
+declare const tags: unique symbol
+export type DbPoint = string & { [tags]: { Point: void } }
+export type DbLineString = string & { [tags]: { LineString: void } }
+export type DbGeometry = DbPoint | DbLineString
+
 export type Database = {
   lane: LaneTable
   lane_vertices_pgr: LaneVerticesPgrTable
@@ -28,7 +33,7 @@ export type BaseLaneTable = {
   length: number
   depth: number
   height: number
-  geom: string
+  geom: DbLineString
 }
 
 export type LaneTable = BaseLaneTable & {
@@ -40,7 +45,7 @@ export type LaneTable = BaseLaneTable & {
 
 export type LaneVerticesPgrTable = {
   id: Generated<number>
-  the_geom: string
+  the_geom: DbPoint
 }
 
 export type ExtraLaneTable = BaseLaneTable & {
@@ -51,19 +56,19 @@ export type EndpointTable = {
   seq: number
   lane: number
   vertex: number
-  geometry: string
-  point: string
+  geometry: DbLineString
+  point: DbPoint
   type: WaypointType
 }
 
 export type SegmentView = {
   seq: number
   source: number
   target: number
-  source_point: string
-  target_point: string
-  source_geometry: string
-  target_geometry: string
+  source_point: DbPoint
+  target_point: DbPoint
+  source_geometry: DbLineString
+  target_geometry: DbLineString
   direct: boolean
   type: WaypointType
 }
@@ -73,7 +78,7 @@ export type SplitLinestring = {
   source: number
   target: number
   length: number
-  geom: string
+  geom: DbLineString
 }
 
 export type PgrDijkstra = {

diff --git a/src/server/db/util.ts b/src/server/db/util.ts
@@ -4,6 +4,7 @@ import {
   Expression,
   ExpressionBuilder,
   ExtractTypeFromReferenceExpression,
+  ExtractTypeFromStringReference,
   FunctionModule,
   Kysely,
   SelectQueryBuilder,
@@ -16,9 +17,13 @@ import {
 } from 'kysely'
 import { LineString } from 'geojson'
 import {
+  DbGeometry,
+  DbLineString,
+  DbPoint,
   PgrDijkstra,
   SplitLinestring,
   TypedReferenceExpression,
+  TypedStringReference,
   ValuesExpression,
 } from './types'
 import { NotEmptyArray } from '../../common/types'
@@ -69,17 +74,36 @@ export const aggBuilder =
     dir: 'asc' | 'desc' = 'asc'
   ) => {
     const [firstArgs, lastArg] = splitLast(args)
-    return fn.agg<O, RE>(name, [
+    return fn.agg<O>(name, [
       ...firstArgs,
-      sql`${parseAggExpression(lastArg)} ORDER BY ${parseAggExpression(orderBy)} ${sql.raw(dir)}` as unknown as RE,
+      sql`${parseAggExpression(lastArg)} ORDER BY ${parseAggExpression(orderBy)} ${sql.raw(dir)}`,
     ])
   }
 
+type ExtractType<Db, Tb extends keyof Db, E extends AggExpression<Db, Tb>> =
+  E extends StringReference<Db, Tb>
+    ? ExtractTypeFromStringReference<Db, Tb, E>
+    : E extends Expression<infer T>
+      ? T
+      : never
+
+export const arrayAgg = <
+  Db,
+  Tb extends keyof Db,
+  RE extends AggExpression<Db, Tb> = AggExpression<Db, Tb>,
+  OE extends AggExpression<Db, Tb> = AggExpression<Db, Tb>,
+>(
+  fn: FunctionModule<Db, Tb>,
+  args: Readonly<NotEmptyArray<RE>>,
+  orderBy: OE,
+  dir: 'asc' | 'desc' = 'asc'
+) => aggBuilder(fn, 'array_agg')<ExtractType<Db, Tb, RE>[]>(args, orderBy, dir)
+
 export const splitLinestring = <Db, Tb extends keyof Db>(
   eb: ExpressionBuilder<Db, Tb>,
   vertexIds: TypedReferenceExpression<Db, Tb, number[]>,
-  points: TypedReferenceExpression<Db, Tb, string[]>,
-  geom: TypedReferenceExpression<Db, Tb, string>,
+  points: TypedReferenceExpression<Db, Tb, DbPoint[]>,
+  geom: TypedReferenceExpression<Db, Tb, DbLineString>,
   source: TypedReferenceExpression<Db, Tb, number>,
   target: TypedReferenceExpression<Db, Tb, number>,
   length: TypedReferenceExpression<Db, Tb, number>
@@ -111,6 +135,12 @@ export const hydrate = <T extends Compilable>(qb: T) => {
   })
 }
 
+export const distance = <Db, Tb extends keyof Db>(
+  eb: ExpressionBuilder<Db, Tb>,
+  from: TypedStringReference<Db, Tb, DbGeometry>,
+  to: TypedStringReference<Db, Tb, DbGeometry>
+) => eb(from, '<->', eb.ref(to)).$castTo<number>()
+
 export const logQuery = <T extends Compilable>(qb: T): T => {
   console.log(qb.compile())
   return qb
@@ -123,26 +153,26 @@ export const logHydratedQuery = <T extends Compilable>(qb: T): T => {
 
 export const asJSON = <Db, Tb extends keyof Db>(
   eb: ExpressionBuilder<Db, Tb>,
-  expr: TypedReferenceExpression<Db, Tb, string>
-) => eb.cast<LineString>(eb.fn('AsJSON', [expr]), 'jsonb')
+  expr: TypedReferenceExpression<Db, Tb, DbLineString>
+) => eb.fn<LineString>('AsJSON', [expr])
 
-export const makeLine = <Db, Tb extends keyof Db>(
+export const makeLine = <Db, Tb extends keyof Db, T extends DbGeometry>(
   eb: ExpressionBuilder<Db, Tb>,
-  ...exprs: TypedReferenceExpression<Db, Tb, string>[]
-) => asJSON(eb, eb.fn('ST_MakeLine', exprs))
+  ...exprs: TypedReferenceExpression<Db, Tb, T>[]
+) => asJSON(eb, eb.fn<DbLineString>('ST_MakeLine', exprs))
 
-export const makeLineAgg = <Db, Tb extends keyof Db>(
+export const makeLineAgg = <Db, Tb extends keyof Db, T extends DbGeometry>(
   eb: ExpressionBuilder<Db, Tb>,
-  expr: TypedReferenceExpression<Db, Tb, string>,
+  expr: TypedReferenceExpression<Db, Tb, T>,
   orderBy: AggExpression<Db, Tb>
-) => asJSON(eb, aggBuilder(eb.fn, 'ST_MakeLine')([expr], orderBy))
+) => asJSON(eb, aggBuilder(eb.fn, 'ST_MakeLine')<DbLineString>([expr], orderBy))
 
 export const makePoint = <Db, Tb extends keyof Db>(
   eb: ExpressionBuilder<Db, Tb>,
   lng: TypedReferenceExpression<Db, Tb, number>,
   lat: TypedReferenceExpression<Db, Tb, number>
 ) =>
-  eb.fn<string>('ST_Transform', [
+  eb.fn<DbPoint>('ST_Transform', [
     eb.fn('ST_SetSRID', [eb.fn('ST_MakePoint', [lng, lat]), eb.lit(4326)]),
     eb.lit(3067),
   ])