Merge pull request #146 from aperture-data/release-0.3.2

Release 0.3.2

Author: gsaluja9

aperturedb/Connector.py

@@ -24,9 +24,10 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 #
-
+from __future__ import annotations
 from . import queryMessage_pb2
 import sys
+import traceback
 import os
 import socket
 import struct
@@ -35,6 +36,8 @@
 import ssl
 import logging
 
+from threading import Lock
+from types import SimpleNamespace
 from dataclasses import dataclass
 
 logger = logging.getLogger(__name__)
@@ -43,6 +46,10 @@
 PROTOCOL_VERSION = 1
 
 
+class UnauthorizedException(Exception):
+    pass
+
+
 @dataclass
 class Session():
 
@@ -55,7 +62,8 @@ class Session():
     def valid(self) -> bool:
         session_age = time.time() - self.session_started
 
-        if session_age > self.session_token_ttl:
+        # This triggers refresh if the session is about to expire.
+        if session_age > self.session_token_ttl - os.getenv("SESSION_EXPIRTY_OFFSET_SEC", 10):
             return False
 
         return True
@@ -81,8 +89,7 @@ class Connector(object):
 
     def __init__(self, host="localhost", port=55555,
                  user="", password="", token="",
-                 session=None,
-                 use_ssl=True):
+                 use_ssl=True, shared_data=None):
 
         self.use_ssl = use_ssl
 
