Merge pull request #604 from aperture-data/release-0.4.49

Release 0.4.49

Includes a change in how GetSchema works for connections with no Classes.

Author: gsaluja9

aperturedb/Connector.py

@@ -49,6 +49,30 @@
 
 
 PROTOCOL_VERSION = 1
+MESSAGE_LENGTH_FORMAT = '@I'
+MESSAGE_LENGTH_SIZE = struct.calcsize(MESSAGE_LENGTH_FORMAT)
+
+# Protocol types
+PROTOCOL_TCP = 1
+PROTOCOL_SSL = 2
+
+DEFAULT_PORT = 55555
+# aperturedb's param ADB_MAX_CONNECTION_MESSAGE_SIZE_MB = 2048 by default
+DEFAULT_MAX_MESSAGE_SIZE_MB = 2048
+
+# Query retry constants
+DEFAULT_RETRY_INTERVAL_SECONDS = 1
+DEFAULT_RETRY_MAX_ATTEMPTS = 3
+DEFAULT_SESSION_EXPIRY_OFFSET_SEC = 10
+DEFAULT_QUERY_CONNECTION_ERROR_SUPPRESSION_DELTA_SEC = 30
+
+# Session renewal constants
+RENEW_SESSION_MAX_ATTEMPTS = 3
+RENEW_SESSION_RETRY_INTERVAL_SEC = 1
+
+# Status codes
+STATUS_OK = 0
+STATUS_ERROR_DEFAULT = -2
 
 
 class UnauthorizedException(Exception):
@@ -73,7 +97,7 @@ def valid(self) -> bool:
 
         # This triggers refresh if the session is about to expire.
         if session_age > self.session_token_ttl - \
-                int(os.getenv("SESSION_EXPIRY_OFFSET_SEC", 10)):
+                int(os.getenv("SESSION_EXPIRY_OFFSET_SEC", DEFAULT_SESSION_EXPIRY_OFFSET_SEC)):
             return False
 
         return True
@@ -106,14 +130,14 @@ class Connector(object):
         str (key): Apeture Key, configuration as a deflated compressed string
     """
 
-    def __init__(self, host="localhost", port=55555,
+    def __init__(self, host="localhost", port=DEFAULT_PORT,
                  user="", password="", token="",
                  use_ssl=True,
                  shared_data=None,
                  authenticate=True,
                  use_keepalive=True,
-                 retry_interval_seconds=1,
-                 retry_max_attempts=3,
+                 retry_interval_seconds=DEFAULT_RETRY_INTERVAL_SECONDS,
+                 retry_max_attempts=DEFAULT_RETRY_MAX_ATTEMPTS,
                  config: Optional[Configuration] = None,
                  key: Optional[str] = None):
         """
@@ -126,7 +150,8 @@ def __init__(self, host="localhost", port=55555,
         self.last_query_timestamp = None
         # suppress connection warnings which occur more than this time
         # after the last query
-        self.query_connection_error_suppression_delta = timedelta(seconds=30)
+        self.query_connection_error_suppression_delta = timedelta(
+            seconds=DEFAULT_QUERY_CONNECTION_ERROR_SUPPRESSION_DELTA_SEC)
 
         if key is not None:
             self.config = Configuration.reinflate(key)
@@ -196,20 +221,20 @@ def __del__(self):
             self.connected = False
 
     def _send_msg(self, data):
-        # aperturedb's param ADB_MAX_CONNECTION_MESSAGE_SIZE_MB = 256 by default
-        if len(data) > (256 * 2**20):
+        if len(data) > (DEFAULT_MAX_MESSAGE_SIZE_MB * 2**20):
             logger.warning(
                 "Message sent is larger than default for ApertureDB Server. Server may disconnect.")
 
-        sent_len = struct.pack('@I', len(data))  # send size first
+        sent_len = struct.pack(MESSAGE_LENGTH_FORMAT,
+                               len(data))  # send size first
         x = self.conn.send(sent_len + data)
-        return x == len(data) + 4
+        return x == len(data) + MESSAGE_LENGTH_SIZE
 
     def _recv_msg(self):
-        recv_len = self.conn.recv(4)  # get message size
+        recv_len = self.conn.recv(MESSAGE_LENGTH_SIZE)  # get message size
         if recv_len == b'':
             return None
-        recv_len = struct.unpack('@I', recv_len)[0]
+        recv_len = struct.unpack(MESSAGE_LENGTH_FORMAT, recv_len)[0]
         response = bytearray(recv_len)
         read = 0
         while read < recv_len:
@@ -250,7 +275,7 @@ def _authenticate(self, user, password="", token=""):
                 "Unexpected response from server upon authenticate request: " +
                 str(response))
         session_info = response[0]["Authenticate"]
