Experimental PostGIS geometry GeoJSON column type for Drizzle ORM

postgisGeometry.ts

import { sql } from "drizzle-orm";
import { customType, CustomTypeValues } from "drizzle-orm/pg-core";

/**
 * Experimental custom type for PostGIS geometry, only supports reads
 */
export const geometry = <
  TType extends GeoJSON.Geometry["type"] = GeoJSON.Geometry["type"],
  T extends CustomTypeValues = CustomTypeValues,
>(
  dbName: string,
  fieldConfig?: T["config"] & { type: TType },
) => {
  const type = fieldConfig?.type;
  return customType<{
    data: GeometryTypes[TType];
  }>({
    dataType() {
      return type ? `geometry(${type},4326)` : "geometry";
    },
    toDriver(value) {
      return sql`ST_GeomFromGeoJSON(${JSON.stringify(value)})`;
    },
    fromDriver(value) {
      const val = value as string;

      // Detect if hex string
      if (!val.startsWith("{")) {
        return parseHexToGeometry(val) as GeometryTypes[TType];
      } else {
        try {
          const data = JSON.parse(value as string);

          if (type && data.type !== type) {
            throw new Error(`Expected geometry type ${type}, got ${data.type}`);
          }

          return data as GeometryTypes[TType];
        } catch (e) {
          throw new Error(`Failed to parse geometry`, {
            cause: e,
          });
        }
      }
    },
  })(dbName, fieldConfig);
};

export enum GeometryType {
  Point = 1,
  LineString = 2,
  Polygon = 3,
  MultiPoint = 4,
  MultiLineString = 5,
  MultiPolygon = 6,
  GeometryCollection = 7,
}

export const parseHexToGeometry = (hex: string): GeoJSON.Geometry => {
  const startMs = performance.now();
  try {
    const byteStr = hex.match(/.{1,2}/g)?.map((byte) => parseInt(byte, 16));
    if (!byteStr) {
      throw new Error("Failed to convert hex to buffer");
    }

    const uint8Array = new Uint8Array(byteStr);
    const buffer = uint8Array.buffer;

    const dataView = new DataView(buffer);

    let byteOffset = 0;

    const byteOrder = dataView.getUint8(0); // 1 byte
    byteOffset += 1; // Move the byte offset past the byte order field

    const littleEndian = byteOrder === 1;

    let geometryType = dataView.getUint32(1, littleEndian); // 4 bytes
    byteOffset += 4; // Move the byte offset past the geometry type field

    const hasSRID = (geometryType & 0x20000000) > 0; // Check if the SRID flag is set

    let srid;

    if (hasSRID) {
      // If SRID is included, read the SRID
      srid = dataView.getUint32(byteOffset, littleEndian);
      // Set geometry type to the actual type, stripping the SRID flag
      geometryType &= ~0x20000000;
      byteOffset += 4; // Move the byte offset past the SRID field
    }

    const geometry = parseGeometry(
      dataView,
      littleEndian,
      geometryType,
      byteOffset,
    );

    return geometry;
  } catch (e) {
    throw e;
  } finally {
    const endMs = performance.now();
    console.debug(
      "[Postgis Geometry] parseHexToGeometry took",
      endMs - startMs,
      "ms",
    );
  }
};

function readPoint(
  dataView: DataView,
  littleEndian: boolean,
  offset: number,
): GeoJSON.Position {
  const x = dataView.getFloat64(offset, littleEndian);
  const y = dataView.getFloat64(offset + 8, littleEndian);
  return [x, y];
}

function readLineString(
  dataView: DataView,
  littleEndian: boolean,
  offset: number,
): GeoJSON.Position[] {
  const numPoints = dataView.getUint32(1, littleEndian);
  const points: GeoJSON.Position[] = [];
  for (let i = 0; i < numPoints; i++) {
    const x = dataView.getFloat64(offset, littleEndian);
    const y = dataView.getFloat64(offset + 8, littleEndian);
    points.push([x, y]);
    offset += 16;
  }
  return points;
}

function readPolygon(
  dataView: DataView,
  littleEndian: boolean,
  offset: number,
): GeoJSON.Position[][] {
  const numRings = dataView.getUint32(1, littleEndian);
  const rings: GeoJSON.Position[][] = [];
  for (let i = 0; i < numRings; i++) {
    const numPoints = dataView.getUint32(offset, littleEndian);
    offset += 4;
    const points: GeoJSON.Position[] = [];
    for (let j = 0; j < numPoints; j++) {
      const x = dataView.getFloat64(offset, littleEndian);
      const y = dataView.getFloat64(offset + 8, littleEndian);
      points.push([x, y]);
      offset += 16;
    }
    rings.push(points);
  }
  return rings;
}

