feat: renew session for expiration, re-execute when error

.github/workflows/run_test.yaml

@@ -9,10 +9,11 @@ on:
     - cron: "0 6 * * *"
 
 jobs:
-  ci-pip-install:
+  ci-pip-install-from-source:
+    # This is to verify the setup.py as a mitigation for remain python 3.6.2+ capability
     runs-on: ubuntu-20.04
     strategy:
-      max-parallel: 1
+      max-parallel: 2
       matrix:
         python-version: [3.6, 3.7]
     steps:
@@ -21,19 +22,19 @@ jobs:
         uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Install dependencies
+      - name: Install nebulagraph-python from source and test dependencies
         run: |
           python setup.py install
           pip install pip-tools pytest
-      - name: Test with pytest==6.2.5
+      - name: Test with pytest
         run: |
           docker-compose -f docker-compose.yaml up -d
           sleep 20
           pytest -s -v -k "not SSL"
         working-directory: tests
 
 
-  ci:
+  build-lint-test:
     runs-on: ubuntu-22.04
     strategy:
       max-parallel: 2
@@ -57,7 +58,10 @@ jobs:
           cache: true
 
       - name: Install dependencies
-        run: pdm install
+        run: |
+          pip install .
+          pdm install -G:dev
+          pdm install -G:test
       - name: lint
         run: pdm fmt-check
       - name: Test with pytest
@@ -75,7 +79,7 @@ jobs:
         with:
           files: coverage.xml
 
-  example:
+  example-test:
     runs-on: ubuntu-latest
     strategy:
       matrix:
@@ -88,11 +92,13 @@ jobs:
           python-version: ${{ matrix.python-version }}
           cache: true
       - name: Install dependencies
-        run: pdm install -G example
+        run: |
+          pip install .
+          pip install prettytable pandas
       - name: Setup containers
         run: |
           docker-compose -f tests/docker-compose.yaml up -d
           sleep 20
       - name: Test example
         run: |
