Skip to content

Commit

Permalink
Add codec for postgis.
Browse files Browse the repository at this point in the history
  • Loading branch information
@dnwpark
dnwpark committed Dec 9, 2024
1 parent d7eeed2 commit 4f8c2e1
Show file tree
Hide file tree
Showing 4 changed files with 324 additions and 1 deletion.
80 changes: 80 additions & 0 deletions edgedb-protocol/src/codec.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,10 @@ pub const STD_PG_TIMESTAMPTZ: UuidVal = UuidVal::from_u128(0x1000002);
pub const STD_PG_TIMESTAMP: UuidVal = UuidVal::from_u128(0x1000003);
pub const STD_PG_DATE: UuidVal = UuidVal::from_u128(0x1000004);
pub const STD_PG_INTERVAL: UuidVal = UuidVal::from_u128(0x1000005);
pub const POSTGIS_GEOMETRY: UuidVal = UuidVal::from_u128(0x44c901c0_d922_4894_83c8_061bd05e4840);
pub const POSTGIS_GEOGRAPHY: UuidVal = UuidVal::from_u128(0x4d738878_3a5f_4821_ab76_9d8e7d6b32c4);
pub const POSTGIS_BOX_2D: UuidVal = UuidVal::from_u128(0x7fae5536_6311_4f60_8eb9_096a5d972f48);
pub const POSTGIS_BOX_3D: UuidVal = UuidVal::from_u128(0xc1a50ff8_fded_48b0_85c2_4905a8481433);

pub(crate) fn uuid_to_known_name(uuid: &UuidVal) -> Option<&'static str> {
match *uuid {
Expand Down Expand Up @@ -78,6 +82,10 @@ pub(crate) fn uuid_to_known_name(uuid: &UuidVal) -> Option<&'static str> {
STD_PG_TIMESTAMP => Some("BaseScalar(std::pg::timestamp)"),
STD_PG_DATE => Some("BaseScalar(std::pg::date)"),
STD_PG_INTERVAL => Some("BaseScalar(std::pg::interval)"),
POSTGIS_GEOMETRY => Some("BaseScalar(ext::postgis::geometry)"),
POSTGIS_GEOGRAPHY => Some("BaseScalar(ext::postgis::geography)"),
POSTGIS_BOX_2D => Some("BaseScalar(ext::postgis::box2d)"),
POSTGIS_BOX_3D => Some("BaseScalar(ext::postgis::box3d)"),
_ => None,
}
}
Expand Down Expand Up @@ -257,6 +265,18 @@ pub struct Enum {
members: HashSet<Arc<str>>,
}

#[derive(Debug)]
pub struct PostGisGeometry {}

#[derive(Debug)]
pub struct PostGisGeography {}

#[derive(Debug)]
pub struct PostGisBox2d {}

#[derive(Debug)]
pub struct PostGisBox3d {}

struct CodecBuilder<'a> {
descriptors: &'a [Descriptor],
}
Expand Down Expand Up @@ -376,6 +396,10 @@ pub fn scalar_codec(uuid: &UuidVal) -> Result<Arc<dyn Codec>, CodecError> {
STD_PG_TIMESTAMP => Ok(Arc::new(LocalDatetime {})),
STD_PG_DATE => Ok(Arc::new(LocalDate {})),
STD_PG_INTERVAL => Ok(Arc::new(RelativeDuration {})),
POSTGIS_GEOMETRY => Ok(Arc::new(PostGisGeometry {})),
POSTGIS_GEOGRAPHY => Ok(Arc::new(PostGisGeography {})),
POSTGIS_BOX_2D => Ok(Arc::new(PostGisBox2d {})),
POSTGIS_BOX_3D => Ok(Arc::new(PostGisBox3d {})),
_ => errors::UndefinedBaseScalar { uuid: uuid.clone() }.fail()?,
}
}
Expand Down Expand Up @@ -1484,3 +1508,59 @@ impl Codec for Enum {
Ok(())
}
}

impl Codec for PostGisGeometry {
fn decode(&self, buf: &[u8]) -> Result<Value, DecodeError> {
RawCodec::decode(buf).map(Value::PostGisGeometry)
}
fn encode(&self, buf: &mut BytesMut, val: &Value) -> Result<(), EncodeError> {
let val = match val {
Value::PostGisGeometry(val) => val,
_ => Err(errors::invalid_value(type_name::<Self>(), val))?,
};
buf.extend(val);
Ok(())
}
}

