Add PackStream v2 support

.gitignore

@@ -14,6 +14,7 @@ __pycache__/
 .Python
 .venv/
 .tox/
+.ruff_cache/
 env/
 build/
 develop-eggs/

bin/target_driver.sh

@@ -12,6 +12,7 @@ git checkout "$version"
 git pull origin "$version"
 cd ..
 cp driver/tests/unit/common/codec/packstream/v1/test_packstream.py tests/codec/packstream/v1/from_driver/test_packstream.py
+cp driver/tests/unit/common/codec/packstream/v2/test_packstream.py tests/codec/packstream/v2/from_driver/test_packstream.py
 cp driver/tests/unit/common/codec/packstream/test_structure.py tests/codec/packstream/from_driver/test_structure.py
 cp -r driver/tests/unit/common/vector/* tests/vector/from_driver
 

changelog.d/86.feature.md

@@ -0,0 +1,2 @@
+Add support for PackStream v2<ISSUES_LIST>.  
+This is required to support Bolt 6.1 which introduces `UUID`s.

src/codec/packstream.rs

@@ -13,7 +13,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod uuid;
 mod v1;
+mod v2;
 
 use pyo3::basic::CompareOp;
 use pyo3::exceptions::{PyIndexError, PyValueError};
@@ -33,6 +35,10 @@ pub(super) fn init_module(m: &Bound<PyModule>, name: &str) -> PyResult<()> {
     m.add_submodule(&mod_v1)?;
     v1::init_module(&mod_v1, format!("{name}.v1").as_str())?;
 
+    let mod_v2 = PyModule::new(py, "v2")?;
+    m.add_submodule(&mod_v2)?;
+    v2::init_module(&mod_v2, format!("{name}.v2").as_str())?;
+
     m.add_class::<Structure>()?;
 
     Ok(())

src/codec/packstream/uuid.rs

@@ -0,0 +1,23 @@
+// Copyright (c) "Neo4j"
+// Neo4j Sweden AB [https://neo4j.com]
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use pyo3::prelude::*;
+use pyo3::sync::PyOnceLock;
+use pyo3::types::PyType;
+
+pub(crate) fn get_uuid_cls(py: Python<'_>) -> PyResult<&Bound<'_, PyType>> {
+    static UUID_CLS: PyOnceLock<Py<PyType>> = PyOnceLock::new();
+    UUID_CLS.import(py, "uuid", "UUID")
+}

src/codec/packstream/v1.rs

@@ -13,10 +13,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod pack;
-mod unpack;
+pub(super) mod extension;
+pub(super) mod pack;
+pub(super) mod unpack;
 
 use pyo3::prelude::*;
+use pyo3::types::{PyByteArray, PyBytes, PyDict};
 use pyo3::wrap_pyfunction;
 
 use crate::register_package;
@@ -46,11 +48,30 @@ const BYTES_8: u8 = 0xCC;
 const BYTES_16: u8 = 0xCD;
 const BYTES_32: u8 = 0xCE;
 
+#[pyfunction]
+#[pyo3(name = "pack", signature = (value, dehydration_hooks=None))]
+fn pack_fn<'py>(
+    value: &Bound<'py, PyAny>,
+    dehydration_hooks: Option<&Bound<'py, PyAny>>,
+) -> PyResult<Bound<'py, PyBytes>> {
+    pack::pack::<extension::PackStreamV1BaseExt>(value, dehydration_hooks)
+}
+
+#[pyfunction]
+#[pyo3(name = "unpack", signature = (bytes, idx, hydration_hooks=None))]
+fn unpack_fn(
+    bytes: Bound<PyByteArray>,
+    idx: usize,
+    hydration_hooks: Option<Bound<PyDict>>,
+) -> PyResult<(Py<PyAny>, usize)> {
+    unpack::unpack::<extension::PackStreamV1BaseExt>(bytes, idx, hydration_hooks)
+}
+
 pub(crate) fn init_module(m: &Bound<PyModule>, name: &str) -> PyResult<()> {
     register_package(m, name)?;
 
-    m.add_function(wrap_pyfunction!(unpack::unpack, m)?)?;
-    m.add_function(wrap_pyfunction!(pack::pack, m)?)?;
+    m.add_function(wrap_pyfunction!(unpack_fn, m)?)?;
+    m.add_function(wrap_pyfunction!(pack_fn, m)?)?;
 
     Ok(())
 }

src/codec/packstream/v1/extension.rs

@@ -0,0 +1,63 @@
+// Copyright (c) "Neo4j"
+// Neo4j Sweden AB [https://neo4j.com]
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::ffi::CStr;
+
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+
+use super::super::uuid::get_uuid_cls;
+use super::pack::PackStreamEncoder;
+use super::unpack::PackStreamDecoder;
+
+pub(crate) trait PackStreamV1Ext: Sized {
+    fn type_mapping_import() -> &'static CStr;
+    fn pack_ext(
+        value: &'_ Bound<PyAny>,
+        encoder: &mut PackStreamEncoder<'_, Self>,
+    ) -> PyResult<bool>;
+    fn unpack_ext(marker: u8, decoder: &mut PackStreamDecoder<Self>)
+        -> PyResult<Option<Py<PyAny>>>;
+}
+
+pub(crate) struct PackStreamV1BaseExt {}
+
+impl PackStreamV1Ext for PackStreamV1BaseExt {
+    #[inline]
+    fn type_mapping_import() -> &'static CStr {
+        c"from neo4j._codec.packstream.v1.types import *"
+    }
+
+    #[inline]
+    fn pack_ext(value: &'_ Bound<PyAny>, _: &mut PackStreamEncoder<'_, Self>) -> PyResult<bool> {
+        let py = value.py();
+
+        let uuid_cls = get_uuid_cls(py)?;
+        if value.is_instance(uuid_cls)? {
+            return Err(PyErr::new::<PyValueError, _>(format!(
+                "Values of type {} are not supported \
+                (requires Bolt protocol version 6.1 or newer)",
+                value.get_type().str()?
+            )));
+        }
+
+        Ok(false)
+    }
+
+    #[inline]
+    fn unpack_ext(_: u8, _: &mut PackStreamDecoder<Self>) -> PyResult<Option<Py<PyAny>>> {
+        Ok(None)
+    }
+}

src/codec/packstream/v1/pack.rs

@@ -14,6 +14,7 @@
 // limitations under the License.
 
 use std::borrow::Cow;
+use std::marker::PhantomData;
 use std::sync::OnceLock;
 
 use pyo3::exceptions::{PyOverflowError, PyTypeError, PyValueError};
@@ -24,6 +25,7 @@ use pyo3::types::{PyByteArray, PyBytes, PyDict, PyString, PyTuple, PyType};
 use pyo3::{intern, IntoPyObjectExt};
 
 use super::super::Structure;
+use super::extension::PackStreamV1Ext;
 use super::{
     BYTES_16, BYTES_32, BYTES_8, FALSE, FLOAT_64, INT_16, INT_32, INT_64, INT_8, LIST_16, LIST_32,
     LIST_8, MAP_16, MAP_32, MAP_8, NULL, STRING_16, STRING_32, STRING_8, TINY_LIST, TINY_MAP,
@@ -130,46 +132,42 @@ impl TypeMappings {
     }
 }
 
-static TYPE_MAPPINGS: OnceLock<PyResult<TypeMappings>> = OnceLock::new();
+fn get_type_mappings<E: PackStreamV1Ext>(py: Python<'_>) -> PyResult<&'static TypeMappings> {
+    static TYPE_MAPPINGS: OnceLock<PyResult<TypeMappings>> = OnceLock::new();
 
-fn get_type_mappings(py: Python<'_>) -> PyResult<&'static TypeMappings> {
     let mappings = TYPE_MAPPINGS.get_or_init_py_attached(py, || {
         let locals = PyDict::new(py);
-        py.run(
-            c"from neo4j._codec.packstream.v1.types import *",
-            None,
-            Some(&locals),
-        )?;
+        py.run(E::type_mapping_import(), None, Some(&locals))?;
         TypeMappings::new(&locals)
     });
     mappings.as_ref().map_err(|e| e.clone_ref(py))
 }
 
-#[pyfunction]
-#[pyo3(signature = (value, dehydration_hooks=None))]
-pub(super) fn pack<'py>(
+pub(crate) fn pack<'py, E: PackStreamV1Ext>(
     value: &Bound<'py, PyAny>,
     dehydration_hooks: Option<&Bound<'py, PyAny>>,
 ) -> PyResult<Bound<'py, PyBytes>> {
     let py = value.py();
-    let type_mappings = get_type_mappings(py)?;
-    let mut encoder = PackStreamEncoder::new(dehydration_hooks, type_mappings);
+    let type_mappings = get_type_mappings::<E>(py)?;
+    let mut encoder = PackStreamEncoder::<E>::new(dehydration_hooks, type_mappings);
     encoder.write(value)?;
     Ok(PyBytes::new(py, &encoder.buffer))
 }
 
-struct PackStreamEncoder<'a> {
+pub(crate) struct PackStreamEncoder<'a, E: PackStreamV1Ext> {
+    ext: PhantomData<E>,
     dehydration_hooks: Option<&'a Bound<'a, PyAny>>,
     type_mappings: &'a TypeMappings,
     buffer: Vec<u8>,
 }
 
-impl<'a> PackStreamEncoder<'a> {
+impl<'a, E: PackStreamV1Ext> PackStreamEncoder<'a, E> {
     fn new(
         dehydration_hooks: Option<&'a Bound<'a, PyAny>>,
         type_mappings: &'a TypeMappings,
     ) -> Self {
         Self {
+            ext: PhantomData,
             dehydration_hooks,
             type_mappings,
             buffer: Default::default(),
@@ -243,6 +241,10 @@ impl<'a> PackStreamEncoder<'a> {
             });
         }
 
+        if E::pack_ext(value, self)? {
+            return Ok(());
+        }
+
         if let Ok(value) = value.extract::<Bound<Structure>>() {
             let value_ref = value.borrow();
             let size = value_ref.fields.len().try_into().map_err(|_| {
@@ -407,4 +409,9 @@ impl<'a> PackStreamEncoder<'a> {
         self.buffer.extend(&[TINY_STRUCT + size, tag]);
         Ok(())
     }
+
+    #[inline]
+    pub(crate) fn write_raw<'b, I: IntoIterator<Item = &'b u8>>(&mut self, iter: I) {
+        self.buffer.extend(iter)
+    }
 }

src/codec/packstream/v1/unpack.rs

@@ -13,47 +13,50 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::marker::PhantomData;
+
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::sync::critical_section::with_critical_section;
 use pyo3::types::{IntoPyDict, PyByteArray, PyBytes, PyDict, PyList, PyTuple};
 use pyo3::{intern, IntoPyObjectExt};
 
 use super::super::Structure;
+use super::extension::PackStreamV1Ext;
 use super::{
     BYTES_16, BYTES_32, BYTES_8, FALSE, FLOAT_64, INT_16, INT_32, INT_64, INT_8, LIST_16, LIST_32,
     LIST_8, MAP_16, MAP_32, MAP_8, NULL, STRING_16, STRING_32, STRING_8, TINY_LIST, TINY_MAP,
     TINY_STRING, TINY_STRUCT, TRUE,
 };
 
-#[pyfunction]
-#[pyo3(signature = (bytes, idx, hydration_hooks=None))]
-pub(super) fn unpack(
+pub(crate) fn unpack<E: PackStreamV1Ext>(
     bytes: Bound<PyByteArray>,
     idx: usize,
     hydration_hooks: Option<Bound<PyDict>>,
 ) -> PyResult<(Py<PyAny>, usize)> {
     let py = bytes.py();
-    let mut decoder = PackStreamDecoder::new(py, bytes, idx, hydration_hooks);
+    let mut decoder = PackStreamDecoder::<E>::new(py, bytes, idx, hydration_hooks);
     let result = decoder.read()?;
     Ok((result, decoder.index))
 }
 
-struct PackStreamDecoder<'a> {
+pub(crate) struct PackStreamDecoder<'a, E: PackStreamV1Ext> {
+    ext: PhantomData<E>,
     py: Python<'a>,
     bytes: Bound<'a, PyByteArray>,
     index: usize,
     hydration_hooks: Option<Bound<'a, PyDict>>,
 }
 
-impl<'a> PackStreamDecoder<'a> {
+impl<'a, E: PackStreamV1Ext> PackStreamDecoder<'a, E> {
     fn new(
         py: Python<'a>,
         bytes: Bound<'a, PyByteArray>,
         idx: usize,
         hydration_hooks: Option<Bound<'a, PyDict>>,
     ) -> Self {
         Self {
+            ext: PhantomData,
             py,
             bytes,
             index: idx,
@@ -133,10 +136,13 @@ impl<'a> PackStreamDecoder<'a> {
             }
             _ if high_nibble == TINY_STRUCT => self.read_struct((marker & 0x0F).into())?,
             _ => {
-                // raise ValueError("Unknown PackStream marker %02X" % marker)
-                return Err(PyErr::new::<PyValueError, _>(format!(
-                    "Unknown PackStream marker {marker:02X}",
-                )));
+                let Some(value) = E::unpack_ext(marker, self)? else {
+                    // raise ValueError("Unknown PackStream marker %02X" % marker)
+                    return Err(PyErr::new::<PyValueError, _>(format!(
+                        "Unknown PackStream marker {marker:02X}",
+                    )));
+                };
+                value
             }
         })
     }
@@ -262,7 +268,7 @@ impl<'a> PackStreamDecoder<'a> {
         Ok(byte)
     }
 
-    fn read_n_bytes<const N: usize>(&mut self) -> PyResult<[u8; N]> {
+    pub(crate) fn read_n_bytes<const N: usize>(&mut self) -> PyResult<[u8; N]> {
         let to = self.index + N;
         with_critical_section(&self.bytes, || {
             // Safety:
@@ -320,4 +326,9 @@ impl<'a> PackStreamDecoder<'a> {
     fn read_f64(&mut self) -> PyResult<f64> {
         self.read_n_bytes().map(f64::from_be_bytes)
     }
+
+    #[inline]
+    pub(crate) fn py(&self) -> Python<'a> {
+        self.py
+    }
 }