export const parseGeometry = (
  dataView: DataView,
  littleEndian: boolean,
  type: GeometryType,
  offset: number,
): GeoJSON.Geometry => {
  switch (type) {
    case GeometryType.Point:
      return {
        type: "Point",
        coordinates: readPoint(
          dataView,
          littleEndian,
          offset,
        ) as GeoJSON.Point["coordinates"],
      };
    case GeometryType.LineString:
      return {
        type: "LineString",
        coordinates: readLineString(
          dataView,
          littleEndian,
          offset,
        ) as GeoJSON.LineString["coordinates"],
      };
    case GeometryType.Polygon:
      return {
        type: "Polygon",
        coordinates: readPolygon(
          dataView,
          littleEndian,
          offset,
        ) as GeoJSON.Polygon["coordinates"],
      };
    case GeometryType.MultiPoint:
      return {
        type: "MultiPoint",
        coordinates: readLineString(
          dataView,
          littleEndian,
          offset,
        ) as GeoJSON.MultiPoint["coordinates"],
      };
    case GeometryType.MultiLineString:
      return {
        type: "MultiLineString",
        coordinates: readPolygon(
          dataView,
          littleEndian,
          offset,
        ) as GeoJSON.MultiLineString["coordinates"],
      };
    default:
      throw new Error("Unsupported geometry type");
  }
};

export type GeometryTypes = {
  Point: GeoJSON.Point;
  LineString: GeoJSON.LineString;
  Polygon: GeoJSON.Polygon;
  MultiPoint: GeoJSON.MultiPoint;
  MultiLineString: GeoJSON.MultiLineString;
  MultiPolygon: GeoJSON.MultiPolygon;
  GeometryCollection: GeoJSON.GeometryCollection;
};

Once drizzle-team/drizzle-orm#1423 lands this should become a lot more simple and faster.

For reference, adding conditions like ST_DWithin shoudn't be too hard:

https://github.com/drizzle-team/drizzle-orm/blob/main/drizzle-orm/src/sql/expressions/conditions.ts

Hi @ItsWendell

ESLint says you never used srid variable. Is it necessary?

Nice experiment! Thanks! 🙏

On my M1 Mac (w/ an Apple Silicon chip), this crashes however with Offset is outside the bounds of the DataView

[Postgis Geometry] parseHexToGeometry took 2.8550829999148846 ms
 ⨯ ../packages/drizzle/lib/geography.ts (114:23) @ getFloat64
 ⨯ RangeError: Offset is outside the bounds of the DataView
    at DataView.getFloat64 (<anonymous>)
    at readLineString (../../packages/drizzle/lib/geography.ts:99:28)
    at parseGeometry (../../packages/drizzle/lib/geography.ts:139:30)
    at parseHexToGeometry (../../packages/drizzle/lib/geography.ts:80:26)

Something must be wrong in the parsing because I adapted your code to use https://github.com/syncpoint/wkx for the parsing and it worked like a charm (well the wkx code is not a ES module... had to hack around a bit).

Something must be wrong in the parsing because I adapted your code to use https://github.com/syncpoint/wkx for the parsing and it worked like a charm (well the wkx code is not a ES module... had to hack around a bit).

FWIW, here is my working code (reads & writes) using the modified wkx library.

import { sql } from "drizzle-orm";
import type { CustomTypeValues } from "drizzle-orm/pg-core";
import { customType } from "drizzle-orm/pg-core";
import type {
  Geometry,
  GeometryCollection,
  LineString,
  MultiLineString,
  MultiPoint,
  MultiPolygon,
  Point,
  Polygon,
} from "geojson";
import { Geometry as WkxGeometry } from "@local-repo/wkx";

interface GeometryTypes {
  Point: Point;
  LineString: LineString;
  Polygon: Polygon;
  MultiPoint: MultiPoint;
  MultiLineString: MultiLineString;
  MultiPolygon: MultiPolygon;
  GeometryCollection: GeometryCollection;
}
/**
 * Custom type for PostGIS geography...
 * 
 * Inspired by & adapted from:
 * https://gist.github.com/ItsWendell/65a107c0f8de00fdd6e4f3935789ce5a
 * https://gist.github.com/ItsWendell/38ebe96b34d00d9138ce23cc363d7009
 */
export const geography = <
  TType extends Geometry["type"] = Geometry["type"],
  T extends CustomTypeValues = CustomTypeValues,
>(
  dbName: string,
  fieldConfig?: T["config"] & { type: TType },
) => {
  const type = fieldConfig?.type;
  return customType<{
    data: GeometryTypes[TType];
  }>({
    dataType() {
      return type ? `geography(${type},4326)` : "geography";
    },
    toDriver(value) {
      return sql`ST_GeomFromGeoJSON(${JSON.stringify(value)})`;
    },
    fromDriver(value) {
      const buffer = Buffer.from(value as string, "hex");
      const geometry = WkxGeometry.parse(buffer);
      const geoJson = geometry.toGeoJSON() as GeometryTypes[TType];
      return geoJson;
    },
  })(dbName, fieldConfig);
};

The modified wkx repo is here: https://github.com/janvorwerk/wkx

@ItsWendell did you ever get to implement a readMultiPolygon function?

Thanks for your solution @janvorwerk. Also, it worked for me without a modified wkx.

Thanks for your solution @janvorwerk. Also, it worked for me without a modified wkx.

@bernatfortet you mean that it worked with the original wkx or with the parser in this gist? In the latter case, what kind of CPU have you got? x86/amd64? Or ARM as in Apple M1/M2/M3 ?

It worked with the original wkx.
I have an ARM M1Max