Merge pull request #94 from TileDB-Inc/jcb/consolidate_key

add key kw arg to array / kv consolidate

tiledb/libtiledb.pxd

@@ -697,7 +697,7 @@ cdef extern from "tiledb/tiledb.h":
         const tiledb_array_schema_t* array_schema,
         tiledb_encryption_type_t key_type,
         const void* key,
-        unsigned int key_len) nogil;
+        unsigned int key_len) nogil
 
     int tiledb_array_is_open(
         tiledb_ctx_t* ctx,
@@ -708,6 +708,13 @@ cdef extern from "tiledb/tiledb.h":
         tiledb_ctx_t* ctx,
         const char* array_path) nogil
 
+    int tiledb_array_consolidate_with_key(
+        tiledb_ctx_t* ctx,
+        const char* uri,
+        tiledb_encryption_type_t key_type,
+        const void* key_ptr,
+        unsigned int key_len) nogil
+
     int tiledb_array_get_schema(
         tiledb_ctx_t* ctx,
         tiledb_array_t* array,
@@ -862,6 +869,13 @@ cdef extern from "tiledb/tiledb.h":
         tiledb_ctx_t* ctx,
         const char* kv_uri) nogil
 
+    int tiledb_kv_consolidate_with_key(
+        tiledb_ctx_t* ctx,
+        const char* kv_uri,
+        tiledb_encryption_type_t encryption_type,
+        const void* encryption_key,
+        unsigned int key_len) nogil
+
     int tiledb_kv_alloc(
         tiledb_ctx_t* ctx,
         const char* kv_uri,

tiledb/libtiledb.pyx

@@ -2310,7 +2310,7 @@ cdef class KV(object):
                 ctx_ptr, uri_ptr, schema_ptr, key_type, key_ptr, key_len)
         if rc != TILEDB_OK:
             _raise_ctx_err(ctx_ptr, rc)
-        return KV(ctx, uri)
+        return
 
     def __init__(self, Ctx ctx, uri, mode='r', key=None, timestamp=None):
         cdef tiledb_ctx_t* ctx_ptr = ctx.ptr
@@ -2475,9 +2475,23 @@ cdef class KV(object):
         cdef tiledb_ctx_t* ctx_ptr = self.ctx.ptr
         cdef bytes buri = unicode_path(self.uri)
         cdef const char* uri_ptr = PyBytes_AS_STRING(buri)
+        # encyrption key
+        cdef bytes bkey
+        cdef tiledb_encryption_type_t key_type = TILEDB_NO_ENCRYPTION
+        cdef void* key_ptr = NULL
+        cdef unsigned int key_len = 0
+        if key is not None:
+            if isinstance(key, str):
+                bkey = key.encode('ascii')
+            else:
+                bkey = bytes(key)
+            key_type = TILEDB_AES_256_GCM
+            key_ptr = <void *> PyBytes_AS_STRING(bkey)
+            #TODO: unsafe cast here ssize_t -> uint64_t
+            key_len = <unsigned int> PyBytes_GET_SIZE(bkey)
         cdef int rc = TILEDB_OK
         with nogil:
-            rc = tiledb_kv_consolidate(ctx_ptr, uri_ptr)
+            rc = tiledb_kv_consolidate_with_key(ctx_ptr, uri_ptr, key_type, key_ptr, key_len)
         if rc != TILEDB_OK:
             _raise_ctx_err(ctx_ptr, rc)
         return
@@ -3469,13 +3483,13 @@ cdef class Array(object):
         return tuple((extents[i, 0].item(), extents[i, 1].item())
                      for i in range(dom.ndim))
 
-    def consolidate(self):
+    def consolidate(self, key=None):
         """Consolidates fragments of an array object for increased read performance.
 
         :raises: :py:exc:`tiledb.TileDBError`
 
         """
-        return consolidate(self.ctx, self.uri)
+        return consolidate(self.ctx, uri=self.uri, key=key)
 
     def dump(self):
         self.schema.dump()
@@ -4456,12 +4470,12 @@ cdef class SparseArray(Array):
         return out
 
 