-        if session_info["status"] != 0:
+        if session_info["status"] != STATUS_OK:
             raise Exception(session_info["info"])
 
         self.shared_data.session = Session(
@@ -280,7 +305,7 @@ def _refresh_token(self):
         logger.info(f"Refresh token response: \r\n{response}")
         if isinstance(response, list):
             session_info = response[0]["RefreshToken"]
-            if session_info["status"] != 0:
+            if session_info["status"] != STATUS_OK:
                 # Refresh token failed, we need to re-authenticate
                 # This is possible with a long lived connector, where
                 # the session token and the refresh token have expired.
@@ -320,7 +345,7 @@ def _connect(self):
 
             # Handshake with server to negotiate protocol
 
-            protocol = 2 if self.use_ssl else 1
+            protocol = PROTOCOL_SSL if self.use_ssl else PROTOCOL_TCP
 
             hello_msg = struct.pack('@II', PROTOCOL_VERSION, protocol)
 
@@ -513,16 +538,16 @@ def query(self, q, blobs=[]):
 
     def _renew_session(self):
         count = 0
-        while count < 3:
+        while count < RENEW_SESSION_MAX_ATTEMPTS:
             try:
                 self._check_session_status()
                 break
             except UnauthorizedException as e:
                 logger.warning(
-                    f"[Attempt {count + 1} of 3] Failed to refresh token.",
+                    f"[Attempt {count + 1} of {RENEW_SESSION_MAX_ATTEMPTS}] Failed to refresh token.",
                     exc_info=True,
                     stack_info=True)
-                time.sleep(1)
+                time.sleep(RENEW_SESSION_RETRY_INTERVAL_SEC)
                 count += 1
 
     def clone(self) -> Connector:
@@ -583,7 +608,7 @@ def check_status(self, json_res: CommandResponses) -> int:
             int: The value recieved from the server, or -2 if not found.
         """
         # Default status is -2, which is an error, but not a server error.
-        status = -2
+        status = STATUS_ERROR_DEFAULT
         if (isinstance(json_res, dict)):
             if ("status" not in json_res):
                 status = self.check_status(json_res[list(json_res.keys())[0]])

aperturedb/ParallelLoader.py

@@ -1,11 +1,19 @@
-
 from aperturedb import ParallelQuery
 from aperturedb.Connector import Connector
 from aperturedb.Utils import Utils
 from aperturedb.Subscriptable import Subscriptable
 
 import numpy as np
 import logging
+
+# For each property of each Entity or Connection,
+# the following information is returned as an array of 3 elements of the form [matched, has_index_flag, type]:
+# - matched: Number of objects that match the search.
+# - has_index_flag: Indicates whether the property is indexed or not.
+# - type: Type for the property. See supported types.
+# https://docs.aperturedata.io/query_language/Reference/db_commands/GetSchema
+PROPERTIES_SCHEMA_INDEX_FLAG = 1
+
 logger = logging.getLogger(__name__)
 
 
@@ -23,20 +31,92 @@ def __init__(self, client: Connector, dry_run: bool = False):
         self.utils = Utils(self.client)
         self.type = "element"
 
+    def get_entity_indexes(self, schema: dict) -> dict:
+        """
+        Returns a dictionary of indexes for entities' properties.
+
+        Args:
+            schema (dict): The schema dictionary to get indexes from.
+
+        Returns:
+            dict: A dictionary of entity indexes.
+        """
+
+        indexes = {}
+
+        entities = schema.get("entities") or {}
+
+        for cls_name, cls_schema in (entities.get("classes") or {}).items():
+            for prop_name, prop_schema in (cls_schema.get("properties") or {}).items():
+                if prop_schema[PROPERTIES_SCHEMA_INDEX_FLAG]:
+                    indexes.setdefault("entity", {}).setdefault(
+                        cls_name, set()).add(prop_name)
+
+        return indexes
+
+    def get_connection_indexes(self, schema: dict) -> dict:
+        """
+        Returns a dictionary of indexes for connections' properties.
+
+        Args:
+            schema (dict): The schema dictionary to get indexes from.
+
+        Returns:
+            dict: A dictionary of connection indexes.
+        """
+
+        indexes = {}
+
+        connections = schema.get("connections") or {}
+        cls_names = connections.get("classes") or {}
+
+        for cls_name, cls_schema in cls_names.items():
+
+            # check if cls_schema is a dict or a list
+            if isinstance(cls_schema, dict):
+                for prop_name, prop_schema in (cls_schema.get("properties") or {}).items():
+                    if prop_schema[PROPERTIES_SCHEMA_INDEX_FLAG]:
+                        indexes.setdefault("connection", {}).setdefault(
+                            cls_name, set()).add(prop_name)
+            elif isinstance(cls_schema, list):
+                # If cls_schema is a list, this occurs when the schema defines multiple connection variants
+                # for the same class. Each element in the list is expected to be a dictionary representing
+                # a connection variant, with a "properties" key containing the property schemas.
+                # Example schema format:
+                # "connections": {
+                #     "classes": {
+                #         "SomeConnectionClass": [
+                #             {"properties": {"prop1": [...], "prop2": [...]}},
+                #             {"properties": {"prop3": [...], "prop4": [...]}},
+                #         ]
+                #     }
+                # }
+                for connection in cls_schema:
+                    for prop_name, prop_schema in (connection.get("properties") or {}).items():
+                        if prop_schema[PROPERTIES_SCHEMA_INDEX_FLAG]:
+                            indexes.setdefault("connection", {}).setdefault(
+                                cls_name, set()).add(prop_name)
+            else:
+                exception_msg = "Unexpected schema format for connection class "
+                exception_msg += f"'{cls_name}': {cls_schema}"
+                logger.error(exception_msg)
+                raise ValueError(exception_msg)
+
+        return indexes
+
     def get_existing_indices(self):
+
+        indexes = {}
         schema = self.utils.get_schema()
-        existing_indices = {}
+
         if schema:
-            for index_type in (("entity", "entities"), ("connection", "connections")):
-                foo = schema.get(index_type[1]) or {}
-                bar = foo.get("classes") or {}
-                for cls_name, cls_schema in bar.items():
-                    props = cls_schema.get("properties") or {}
-                    for prop_name, prop_schema in props.items():
-                        if prop_schema[1]:  # indicates property has an index
-                            existing_indices.setdefault(index_type[0], {}).setdefault(
-                                cls_name, set()).add(prop_name)
-        return existing_indices
+            entity_indexes     = self.get_entity_indexes(schema)
+            connection_indexes = self.get_connection_indexes(schema)
+
+            # Combine both entity and connection indexes
+            indexes = {**entity_indexes, **connection_indexes}
+
+        return indexes
 
     def query_setup(self, generator: Subscriptable) -> None:
         """

aperturedb/__init__.py

@@ -10,7 +10,7 @@
 import signal
 import sys
 
-__version__ = "0.4.48"
+__version__ = "0.4.49"
 
 logger = logging.getLogger(__name__)