impl Codec for PostGisGeography {
fn decode(&self, buf: &[u8]) -> Result<Value, DecodeError> {
RawCodec::decode(buf).map(Value::PostGisGeography)
}
fn encode(&self, buf: &mut BytesMut, val: &Value) -> Result<(), EncodeError> {
let val = match val {
Value::PostGisGeography(val) => val,
_ => Err(errors::invalid_value(type_name::<Self>(), val))?,
};
buf.extend(val);
Ok(())
}
}

impl Codec for PostGisBox2d {
fn decode(&self, buf: &[u8]) -> Result<Value, DecodeError> {
RawCodec::decode(buf).map(Value::PostGisBox2d)
}
fn encode(&self, buf: &mut BytesMut, val: &Value) -> Result<(), EncodeError> {
let val = match val {
Value::PostGisBox2d(val) => val,
_ => Err(errors::invalid_value(type_name::<Self>(), val))?,
};
buf.extend(val);
Ok(())
}
}

impl Codec for PostGisBox3d {
fn decode(&self, buf: &[u8]) -> Result<Value, DecodeError> {
RawCodec::decode(buf).map(Value::PostGisBox3d)
}
fn encode(&self, buf: &mut BytesMut, val: &Value) -> Result<(), EncodeError> {
let val = match val {
Value::PostGisBox3d(val) => val,
_ => Err(errors::invalid_value(type_name::<Self>(), val))?,
};
buf.extend(val);
Ok(())
}
}
8 changes: 8 additions & 0 deletions edgedb-protocol/src/query_arg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,10 @@ impl QueryArg for Value {
Enum(v) => v.encode_slot(enc)?,
Range(v) => v.encode_slot(enc)?,
Vector(v) => v.encode_slot(enc)?,
PostGisGeometry(v) => v.encode_slot(enc)?,
PostGisGeography(v) => v.encode_slot(enc)?,
PostGisBox2d(v) => v.encode_slot(enc)?,
PostGisBox3d(v) => v.encode_slot(enc)?,
}

Ok(())
Expand Down Expand Up @@ -238,6 +242,10 @@ impl QueryArg for Value {
let val = val.deref();
check_enum(val, &members)
}
(PostGisGeometry(_), BaseScalar(d)) if d.id == codec::POSTGIS_GEOMETRY => Ok(()),
(PostGisGeography(_), BaseScalar(d)) if d.id == codec::POSTGIS_GEOGRAPHY => Ok(()),
(PostGisBox2d(_), BaseScalar(d)) if d.id == codec::POSTGIS_BOX_2D => Ok(()),
(PostGisBox3d(_), BaseScalar(d)) if d.id == codec::POSTGIS_BOX_3D => Ok(()),
// TODO(tailhook) all types
(_, desc) => Err(ctx.wrong_type(&desc, self.kind())),
}
Expand Down
8 changes: 8 additions & 0 deletions edgedb-protocol/src/value.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,10 @@ pub enum Value {
Vector(Vec<f32>),
Enum(EnumValue),
Range(Range<Box<Value>>),
PostGisGeometry(bytes::Bytes),
PostGisGeography(bytes::Bytes),
PostGisBox2d(bytes::Bytes),
PostGisBox3d(bytes::Bytes),
}

#[derive(Clone, Debug)]
Expand Down Expand Up @@ -88,6 +92,10 @@ impl Value {
Tuple(..) => "tuple",
Uuid(..) => "uuid",
Vector(..) => "ext::pgvector::vector",
PostGisGeometry(..) => "ext::postgis::geometry",
PostGisGeography(..) => "ext::postgis::geography",
PostGisBox2d(..) => "ext::postgis::box2d",
PostGisBox3d(..) => "ext::postgis::box3d",
}
}
pub fn empty_tuple() -> Value {
Expand Down
229 changes: 228 additions & 1 deletion edgedb-protocol/tests/codecs.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ extern crate pretty_assertions;
use std::error::Error;
use std::sync::Arc;

use bytes::Bytes;
use bytes::{Buf, Bytes};

use edgedb_protocol::codec::build_codec;
use edgedb_protocol::codec::{Codec, ObjectShape};
Expand Down Expand Up @@ -1201,3 +1201,230 @@ fn multi_range() -> Result<(), Box<dyn Error>> {
);
Ok(())
}

