fix: use sendall() to prevent partial sends and reconnect if conn is None before query

.dockerignore

@@ -0,0 +1,156 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# VSCode
+.vscode/
+
+#Data files
+*.adb.csv
+*.jpg
+*.npy
+test/aperturedb/db*/
+test/input/blobs/
+docs/examples/
+examples/*/coco
+examples/*/classification.txt
+kaggleds/
+examples/*/kaggleds/
+docs/*/*.svg
+test/aperturedb/log*
+adb-python/*
+/test/input/
+/test/input/images/
+
+.aperturedb
+test/data/
+test/aperturedb/certificate*/
+.devcontainer/aperturedb/
+.devcontainer/ca/
+test/*_ca/

.github/workflows/pr.yaml

@@ -13,6 +13,10 @@ jobs:
 
     steps:
 
+    - name: Cleanup previous run
+      run: docker run --rm -v ${{ github.workspace }}:/workspace alpine sh -c "rm -rf /workspace/test/aperturedb/db*"
+      continue-on-error: true
+
     - uses: actions/checkout@v3
 
     - name: Login to DockerHub

aperturedb/Connector.py

@@ -236,8 +236,7 @@ def _send_msg(self, data):
 
         sent_len = struct.pack(MESSAGE_LENGTH_FORMAT,
                                len(data))  # send size first
-        x = self.conn.send(sent_len + data)
-        return x == len(data) + MESSAGE_LENGTH_SIZE
+        self.conn.sendall(sent_len + data)
 
     def _recv_msg(self):
         recv_len = self.conn.recv(MESSAGE_LENGTH_SIZE)  # get message size
@@ -451,7 +450,8 @@ def connect(self, details: str = None):
                 self._connect()
             except socket.error as e:
                 logger.error(
-                    f"Error connecting to server: {self.config} \r\n{details}. {e=}",
+                    f"Error connecting to server: "
+                    f"{self.config} \r\n{details}. {e=}",
                     exc_info=True,
                     stack_info=True)
 
@@ -477,18 +477,21 @@ def _query(self, query, blob_array = [], try_resume=True):
         # Serialize with protobuf and send
         data = query_msg.SerializeToString()
 
+        if self.conn is None:
+            self.connect(details="Initial connect from _query")
+
         # this is for session refresh attempts
         tries = 0
         while tries < self.config.retry_max_attempts:
             try:
-                if self._send_msg(data):
-                    response = self._recv_msg()
-                    if response is not None:
-                        querRes = queryMessage.queryMessage()
-                        queryMessage.ParseFromString(querRes, response)
-                        response_blob_array = [b for b in querRes.blobs]
-                        self.last_response = json.loads(querRes.json)
-                        break
+                self._send_msg(data)
+                response = self._recv_msg()
+                if response is not None:
+                    querRes = queryMessage.queryMessage()
+                    queryMessage.ParseFromString(querRes, response)
+                    response_blob_array = [b for b in querRes.blobs]
+                    self.last_response = json.loads(querRes.json)
+                    break
             except ssl.SSLEOFError as ssle:
                 # this can happen when working in a notebook.
                 # we log if this isn't the first try, or if
@@ -537,7 +540,7 @@ def _query(self, query, blob_array = [], try_resume=True):
             # For example aperturedb server is restarted, or network is lost.
             # While this is useful bit of code, when executed in a refresh token
             # path, this can cause a deadlock. Hence the try_resume flag.
-            if try_resume:
+            if try_resume and self.connected:
                 self._renew_session()
         if tries == self.config.retry_max_attempts:
             # We have tried enough times, and failed. Log some state info.
@@ -565,6 +568,7 @@ def query(self, q, blobs=[]):
         Returns:
             _type_: _description_
         """
+        self._renew_session()
         if self.should_authenticate:
             self.authenticate(
                 shared_data=self.shared_data,
@@ -602,7 +606,9 @@ def _renew_session(self):
                 break
             except UnauthorizedException as e:
                 logger.warning(
-                    f"[Attempt {count + 1} of {RENEW_SESSION_MAX_ATTEMPTS}] Failed to refresh token.",
+                    f"[Attempt {count + 1} of "
+                    f"{RENEW_SESSION_MAX_ATTEMPTS}] "
+                    "Failed to refresh token.",
                     exc_info=True,
                     stack_info=True)
                 time.sleep(RENEW_SESSION_RETRY_INTERVAL_SEC)

test/docker-compose.yml

@@ -14,7 +14,7 @@ services:
       openssl req -new -key /cert/tls.key -out /ca/http.csr -days 3650 -subj \"/C=US/ST=NY/L=NYC/O=instance/OU=instanceDB/CN=${DB_HTTP_CN:-localhost}\"
       openssl x509 -req  -CA /ca/ca.crt -CAkey /ca/ca.key -in /ca/http.csr -out /cert/http.crt -passin pass:1234"
     volumes:
-      - ./aperturedb/certificate:/cert
+      - ./aperturedb/certificate_${RUNNER_NAME}:/cert
       - ./${RUNNER_NAME}_ca:/ca
 
   lenz:
@@ -38,7 +38,7 @@ services:
       LNZ_CERTIFICATE_PATH: /etc/lenz/certificate/tcp.crt
       LNZ_PRIVATE_KEY_PATH: /etc/lenz/certificate/tls.key
     volumes:
-      - ./aperturedb/certificate:/etc/lenz/certificate
+      - ./aperturedb/certificate_${RUNNER_NAME}:/etc/lenz/certificate
 
   aperturedb:
     image: $ADB_REPO:$ADB_TAG
@@ -71,7 +71,7 @@ services:
       - source: nginx.conf
         target: /etc/nginx/conf.d/default.conf
     volumes:
-      - ./aperturedb/certificate:/etc/nginx/certificate
+      - ./aperturedb/certificate_${RUNNER_NAME}:/etc/nginx/certificate
 
 configs:
   nginx.conf:

test/run_test_container.sh

@@ -1,6 +1,5 @@
 #!/bin/bash
 
-set -u
 set -e
 
 function check_containers_networks(){

test/test_Session.py

@@ -114,8 +114,7 @@ def mock_connect(host, port):
             # Check the exception is not an obscure one.
             assert "self.connected=False" in e.args[0]
 
-        # Check that we tried to connect 3 times.
-        assert connect_attempts == 3
+        assert connect_attempts == 4
 
     def test_socket_send_error_initial(self, monkeypatch):
         send_attempts = 0
@@ -124,7 +123,7 @@ def mock_send(x, buff):
             nonlocal send_attempts
             send_attempts += 1
             raise socket.error("Connection broke when send")
-        monkeypatch.setattr(socket.socket, "send", mock_send)
+        monkeypatch.setattr(socket.socket, "sendall", mock_send)
 
         # Create new db connection.
         new_db = Connector(
@@ -141,8 +140,8 @@ def mock_send(x, buff):
             # Check the exception is not an obscure one.
             assert "self.connected=False" in e.args[0]
 
-        # Check that we tried to send 5 (connect hello:2) + query:3) times.
-        assert send_attempts == 5
+        # Should be exactly 7 attempts (1 initial + 3 retries, each doing query send + reconnect send)
+        assert send_attempts == 7
 
     def test_socket_recv_error_initial(self, monkeypatch):
         connect_attempts = 0
@@ -170,11 +169,12 @@ def mock_recv(x, buff):
             # Check the exception is not an obscure one.
             assert "self.connected=False" in e.args[0]
 
-        # Check that we tried to connect 3 times.
-        assert connect_attempts == 3
+        assert connect_attempts == 4
 
     def test_con_close_on_send_query(self, db: Connector, monkeypatch):
         if not isinstance(db, ConnectorRest):
+            if db.conn is None:
+                db.connect()
             original_send_msg = db._send_msg
             count = 0
 
@@ -205,6 +205,8 @@ def mock_send_msg(msg):
 
     def test_con_close_on_recv_query(self, db: Connector, monkeypatch):
         if not isinstance(db, ConnectorRest):
+            if db.conn is None:
+                db.connect()
             original_recv_msg = db._recv_msg
             count = 0