-def consolidate(Ctx ctx, uri):
+def consolidate(Ctx ctx, uri=None, key=None):
     """Consolidates a TileDB Array updates for improved read performance
 
     :param tiledb.Ctx ctx: The TileDB Context
     :param str uri: URI to the TileDB Array
-    :param int timestamp: (default None) timestamp
+    :param str: (default None) Key to decrypt array if the array is encrypted
     :rtype: str
     :return: path (URI) to the consolidated TileDB Array
     :raises TypeError: cannot convert path to unicode string
@@ -4471,9 +4485,23 @@ def consolidate(Ctx ctx, uri):
     cdef tiledb_ctx_t* ctx_ptr = ctx.ptr
     cdef bytes buri = unicode_path(uri)
     cdef const char* uri_ptr = PyBytes_AS_STRING(buri)
+    # encyrption key
+    cdef bytes bkey
+    cdef tiledb_encryption_type_t key_type = TILEDB_NO_ENCRYPTION
+    cdef void* key_ptr = NULL
+    cdef unsigned int key_len = 0
+    if key is not None:
+        if isinstance(key, str):
+            bkey = key.encode('ascii')
+        else:
+            bkey = bytes(key)
+        key_type = TILEDB_AES_256_GCM
+        key_ptr = <void *> PyBytes_AS_STRING(bkey)
+        #TODO: unsafe cast here ssize_t -> uint64_t
+        key_len = <unsigned int> PyBytes_GET_SIZE(bkey)
     cdef int rc = TILEDB_OK
     with nogil:
-        rc = tiledb_array_consolidate(ctx_ptr, uri_ptr)
+        rc = tiledb_array_consolidate_with_key(ctx_ptr, uri_ptr, key_type, key_ptr, key_len)
     if rc != TILEDB_OK:
         _raise_ctx_err(ctx_ptr, rc)
     return uri

tiledb/tests/test_libtiledb.py

@@ -488,6 +488,7 @@ def test_array_create_encrypted(self):
                 self.assertTrue(array.isopen)
                 self.assertEqual(array.schema, schema)
                 self.assertEqual(array.mode, 'r')
+            tiledb.consolidate(ctx, uri=self.path("foo"), key=key)
 
         # check that opening the array with the wrong key fails:
         with self.assertRaises(tiledb.TileDBError):
@@ -499,6 +500,11 @@ def test_array_create_encrypted(self):
             tiledb.libtiledb.Array(ctx, self.path("foo"), mode='r',
                                    key=b"0123456789abcdeF0123456789abcde")
 
+        # check that consolidating the array with the wrong key fails:
+        with self.assertRaises(tiledb.TileDBError):
+            tiledb.consolidate(ctx, uri=self.path("foo"),
+                               key=b"0123456789abcdeF0123456789abcde")
+
     def test_array_doesnt_exist(self):
         ctx = tiledb.Ctx()
         with self.assertRaises(tiledb.TileDBError):
@@ -1360,6 +1366,33 @@ def test_kv_write_load_read(self):
           self.assertTrue('foo' in kv)
           self.assertFalse('bar' in kv)
 
+    def test_kv_write_load_read_encrypted(self):
+         # create a kv array
+        ctx = tiledb.Ctx()
+        a1 = tiledb.Attr(ctx, "value", dtype=bytes)
+        schema = tiledb.KVSchema(ctx, attrs=(a1,))
+        # persist kv schema
+        tiledb.KV.create(ctx, self.path("foo"), schema, key=b"0123456789abcdeF0123456789abcdeF")
+
+        with tiledb.KV(ctx, self.path("foo"), mode='w', key=b"0123456789abcdeF0123456789abcdeF") as kv:
+            kv['foo'] = 'bar'
+
+        # try to load it
+        with tiledb.KV(ctx, self.path("foo"), mode='r', key=b"0123456789abcdeF0123456789abcdeF") as kv:
+          self.assertEqual(kv["foo"], 'bar')
+          self.assertTrue('foo' in kv)
+          self.assertFalse('bar' in kv)
+
+        # loading with no key fails
+        with self.assertRaises(tiledb.TileDBError):
+            with tiledb.KV(ctx, self.path("foo")) as kv:
+                self.assertTrue('foo' in kv)
+
+        # loading with wrong key fails
+        with self.assertRaises(tiledb.TileDBError):
+            with tiledb.KV(ctx, self.path("foo"), key=b"0123456789abcdeF0123456789abcdeZ") as kv:
+                self.assertTrue('foo' in kv)
+
     def test_kv_update_reload(self):
         # create a kv array
         ctx = tiledb.Ctx()