pybind/rados: enable alternate types for key names

Specifically, allow the caller to select the Python bytes type as an
acceptable key type.

This is important when the omap keys contain data that cannot be decoded
as a UTF-8 string. Encountering such a key currently stops the iteration
and makes the key virtually inaccessible by pyrados.

Fixes: https://tracker.ceph.com/issues/59716
Signed-off-by: Patrick Donnelly <pdonnell@redhat.com>

1 files changed, 25 insertions, 14 deletions

diff --git a/src/pybind/rados/rados.pyx b/src/pybind/rados/rados.pyx
index d7b68b4ac90..15198f6e642 100644
--- a/src/pybind/rados/rados.pyx
+++ b/src/pybind/rados/rados.pyx
@@ -76,6 +76,7 @@ MAX_ERRNO = _MAX_ERRNO
 ANONYMOUS_AUID = 0xffffffffffffffff
 ADMIN_AUID = 0
 
+OMAP_KEY_TYPE = Union[str,bytes]
 
 class Error(Exception):
     """ `Error` class, derived from `Exception` """
@@ -1369,10 +1370,12 @@ cdef class OmapIterator(object):
     """Omap iterator"""
 
     cdef public Ioctx ioctx
+    cdef public object omap_key_type
     cdef rados_omap_iter_t ctx
 
-    def __cinit__(self, Ioctx ioctx):
+    def __cinit__(self, Ioctx ioctx, omap_key_type):
         self.ioctx = ioctx
+        self.omap_key_type = omap_key_type
 
     def __iter__(self):
         return self
@@ -1394,7 +1397,7 @@ cdef class OmapIterator(object):
             raise make_ex(ret, "error iterating over the omap")
         if key_ == NULL:
             raise StopIteration()
-        key = decode_cstr(key_)
+        key = self.omap_key_type(key_)
         val = None
         if val_ != NULL:
             val = val_[:len_]
@@ -1929,7 +1932,7 @@ cdef class WriteOp(object):
         with nogil:
             rados_write_op_cmpext(self.write_op, _cmp_buf, _cmp_buf_len, _offset, NULL)
 
-    def omap_cmp(self, key: str, val: str, cmp_op: int = LIBRADOS_CMPXATTR_OP_EQ):
+    def omap_cmp(self, key: OMAP_KEY_TYPE, val: OMAP_KEY_TYPE, cmp_op: int = LIBRADOS_CMPXATTR_OP_EQ):
         """
         Ensure that an omap key value satisfies comparison
         :param key: omap key whose associated value is evaluated for comparison
@@ -3605,7 +3608,7 @@ returned %d, but should return zero on success." % (self.name, ret))
         """
         read_op.release()
 
-    def set_omap(self, write_op: WriteOp, keys: Sequence[str], values: Sequence[bytes]):
+    def set_omap(self, write_op: WriteOp, keys: Sequence[OMAP_KEY_TYPE], values: Sequence[bytes]):
         """
         set keys values to write_op
         :para write_op: write_operation object
@@ -3752,9 +3755,10 @@ returned %d, but should return zero on success." % (self.name, ret))
 
     def get_omap_vals(self,
                       read_op: ReadOp,
-                      start_after: str,
-                      filter_prefix: str,
-                      max_return: int) -> Tuple[OmapIterator, int]:
+                      start_after: OMAP_KEY_TYPE,
+                      filter_prefix: OMAP_KEY_TYPE,
+                      max_return: int,
+                      omap_key_type = bytes.decode) -> Tuple[OmapIterator, int]:
         """
         get the omap values
         :para read_op: read operation object
@@ -3776,11 +3780,15 @@ returned %d, but should return zero on success." % (self.name, ret))
         with nogil:
             rados_read_op_omap_get_vals2(_read_op.read_op, _start_after, _filter_prefix,
                                          _max_return, &iter_addr, NULL, NULL)
-        it = OmapIterator(self)
+        it = OmapIterator(self, omap_key_type)
         it.ctx = iter_addr
         return it, 0   # 0 is meaningless; there for backward-compat
 
-    def get_omap_keys(self, read_op: ReadOp, start_after: str, max_return: int) -> Tuple[OmapIterator, int]:
+    def get_omap_keys(self,
+                      read_op: ReadOp,
+                      start_after: OMAP_KEY_TYPE,
+                      max_return: int,
+                      omap_key_type = bytes.decode) -> Tuple[OmapIterator, int]:
         """
         get the omap keys
         :para read_op: read operation object
@@ -3798,11 +3806,14 @@ returned %d, but should return zero on success." % (self.name, ret))
         with nogil:
             rados_read_op_omap_get_keys2(_read_op.read_op, _start_after,
                                          _max_return, &iter_addr, NULL, NULL)
-        it = OmapIterator(self)
+        it = OmapIterator(self, omap_key_type)
         it.ctx = iter_addr
         return it, 0   # 0 is meaningless; there for backward-compat
 
-    def get_omap_vals_by_keys(self, read_op: ReadOp, keys: Sequence[str]) -> Tuple[OmapIterator, int]:
+    def get_omap_vals_by_keys(self,
+                              read_op: ReadOp,
+                              keys: Sequence[OMAP_KEY_TYPE],
+                              omap_key_type = bytes.decode) -> Tuple[OmapIterator, int]:
         """
         get the omap values by keys
         :para read_op: read operation object
@@ -3821,13 +3832,13 @@ returned %d, but should return zero on success." % (self.name, ret))
                 rados_read_op_omap_get_vals_by_keys(_read_op.read_op,
                                                     <const char**>_keys,
                                                     key_num, &iter_addr, NULL)
-            it = OmapIterator(self)
+            it = OmapIterator(self, omap_key_type)
             it.ctx = iter_addr
             return it, 0   # 0 is meaningless; there for backward-compat
         finally:
             free(_keys)
 
-    def remove_omap_keys(self, write_op: WriteOp, keys: Sequence[str]):
+    def remove_omap_keys(self, write_op: WriteOp, keys: Sequence[OMAP_KEY_TYPE]):
         """
         remove omap keys specifiled
         :para write_op: write operation object
@@ -3858,7 +3869,7 @@ returned %d, but should return zero on success." % (self.name, ret))
         with nogil:
             rados_write_op_omap_clear(_write_op.write_op)
 
-    def remove_omap_range2(self, write_op: WriteOp, key_begin: str, key_end: str):
+    def remove_omap_range2(self, write_op: WriteOp, key_begin: OMAP_KEY_TYPE, key_end: OMAP_KEY_TYPE):
         """
         Remove key/value pairs from an object whose keys are in the range
         [key_begin, key_end)