Support for interactive transactions, and different column names mode (#46)

* Interactive session support

* Bumped package versions

* Column names in SQL result

* Tests

Co-authored-by: Bartłomiej Święcki <[email protected]>
Co-authored-by: Michael Abmayer <[email protected]>

Author: 3 people

.github/workflows/ci.yml

@@ -21,7 +21,7 @@ jobs:
           make dev
       - name: Start immudb container
         run: |
-          docker run -d --health-cmd "immuadmin status" --health-interval 10s --health-timeout 5s --health-retries 5 -v ${{ github.workspace }}/tests/certs/my.key.pem:/key.pem -p 3322:3322 codenotary/immudb:1.2.3 --signingKey=/key.pem
+          docker run -d --health-cmd "immuadmin status" --health-interval 10s --health-timeout 5s --health-retries 5 -v ${{ github.workspace }}/tests/certs/my.key.pem:/key.pem -p 3322:3322 codenotary/immudb:1.2.4 --signingKey=/key.pem
           docker run -d --health-cmd "immuadmin status" --health-interval 10s --health-timeout 5s --health-retries 5 -v ${{ github.workspace }}/tests/certs/my.key.pem:/key.pem -p 3333:3322 codenotary/immudb:1.1.0 --signingKey=/key.pem
       - name: Run tests
         run: |

immudb/client.py

@@ -13,19 +13,22 @@
 import grpc
 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 
-from immudb import header_manipulator_client_interceptor
+from immudb import grpcutils
+from immudb.grpc.schema_pb2 import OpenSessionResponse
 from immudb.handler import (batchGet, batchSet, changePassword, changePermission, createUser,
                             currentRoot, createDatabase, databaseList, deleteKeys, useDatabase,
                             get, listUsers, sqldescribe, verifiedGet, verifiedSet, setValue, history,
                             scan, reference, verifiedreference, zadd, verifiedzadd,
                             zscan, healthcheck, health, txbyid, verifiedtxbyid, sqlexec, sqlquery,
-                            listtables, execAll)
+                            listtables, execAll, transaction)
 from immudb.rootService import *
 from immudb.grpc import schema_pb2_grpc
 import warnings
 import ecdsa
 from immudb.datatypes import DeleteKeysRequest
 from immudb.embedded.store import KVMetadata
+import threading
+import queue
 
 import datetime
 