#[test]
fn postgis_geometry() -> Result<(), Box<dyn Error>> {
let codec = build_codec(
Some(TypePos(0)),
&[Descriptor::BaseScalar(BaseScalarTypeDescriptor {
id: "44c901c0-d922-4894-83c8-061bd05e4840"
.parse::<Uuid>()?
.into(),
})],
)?;

encoding_eq!(
&codec,
/*
* Point
* 01 - byteOrder, Little Endian
* 01000000 - wkbType, WKBPoint
* 0000000000000040 - x, 2.0
* 000000000000F03F - y, 1.0
*/
b"\
\x01\
\x01\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
",
Value::PostGisGeometry(
b"\
\x01\
\x01\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
"[..]
.into()
)
);
Ok(())
}

#[test]
fn postgis_geography() -> Result<(), Box<dyn Error>> {
let codec = build_codec(
Some(TypePos(0)),
&[Descriptor::BaseScalar(BaseScalarTypeDescriptor {
id: "4d738878-3a5f-4821-ab76-9d8e7d6b32c4"
.parse::<Uuid>()?
.into(),
})],
)?;
encoding_eq!(
&codec,
/*
* Point
* 01 - byteOrder, Little Endian
* 01000000 - wkbType, WKBPoint
* 0000000000000040 - x, 2.0
* 000000000000F03F - y, 1.0
*/
b"\
\x01\
\x01\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
",
Value::PostGisGeometry(
b"\
\x01\
\x01\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
"[..]
.into()
)
);
Ok(())
}

#[test]
fn postgis_box_2d() -> Result<(), Box<dyn Error>> {
let codec = build_codec(
Some(TypePos(0)),
&[Descriptor::BaseScalar(BaseScalarTypeDescriptor {
id: "7fae5536-6311-4f60-8eb9-096a5d972f48"
.parse::<Uuid>()?
.into(),
})],
)?;
encoding_eq!(
&codec,
/*
* Polygon
* 01 - byteOrder, Little Endian
* 03000000 - wkbType, wkbPolygon
* 01000000 - numRings, 1
* 05000000 - numPoints, 5
* 000000000000F03F - x, 1.0
* 000000000000F03F - y, 1.0
* 0000000000000040 - x, 2.0
* 000000000000F03F - y, 1.0
* 0000000000000040 - x, 2.0
* 0000000000000040 - y, 2.0
* 000000000000F03F - x, 1.0
* 0000000000000040 - y, 2.0
* 000000000000F03F - x, 1.0
* 000000000000F03F - y, 1.0
*/
b"\
\x01\
\x03\x00\x00\x00\
\x01\x00\x00\x00\
\x05\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
",
Value::PostGisBox2d(
b"\
\x01\
\x03\x00\x00\x00\
\x01\x00\x00\x00\
\x05\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
"[..]
.into()
)
);
Ok(())
}

#[test]
fn postgis_box_3d() -> Result<(), Box<dyn Error>> {
let codec = build_codec(
Some(TypePos(0)),
&[Descriptor::BaseScalar(BaseScalarTypeDescriptor {
id: "c1a50ff8-fded-48b0-85c2-4905a8481433"
.parse::<Uuid>()?
.into(),
})],
)?;
encoding_eq!(
&codec,
/*
* Polygon
* 01 - byteOrder, Little Endian
* 03000080 - wkbType, wkbPolygonZ
* 01000000 - numRings, 1
* 05000000 - numPoints, 5
* 000000000000F03F - x, 1.0
* 000000000000F03F - y, 1.0
* 0000000000000000 - z, 0.0
* 0000000000000040 - x, 2.0
* 000000000000F03F - y, 1.0
* 0000000000000000 - z, 0.0
* 0000000000000040 - x, 2.0
* 0000000000000040 - y, 2.0
* 000000000000F03F - x, 1.0
* 0000000000000000 - z, 0.0
* 0000000000000040 - y, 2.0
* 000000000000F03F - x, 1.0
* 000000000000F03F - y, 1.0
* 0000000000000000 - z, 0.0
*/
b"\
\x01\
\x03\x00\x00\x80\
\x01\x00\x00\x00\
\x05\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x00\
",
Value::PostGisBox3d(
b"\
\x01\
\x03\x00\x00\x80\
\x01\x00\x00\x00\
\x05\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x40\
\x00\x00\x00\x00\x00\x00\x00\x00\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\xF0\x3F\
\x00\x00\x00\x00\x00\x00\x00\x00\
"[..]
.into()
)
);
Ok(())
}

0 comments on commit 4f8c2e1

Please sign in to comment.