From 8e848f16df33b5a128ad1280d2fde4be8df49991 Mon Sep 17 00:00:00 2001
From: Jon Fu <jon.fu@snowflake.com>
Date: Mon, 28 Feb 2022 18:15:10 -0500
Subject: [PATCH] Support tuples in python tenants

---
 bindings/python/fdb/impl.py                   | 19 ++++++--
 .../python/tests/tenant_tuple_name_tests.py   | 47 +++++++++++++++++++
 bindings/python/tests/tester.py               |  3 ++
 3 files changed, 64 insertions(+), 5 deletions(-)
 create mode 100644 bindings/python/tests/tenant_tuple_name_tests.py

diff --git a/bindings/python/fdb/impl.py b/bindings/python/fdb/impl.py
index 3a7ea07c9e..47e22f043d 100644
--- a/bindings/python/fdb/impl.py
+++ b/bindings/python/fdb/impl.py
@@ -34,6 +34,7 @@ import traceback
 
 import fdb
 from fdb import six
+from fdb.tuple import pack, unpack
 
 _network_thread = None
 _network_thread_reentrant_lock = threading.RLock()
@@ -1146,6 +1147,13 @@ class _TransactionCreator(_FDBBase):
                 yield None
         return TransactionCreator
 
+def process_tenant_name(name):
+    if isinstance(name, tuple):
+        return pack(name)
+    elif isinstance(name, bytes):
+        return name
+    else:
+        raise TypeError('Tenant name must be of type ' + bytes.__name__ + ' or of type ' + tuple.__name__)
 
 class Database(_TransactionCreator):
     def __init__(self, dpointer):
@@ -1160,10 +1168,9 @@ class Database(_TransactionCreator):
         self.capi.fdb_database_set_option(self.dpointer, option, param, length)
 
     def open_tenant(self, name):
-        if not isinstance(name, bytes):
-            raise TypeError('Tenant name must be of type ' + bytes.__name__)
+        tname = process_tenant_name(name)
         pointer = ctypes.c_void_p()
-        self.capi.fdb_database_open_tenant(self.dpointer, name, len(name), ctypes.byref(pointer))
+        self.capi.fdb_database_open_tenant(self.dpointer, tname, len(tname), ctypes.byref(pointer))
         return Tenant(pointer.value)
 
     def create_transaction(self):
@@ -1172,10 +1179,12 @@ class Database(_TransactionCreator):
         return Transaction(pointer.value, self)
 
     def allocate_tenant(self, name):
-        return FutureVoid(self.capi.fdb_database_allocate_tenant(self.dpointer, name, len(name)))
+        tname = process_tenant_name(name)
+        return FutureVoid(self.capi.fdb_database_allocate_tenant(self.dpointer, tname, len(tname)))
 
     def delete_tenant(self, name):
-        return FutureVoid(self.capi.fdb_database_remove_tenant(self.dpointer, name, len(name)))
+        tname = process_tenant_name(name)
+        return FutureVoid(self.capi.fdb_database_remove_tenant(self.dpointer, tname, len(tname)))
 
 
 class Tenant(_TransactionCreator):
diff --git a/bindings/python/tests/tenant_tuple_name_tests.py b/bindings/python/tests/tenant_tuple_name_tests.py
new file mode 100644
index 0000000000..3b64675c68
--- /dev/null
+++ b/bindings/python/tests/tenant_tuple_name_tests.py
@@ -0,0 +1,47 @@
+#!/usr/bin/python
+#
+# tenant_tuple_name_tests.py
+#
+# This source file is part of the FoundationDB open source project
+#
+# Copyright 2013-2022 Apple Inc. and the FoundationDB project authors
+#
+# 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
+#
+#     http://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.
+#
+import fdb
+import sys
+from fdb.tuple import pack
+
+if __name__ == '__main__':
+    fdb.api_version(710)
+
+def test_tenant_tuple_name(db):
+    tuplename=(b'test', b'level', b'hierarchy', 3, 1.24, 'str')
+    db.allocate_tenant(tuplename).wait()
+
+    tenant=db.open_tenant(tuplename)
+    tenant[b'foo'] = b'bar'
+
+    assert tenant[b'foo'] == b'bar'
+
+    del tenant[b'foo']
+    db.delete_tenant(tuplename).wait()
+
+# Expect a cluster file as input. This test will write to the FDB cluster, so
+# be aware of potential side effects.
+if __name__ == '__main__':
+    clusterFile = sys.argv[1]
+    db = fdb.open(clusterFile)
+    db.options.set_transaction_timeout(2000)  # 2 seconds
+    db.options.set_transaction_retry_limit(3)
+    test_tenant_tuple_name(db)
diff --git a/bindings/python/tests/tester.py b/bindings/python/tests/tester.py
index e7ce61d766..3f1f3f10d9 100644
--- a/bindings/python/tests/tester.py
+++ b/bindings/python/tests/tester.py
@@ -49,6 +49,7 @@ from cancellation_timeout_tests import test_db_retry_limits
 from cancellation_timeout_tests import test_combinations
 
 from size_limit_tests import test_size_limit_option, test_get_approximate_size
+from tenant_tuple_name_tests import test_tenant_tuple_name
 
 random.seed(0)
 
@@ -618,6 +619,8 @@ class Tester:
                         test_size_limit_option(db)
                         test_get_approximate_size(db)
 
+                        test_tenant_tuple_name(db)
+
                     except fdb.FDBError as e:
                         print("Unit tests failed: %s" % e.description)
                         traceback.print_exc()