@@ -36,13 +39,15 @@ def __init__(self, immudUrl=None, rs: RootService = None, publicKeyFile: str = N
         if immudUrl is None:
             immudUrl = "localhost:3322"
         self.channel = grpc.insecure_channel(immudUrl)
-        self.__stub = schema_pb2_grpc.ImmuServiceStub(self.channel)
+        self._resetStub()
         if rs is None:
             self.__rs = RootService()
         else:
             self.__rs = rs
         self.__url = immudUrl
         self.loadKey(publicKeyFile)
+        self.__login_response = None
+        self._session_response = None
 
     def loadKey(self, kfile: str):
         if kfile is None:
@@ -58,22 +63,29 @@ def shutdown(self):
         self.intercept_channel = None
         self.__rs = None
 
+    def set_session_id_interceptor(self, openSessionResponse):
+        sessionId = openSessionResponse.sessionID
+        self.headersInterceptors = [
+            grpcutils.header_adder_interceptor('sessionid', sessionId)]
+        return self.get_intercepted_stub()
+
     def set_token_header_interceptor(self, response):
         try:
             token = response.token
         except AttributeError:
             token = response.reply.token
-        self.header_interceptor = \
-            header_manipulator_client_interceptor.header_adder_interceptor(
+        self.headersInterceptors = [
+            grpcutils.header_adder_interceptor(
                 'authorization', "Bearer " + token
             )
-        try:
-            self.intercept_channel = grpc.intercept_channel(
-                self.channel, self.header_interceptor)
-        except ValueError as e:
-            raise Exception(
-                "Attempted to login on termninated client, channel has been shutdown") from e
-        return schema_pb2_grpc.ImmuServiceStub(self.intercept_channel)
+        ]
+        return self.get_intercepted_stub()
+
+    def get_intercepted_stub(self):
+        intercepted, newStub = grpcutils.get_intercepted_stub(
+            self.channel, self.headersInterceptors)
+        self.intercept_channel = intercepted
+        return newStub
 
     @property
     def stub(self):
@@ -88,12 +100,17 @@ def healthCheck(self):
         return healthcheck.call(self.__stub, self.__rs)
 
     # Not implemented: connect
+    def _convertToBytes(self, what):
+        if(type(what) != bytes):
+            return bytes(what, encoding='utf-8')
+        return what
 
     def login(self, username, password, database=b"defaultdb"):
-        req = schema_pb2_grpc.schema__pb2.LoginRequest(user=bytes(
-            username, encoding='utf-8'), password=bytes(
-                password, encoding='utf-8'
-        ))
+        convertedUsername = self._convertToBytes(username)
+        convertedPassword = self._convertToBytes(password)
+        convertedDatabase = self._convertToBytes(database)
+        req = schema_pb2_grpc.schema__pb2.LoginRequest(
+            user=convertedUsername, password=convertedPassword)
         try:
             self.__login_response = schema_pb2_grpc.schema__pb2.LoginResponse = \
                 self.__stub.Login(
@@ -105,7 +122,8 @@ def login(self, username, password, database=b"defaultdb"):
 
         self.__stub = self.set_token_header_interceptor(self.__login_response)
         # Select database, modifying stub function accordingly
-        request = schema_pb2_grpc.schema__pb2.Database(databaseName=database)
+        request = schema_pb2_grpc.schema__pb2.Database(
+            databaseName=convertedDatabase)
         resp = self.__stub.UseDatabase(request)
         self.__stub = self.set_token_header_interceptor(resp)
 
@@ -115,9 +133,62 @@ def login(self, username, password, database=b"defaultdb"):
     def logout(self):
         self.__stub.Logout(google_dot_protobuf_dot_empty__pb2.Empty())
         self.__login_response = None
+        self._resetStub()
+
+    def _resetStub(self):
+        self.headersInterceptors = []
+        self.__stub = schema_pb2_grpc.ImmuServiceStub(self.channel)
 
-    # Not implemented: openSession
-    # Not implemented: closeSession
+    def keepAlive(self):
+        self.__stub.KeepAlive(google_dot_protobuf_dot_empty__pb2.Empty())
+
+    def openManagedSession(self, username, password, database=b"defaultdb", keepAliveInterval=60):
+        class ManagedSession:
+            def __init__(this, keepAliveInterval):
+                this.keepAliveInterval = keepAliveInterval
+                this.keepAliveStarted = False
+                this.keepAliveProcess = None
+                this.queue = queue.Queue()
+
+            def manage(this):
+                while this.keepAliveStarted:
+                    try:
+                        what = this.queue.get(True, this.keepAliveInterval)
+                    except queue.Empty:
+                        self.keepAlive()
+
+            def __enter__(this):
+                interface = self.openSession(username, password, database)
+                this.keepAliveStarted = True
+                this.keepAliveProcess = threading.Thread(target=this.manage)
+                this.keepAliveProcess.start()
+                return interface
+
+            def __exit__(this, type, value, traceback):
+                this.keepAliveStarted = False
+                this.queue.put(b'0')
+                self.closeSession()
+
+        return ManagedSession(keepAliveInterval)
+
+    def openSession(self, username, password, database=b"defaultdb"):
+        convertedUsername = self._convertToBytes(username)
+        convertedPassword = self._convertToBytes(password)
+        convertedDatabase = self._convertToBytes(database)
+        req = schema_pb2_grpc.schema__pb2.OpenSessionRequest(
+            username=convertedUsername,
+            password=convertedPassword,
+            databaseName=convertedDatabase
+        )
+        self._session_response = schema_pb2_grpc.schema__pb2.OpenSessionResponse = self.__stub.OpenSession(
+            req)
+        self.__stub = self.set_session_id_interceptor(self._session_response)
+        return transaction.Tx(self.__stub, self._session_response, self.channel)
+
+    def closeSession(self):
+        self.__stub.CloseSession(google_dot_protobuf_dot_empty__pb2.Empty())
+        self._session_response = None
+        self._resetStub()
 
     def createUser(self, user, password, permission, database):
         request = schema_pb2_grpc.schema__pb2.CreateUserRequest(
@@ -213,7 +284,7 @@ def verifiedGet(self, key: bytes):
         return verifiedGet.call(self.__stub, self.__rs, key, verifying_key=self.__vk)
 
     def verifiedGetSince(self, key: bytes, sinceTx: int):
-        return verifiedGet.call(self.__stub, self.__rs, key, sinceTx, self.__vk)
+        return verifiedGet.call(self.__stub, self.__rs, key, sinceTx=sinceTx, verifying_key=self.__vk)
 
     def verifiedGetAt(self, key: bytes, atTx: int):
         return verifiedGet.call(self.__stub, self.__rs, key, atTx, self.__vk)
@@ -299,7 +370,7 @@ def sqlExec(self, stmt, params={}, noWait=False):
 
         return sqlexec.call(self.__stub, self.__rs, stmt, params, noWait)
 
-    def sqlQuery(self, query, params={}):
+    def sqlQuery(self, query, params={}, columnNameMode=constants.COLUMN_NAME_MODE_NONE):
         """Queries the database using SQL
         Args:
             query: a query in immudb SQL dialect.
@@ -310,7 +381,7 @@ def sqlQuery(self, query, params={}):
 
             ['table1', 'table2']
         """
-        return sqlquery.call(self.__stub, self.__rs, query, params)
+        return sqlquery.call(self.__stub, self.__rs, query, params, columnNameMode)
 
     def listTables(self):
         """List all tables in the current database
@@ -326,7 +397,6 @@ def describeTable(self, table):
         return sqldescribe.call(self.__stub, self.__rs, table)
 
     # Not implemented: verifyRow
-    # Not implemented: newTx
 
 # deprecated
     def databaseCreate(self, dbName: bytes):
@@ -360,7 +430,6 @@ def safeSet(self, key: bytes, value: bytes):  # deprecated
 
 # immudb-py only
 
-
     def getAllValues(self, keys: list):  # immudb-py only
         resp = batchGet.call(self.__stub, self.__rs, keys)
         return resp

immudb/constants.py

@@ -28,3 +28,9 @@
 
 OLDEST_FIRST = False
 NEWEST_FIRST = True
+
+COLUMN_NAME_MODE_FULL = 255
+COLUMN_NAME_MODE_DATABASE = 3
+COLUMN_NAME_MODE_TABLE = 2
+COLUMN_NAME_MODE_FIELD = 1
+COLUMN_NAME_MODE_NONE = 0

immudb/datatypes.py

@@ -11,7 +11,14 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-import struct
+from enum import IntEnum
+from immudb.grpc.schema_pb2 import ReadOnly, WriteOnly, ReadWrite
+
+
+class TxMode(IntEnum):
+    ReadOnly = ReadOnly
+    WriteOnly = WriteOnly
+    ReadWrite = ReadWrite
 
 
 @dataclass

immudb/exceptions.py

@@ -45,3 +45,7 @@ class ErrKeyNotFound(Exception):
 
 class ErrMetadataUnsupported(Exception):
     pass
+
+
+class ErrPySDKInvalidColumnMode(Exception):
+    pass

immudb/header_manipulator_client_interceptor.py renamed to immudb/grpcutils.py

@@ -17,6 +17,7 @@
 
 import grpc
 import immudb.generic_client_interceptor as generic_client_interceptor
+import immudb.grpc.schema_pb2_grpc as schema_pb2_grpc
 
 
 class _ClientCallDetails(
@@ -44,3 +45,13 @@ def intercept_call(client_call_details, request_iterator, request_streaming,
         return client_call_details, request_iterator, None
 
     return generic_client_interceptor.create(intercept_call)
+
+
+def get_intercepted_stub(channel, headers):
+    try:
+        intercepted = grpc.intercept_channel(channel, *headers)
+        newStub = schema_pb2_grpc.ImmuServiceStub(intercepted)
+        return intercepted, newStub
+    except ValueError as e:
+        raise Exception(
+            "Attempted to login on termninated client, channel has been shutdown") from e

immudb/handler/sqlexec.py

@@ -10,16 +10,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
-
 from immudb.grpc import schema_pb2
 from immudb.grpc import schema_pb2_grpc
 from immudb.rootService import RootService
-from immudb import datatypes
 from immudb.typeconv import py_to_sqlvalue
 
 
 def call(service: schema_pb2_grpc.ImmuServiceStub, rs: RootService, stmt, params, noWait):
+    return _call_with_executor(stmt, params, noWait, service.SQLExec)
+
+
+def _call_with_executor(stmt, params, noWait, executor):
     paramsObj = []
     for key, value in params.items():
         paramsObj.append(schema_pb2.NamedParam(
@@ -30,5 +31,5 @@ def call(service: schema_pb2_grpc.ImmuServiceStub, rs: RootService, stmt, params
         params=paramsObj,
         noWait=noWait)
 
-    result = service.SQLExec(request)
+    result = executor(request)
     return result

immudb/handler/sqlquery.py

@@ -15,9 +15,15 @@
 from immudb.rootService import RootService
 from immudb.typeconv import py_to_sqlvalue
 from immudb.typeconv import sqlvalue_to_py
+from immudb import constants
+from immudb.exceptions import ErrPySDKInvalidColumnMode
 
 
-def call(service: schema_pb2_grpc.ImmuServiceStub, rs: RootService, query, params):
+def call(service: schema_pb2_grpc.ImmuServiceStub, rs: RootService, query, params, columnNameMode):
+    return _call_with_executor(query, params, columnNameMode, service.SQLQuery)
+
+
+def _call_with_executor(query, params, columnNameMode, executor):
     paramsObj = []
     for key, value in params.items():
         paramsObj.append(schema_pb2.NamedParam(
@@ -27,8 +33,32 @@ def call(service: schema_pb2_grpc.ImmuServiceStub, rs: RootService, query, param
         sql=query,
         params=paramsObj)
 
-    resp = service.SQLQuery(request)
+    resp = executor(request)
     result = []
+
+    columnNames = getColumnNames(resp, columnNameMode)
+
     for row in resp.rows:
-        result.append(tuple([sqlvalue_to_py(i) for i in row.values]))
+        if columnNameMode == constants.COLUMN_NAME_MODE_NONE:
+            result.append(tuple([sqlvalue_to_py(i) for i in row.values]))
+        else:
+            result.append(
+                dict(zip(columnNames, tuple([sqlvalue_to_py(i) for i in row.values]))))
     return result
+
+
+def getColumnNames(resp, columnNameMode):
+    columnNames = []
+    if columnNameMode:
+        for column in resp.columns:
+            if columnNameMode == constants.COLUMN_NAME_MODE_FIELD:
+                columnNames.append(column.name.strip("()").split(".")[2])
+            elif columnNameMode == constants.COLUMN_NAME_MODE_TABLE:
+                columnNames.append(column.name.strip("()").split(".", 1)[1])
+            elif columnNameMode == constants.COLUMN_NAME_MODE_DATABASE:
+                columnNames.append(column.name.strip("()"))
+            elif columnNameMode == constants.COLUMN_NAME_MODE_FULL:
+                columnNames.append(column.name)
+            else:
+                raise ErrPySDKInvalidColumnMode
+    return columnNames