@@ -93,17 +100,18 @@ def __init__(self, host="localhost", port=55555,
         self.last_response   = ''
         self.last_query_time = 0
 
-        self.session = None
-
         self._connect()
 
-        if session:
-            self.session = session
-        else:
+        if shared_data is None:
+            self.shared_data = SimpleNamespace()
+            self.shared_data.session = None
+            self.shared_data.lock = Lock()
             try:
                 self._authenticate(user, password, token)
             except Exception as e:
                 raise Exception("Authentication failed:", str(e))
+        else:
+            self.shared_data = shared_data
 
     def __del__(self):
 
@@ -152,39 +160,44 @@ def _authenticate(self, user, password="", token=""):
         if session_info["status"] != 0:
             raise Exception(session_info["info"])
 
-        self.session = Session(session_info["session_token"],
-                               session_info["refresh_token"],
-                               session_info["session_token_expires_in"],
-                               session_info["refresh_token_expires_in"],
-                               )
+        self.shared_data.session = Session(session_info["session_token"],
+                                           session_info["refresh_token"],
+                                           session_info["session_token_expires_in"],
+                                           session_info["refresh_token_expires_in"],
+                                           time.time()
+                                           )
 
     def _check_session_status(self):
-
-        if not self.session:
+        if not self.shared_data.session:
             return
 
-        if not self.session.valid():
-            self._refresh_token()
+        if not self.shared_data.session.valid():
+            with self.shared_data.lock:
+                self._refresh_token()
 
     def _refresh_token(self):
-
         query = [{
             "RefreshToken": {
-                "refresh_token": self.session.refresh_token
+                "refresh_token": self.shared_data.session.refresh_token
             }
         }]
 
         response, _ = self._query(query, [])
 
-        session_info = response[0]["RefreshToken"]
-        if session_info["status"] != 0:
-            raise Exception(session_info["info"])
-
-        self.session = Session(session_info["session_token"],
-                               session_info["refresh_token"],
-                               session_info["session_token_expires_in"],
-                               session_info["refresh_token_expires_in"],
-                               )
+        logger.info(f"Refresh token response: \r\n{response}")
+        if isinstance(response, list):
+            session_info = response[0]["RefreshToken"]
+            if session_info["status"] != 0:
+                raise UnauthorizedException(response)
+
+            self.shared_data.session = Session(session_info["session_token"],
+                                               session_info["refresh_token"],
+                                               session_info["session_token_expires_in"],
+                                               session_info["refresh_token_expires_in"],
+                                               time.time()
+                                               )
+        else:
+            raise UnauthorizedException(response)
 
     def _connect(self):
 
@@ -258,8 +271,8 @@ def _query(self, query, blob_array = []):
         query_msg.json = query_str
 
         # Set Auth token, only when not authenticated before
-        if self.session:
-            query_msg.token = self.session.session_token
+        if self.shared_data.session:
+            query_msg.token = self.shared_data.session.session_token
 
         for blob in blob_array:
             query_msg.blobs.append(blob)
@@ -280,21 +293,53 @@ def _query(self, query, blob_array = []):
         return (self.last_response, response_blob_array)
 
     def query(self, q, blobs=[]):
-
-        self._check_session_status()
-
+        """
+        Query the database with a query string or a json object.
+        First it checks if the session is valid, if not, it refreshes the token.
+        Then it sends the query to the server and returns the response.
+
+        Args:
+            q (json): native query to be sent
+            blobs (list, optional): Blobs if needed with the query. Defaults to [].
+
+        Raises:
+            ConnectionError: Fatal error, connection to server lost
+
+        Returns:
+            _type_: _description_
+        """
+        self._renew_session()
         try:
             start = time.time()
             self.response, self.blobs = self._query(q, blobs)
+            if not isinstance(self.response, list) and self.response["info"] == "Not Authenticated!":
+                # The case where session is valid, but expires while query is sent.
+                # Hope is that the query send won't be longer than the session ttl.
+                logger.warn(
+                    f"Session expired while query was sent. Retrying... \r\n{traceback.format_stack(limit=5)}")
+                self._renew_session()
+                start = time.time()
+                self.response, self.blobs = self._query(q, blobs)
             self.last_query_time = time.time() - start
             return self.response, self.blobs
         except BaseException as e:
-            print(e)
+            logger.critical(e)
             raise ConnectionError("ApertureDB disconnected")
 
-    def create_new_connection(self):
-
-        return Connector(self.host, self.port, session=self.session)
+    def _renew_session(self):
+        count = 0
+        while count < 3:
+            try:
+                self._check_session_status()
+                break
+            except UnauthorizedException as e:
+                logger.warn(
+                    f"[Attempt {count + 1} of 3] Failed to refresh token. Details: \r\n{traceback.format_exc(limit=5)}")
+                time.sleep(1)
+                count += 1
+
+    def create_new_connection(self) -> Connector:
+        return Connector(self.host, self.port, shared_data=self.shared_data)
 
     def get_last_response_str(self):
 

aperturedb/__init__.py

@@ -2,15 +2,17 @@
 import logging
 import datetime
 import os
+import json
+import requests
 
 logger = logging.getLogger(__name__)
 
-__version__ = "0.3.1"
+__version__ = "0.3.2"
 
 # set log level
 logger.setLevel(logging.DEBUG)
 formatter    = logging.Formatter(
-    "%(asctime)s : %(levelname)s : %(name)s : %(lineno)d : %(message)s")
+    "%(asctime)s : %(levelname)s : %(name)s : %(thread)d : %(lineno)d : %(message)s")
 
 log_file_level = logging.getLevelName(os.getenv("LOG_FILE_LEVEL", "WARN"))
 log_console_level = logging.getLevelName(
@@ -27,3 +29,16 @@
 error_console_handler.setLevel(log_console_level)
 error_console_handler.setFormatter(formatter)
 logger.addHandler(error_console_handler)
+
+try:
+    latest_version = json.loads(requests.get(
+        "https://pypi.org/pypi/aperturedb/json").text)["info"]["version"]
+except Exception as e:
+    logger.warning(
+        f"Failed to get latest version: {e}. You are using version {__version__}")
+    latest_version = None
+if __version__ != latest_version:
+    logger.warning(
+        f"The latest version of aperturedb is {latest_version}. You are using version {__version__}. It is recommended to upgrade.")
+    print(
+        f"The latest version of aperturedb is {latest_version}. You are using version {__version__}. It is recommended to upgrade.")

ci.sh

@@ -29,7 +29,7 @@ update_version() {
     fi
     echo "Updating version $BUILD_VERSION to $VERSION_BUMP"
     # Replace version in __init__.py
-    printf '%s\n' "%s/__version__.*/__version__ = \"$VERSION_BUMP\"/g" 'x' | ex aperturedb/__init__.py
+    printf '%s\n' "%s/__version__ = .*/__version__ = \"$VERSION_BUMP\"/g" 'x' | ex aperturedb/__init__.py
     printf '%s\n' "%s/version=.*/version=\"$VERSION_BUMP\",/g" 'x' | ex setup.py
 
     # Commit and push version bump
@@ -65,7 +65,7 @@ IMAGE_EXTENSION_WITH_VERSION="-${BRANCH_NAME}:v${BUILD_VERSION}"
 if [ $BRANCH_NAME == 'master' ] || [ $BRANCH_NAME == 'main' ]
 then
     # when merging to master remove branch name
-    IMAGE_EXTENSION_WITH_VERSION=":v${APP_VERSION}"
+    IMAGE_EXTENSION_WITH_VERSION=":v${BUILD_VERSION}"
     IMAGE_EXTENSION_LATEST=":latest"
 fi
 
@@ -79,11 +79,12 @@ echo "Repository: $DOCKER_REPOSITORY"
 # Build notebook image
 build_notebook_image(){
     NOTEBOOK_IMAGE=$DOCKER_REPOSITORY/aperturedb-notebook${IMAGE_EXTENSION_WITH_VERSION}
+    LATEST_IMAGE=$DOCKER_REPOSITORY/aperturedb-notebook${IMAGE_EXTENSION_LATEST}
     echo "Building image $NOTEBOOK_IMAGE"
-    docker build -t $NOTEBOOK_IMAGE -f docker/notebook/Dockerfile .
+    docker build -t $NOTEBOOK_IMAGE -t $LATEST_IMAGE -f docker/notebook/Dockerfile .
     if [ -z ${NO_PUSH+x} ]
     then
-        docker push $NOTEBOOK_IMAGE
+        docker push --all-tags $NOTEBOOK_IMAGE
     fi
 }
 
@@ -95,11 +96,12 @@ build_docs_image(){
     cp -r ./docs/{*.*,Makefile,_static} docs/docker/build/docs
     find examples/ -name *.ipynb | xargs -i cp {} docs/docker/build/examples
     DOCS_IMAGE=$DOCKER_REPOSITORY/aperturedb-python-docs${IMAGE_EXTENSION_WITH_VERSION}
+    LATEST_IMAGE=$DOCKER_REPOSITORY/aperturedb-python-docs${IMAGE_EXTENSION_LATEST}
     echo "Building image $DOCS_IMAGE"
     docker build -t $DOCS_IMAGE -f docs/docker/Dockerfile .
     if [ -z ${NO_PUSH+x} ]
     then
-        docker push $DOCS_IMAGE
+        docker push --all-tags $DOCS_IMAGE
     fi
 }
 

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="aperturedb",
-    version="0.3.1",
+    version="0.3.2",
     description="ApertureDB Client Module",
     install_requires=['protobuf>=3.20.0', 'scikit-image', 'image', 'requests', 'boto3',
                       'opencv-python', 'numpy', 'matplotlib', 'pandas', 'kaggle', 'google-cloud-storage'],

test/conftest.py

@@ -13,15 +13,15 @@
 import dbinfo
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def db():
     return Connector(
         port=dbinfo.DB_PORT,
         user=dbinfo.DB_USER,
         password=dbinfo.DB_PASSWORD)
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def insert_data_from_csv(db):
     def insert_data_from_csv(in_csv_file, rec_count=-1):
         file_data_pair = {
@@ -53,7 +53,7 @@ def insert_data_from_csv(in_csv_file, rec_count=-1):
     return insert_data_from_csv
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def utils(db):
     return Utils(db)
 

test/test_Session.py

@@ -20,14 +20,14 @@ def test_sessionRenew(self):
         try:
             db = self.create_connection()
             # force session token expiry
-            db.session.session_token_ttl = 1
+            db.shared_data.session.session_token_ttl = 1
             logging.debug("Connected? {0}".format(
                 "yes" if db.connected else "no"))
             logging.debug(
                 "Session valid? {0}".format(
-                    "yes" if db.session.valid() else "no"))
+                    "yes" if db.shared_data.session.valid() else "no"))
             logging.debug("Valid length: {0}".format(
-                db.session.session_token_ttl))
+                db.shared_data.session.session_token_ttl))
             time.sleep(2)
             query = [{
                 "FindImage": {
@@ -38,6 +38,7 @@ def test_sessionRenew(self):
             }]
             responses, blobs = db.query(query)
             logging.debug(responses)
-            self.assertTrue(db.session.valid(), "Failed to renew Session")
+            self.assertTrue(db.shared_data.session.valid(),
+                            "Failed to renew Session")
         except Exception:
             self.fail("Failed to renew Session")

test/test_torch_connector.py

@@ -60,7 +60,7 @@ def test_nativeContraints(self, db, utils, images):
         assert len(dataset) == utils.count_images()
         self.validate_dataset(dataset)
 
-    def test_datasetWithMultiprocessing(self, db, utils, images):
+    def test_datasetWithMultiprocessing(self, db, utils):
         query = [{
             "FindImage": {
                 "constraints": {