-          for f in example/*.py; do pdm run python "$f"; done
+          for f in example/*.py; do python3 "$f"; done

README.md

@@ -4,7 +4,7 @@ This repository holds the official Python API for NebulaGraph.
 
 [![pdm-managed](https://img.shields.io/badge/pdm-managed-blueviolet)](https://pdm.fming.dev)
 [![pypi-version](https://img.shields.io/pypi/v/nebula3-python)](https://pypi.org/project/nebula3-python/)
-[![python-version](https://img.shields.io/badge/python-3.6.2%20|%203.7%20|%203.8%20|%203.9%20|%203.10%20|%203.11%20|%203.12-blue)](https://www.python.org/)
+[![python-version](https://img.shields.io/badge/python-3.6.2+%20|%203.7%20|%203.8%20|%203.9%20|%203.10%20|%203.11%20|%203.12-blue)](https://www.python.org/)
 
 
 ## Before you start

nebula3/Exception.py

@@ -72,6 +72,25 @@ def __init__(self, message):
         self.message = 'Invalid hostname: {}'.format(message)
 
 
+class SessionException(Exception):
+    E_SESSION_INVALID = -1002
+    E_SESSION_TIMEOUT = -1003
+
+    def __init__(self, code=E_SESSION_INVALID, message=None):
+        Exception.__init__(self, message)
+        self.type = code
+        self.message = message
+
+
+class ExecutionErrorException(Exception):
+    E_EXECUTION_ERROR = -1005
+
+    def __init__(self, message=None):
+        Exception.__init__(self, message)
+        self.type = self.E_EXECUTION_ERROR
+        self.message = message
+
+
 class IOErrorException(Exception):
     E_UNKNOWN = 0
     E_ALL_BROKEN = 1

nebula3/gclient/net/Connection.py

@@ -25,6 +25,8 @@
     AuthFailedException,
     IOErrorException,
     ClientServerIncompatibleException,
+    SessionException,
+    ExecutionErrorException,
 )
 
 from nebula3.gclient.net.AuthResult import AuthResult
@@ -146,6 +148,12 @@ def execute_parameter(self, session_id, stmt, params):
         """
         try:
             resp = self._connection.executeWithParameter(session_id, stmt, params)
+            if resp.error_code == ErrorCode.E_SESSION_INVALID:
+                raise SessionException(resp.error_code, resp.error_msg)
+            if resp.error_code == ErrorCode.E_SESSION_TIMEOUT:
+                raise SessionException(resp.error_code, resp.error_msg)
+            if resp.error_code == ErrorCode.E_EXECUTION_ERROR:
+                raise ExecutionErrorException(resp.error_msg)
             return resp
         except Exception as te:
             if isinstance(te, TTransportException):

nebula3/gclient/net/Session.py

@@ -10,6 +10,7 @@
 from nebula3.Exception import (
     IOErrorException,
     NotValidConnectionException,
+    ExecutionErrorException,
 )
 
 from nebula3.data.ResultSet import ResultSet
@@ -18,14 +19,34 @@
 
 
 class Session(object):
-    def __init__(self, connection, auth_result: AuthResult, pool, retry_connect=True):
+    def __init__(
+        self,
+        connection,
+        auth_result: AuthResult,
+        pool,
+        retry_connect=True,
+        retry_times=3,
+        retry_interval_sec=1,
+    ):
+        """
+        Initialize the Session object.
+
+        :param connection: The connection object associated with the session.
+        :param auth_result: The result of the authentication process.
+        :param pool: The pool object where the session was created.
+        :param retry_connect: A boolean indicating whether to retry the connection if it fails.
+        :param retry_times: The number of times to retry the connection.
+        :param retry_interval_sec: The interval between connection retries in seconds.
+        """
         self._session_id = auth_result.get_session_id()
         self._timezone_offset = auth_result.get_timezone_offset()
         self._connection = connection
         self._timezone = 0
         # connection the where the session was created, if session pool was used
         self._pool = pool
         self._retry_connect = retry_connect
+        self._retry_times = retry_times
+        self._retry_interval_sec = retry_interval_sec
         # the time stamp when the session was added to the idle list of the session pool
         self._idle_time_start = 0
 
@@ -65,6 +86,27 @@ def execute_parameter(self, stmt, params):
                         timezone_offset=self._timezone_offset,
                     )
             raise
+        except ExecutionErrorException as eee:
+            retry_count = 0
+            while retry_count < self._retry_times:
+                try:
+                    # TODO: add exponential backoff
+                    time.sleep(self._retry_interval_sec)
+                    resp = self._connection.execute_parameter(
+                        self._session_id, stmt, params
+                    )
+                    end_time = time.time()
+                    return ResultSet(
+                        resp,
+                        all_latency=int((end_time - start_time) * 1000000),
+                        timezone_offset=self._timezone_offset,
+                    )
+                except ExecutionErrorException:
+                    if retry_count >= self._retry_times - 1:
+                        raise eee
+                    else:
+                        retry_count += 1
+                        continue
         except Exception:
             raise
 
@@ -222,6 +264,22 @@ def execute_json_with_parameter(self, stmt, params):
                     )
                     return resp_json
             raise
+        except ExecutionErrorException as eee:
+            retry_count = 0
+            while retry_count < self._retry_times:
+                try:
+                    # TODO: add exponential backoff
+                    time.sleep(self._retry_interval_sec)
+                    resp = self._connection.execute_json_with_parameter(
+                        self._session_id, stmt, params
+                    )
+                    return resp
+                except ExecutionErrorException:
+                    if retry_count >= self._retry_times - 1:
+                        raise eee
+                    else:
+                        retry_count += 1
+                        continue
         except Exception:
             raise
 

nebula3/gclient/net/SessionPool.py

@@ -15,6 +15,7 @@
     AuthFailedException,
     NoValidSessionException,
     InValidHostname,
+    SessionException,
 )
 
 from nebula3.gclient.net.Session import Session
@@ -170,6 +171,18 @@ def execute_parameter(self, stmt, params):
             self._return_session(session)
 
             return resp
+        except SessionException as se:
+            if se.type in [
+                SessionException.E_SESSION_INVALID,
+                SessionException.E_SESSION_TIMEOUT,
+            ]:
+                self._active_sessions.remove(session)
+                session = self._get_idle_session()
+                if session is None:
+                    raise RuntimeError('Get session failed')
+                self._add_session_to_idle(session)
+            raise se
+
         except Exception as e:
             logger.error('Execute failed: {}'.format(e))
             # remove the session from the pool if it is invalid
@@ -257,6 +270,17 @@ def execute_json_with_parameter(self, stmt, params):
             self._return_session(session)
 
             return resp
+        except SessionException as se:
+            if se.type in [
+                SessionException.E_SESSION_INVALID,
+                SessionException.E_SESSION_TIMEOUT,
+            ]:
+                self._active_sessions.remove(session)
+                session = self._get_idle_session()
+                if session is None:
+                    raise RuntimeError('Get session failed')
+                self._add_session_to_idle(session)
+            raise se
         except Exception as e:
             logger.error('Execute failed: {}'.format(e))
             # remove the session from the pool if it is invalid