Merge pull request #591 from aperture-data/release-0.4.47

Author: gsaluja9

aperturedb/Configuration.py

@@ -4,8 +4,14 @@
 import re
 from base64 import b64encode, b64decode
 
-APERTURE_CLOUD = ".cloud.aperturedata.io"
-AD_CLOUD_SCHEME = "adbc://"
+APERTUREDB_CLOUD = ".cloud.aperturedata.io"
+APERTUREDB_KEY_VERSION = 1
+FLAG_USE_COMPRESSED_HOST = 4
+FLAG_USE_REST = 2
+FLAG_USE_SSL = 1
+DEFAULT_HTTP_PORT = 80
+DEFAULT_HTTPS_PORT = 443
+DEFAULT_TCP_PORT = 55555
 
 
 @dataclass(repr=False)
@@ -34,16 +40,8 @@ def __repr__(self) -> str:
         return f"[{self.host}:{self.port} as {self.username} using {mode} with SSL={self.use_ssl} auth={auth_mode}]"
 
     def deflate(self) -> list:
-        deflate_version = 1
-        host = self.host
-        if host.endswith(APERTURE_CLOUD):
-            host = "adbc://{}".format(host[:-1 * len(APERTURE_CLOUD)])
-        as_list = [deflate_version, host, self.port, self.username, self.password, self.name, int(self.use_ssl),
-                   int(self.use_rest), int(self.use_keepalive),
-                   self.retry_interval_seconds, self.retry_max_attempts]
-        simplified = json.dumps(as_list)
-        encoded_key = b64encode(simplified.encode('utf-8')).decode('utf-8')
-        return encoded_key
+        return self.create_aperturedb_key(self.host, self.port, self.token,
+                                          self.use_rest, self.use_ssl, self.username, self.password)
 
     def has_user_keys(self) -> bool:
         return self.user_keys is not None
@@ -65,76 +63,90 @@ def set_user_keys(self, keys: dict) -> None:
         self.user_keys = keys
 
     @classmethod
-    def create_web_token(cls, host: str, port: int,  token_string: str) -> None:
-        if token_string.startswith("adbp_"):
+    def config_to_key_type(cls, compressed_host: bool,  use_rest: bool, use_ssl: bool):
+        return (FLAG_USE_COMPRESSED_HOST if compressed_host else 0) + \
+               (FLAG_USE_REST if use_rest else 0) + \
+               (FLAG_USE_SSL if use_ssl else 0)
+
+    @classmethod
+    def key_type_to_config(cls, key_type: int): \
+        return [bool(key_type & FLAG_USE_COMPRESSED_HOST),
+                bool(key_type & FLAG_USE_REST),
+                bool(key_type & FLAG_USE_SSL)]
+
+    @classmethod
+    def config_default_port(cls, use_rest: bool, use_ssl: bool):
+        if use_rest:
+            return DEFAULT_HTTPS_PORT if use_ssl else DEFAULT_HTTP_PORT
+        else:
+            return DEFAULT_TCP_PORT
+
+    @classmethod
+    def create_aperturedb_key(cls, host: str, port: int,  token_string: str,
+                              use_rest: bool, use_ssl: bool, username: str = None, password: str = None) -> None:
+        compressed = False
+        if token_string is not None and token_string.startswith("adbp_"):
             token_string = token_string[5:]
 
-        if host.endswith(APERTURE_CLOUD):
-            host = host[:-1 * len(APERTURE_CLOUD)]
+        if host.endswith(APERTUREDB_CLOUD):
+            host = host[:-1 * len(APERTUREDB_CLOUD)]
             m = re.match("(.*)\.farm(\d+)$", host)
             if m is not None:
                 host = "{}.{}".format(m.group(1), int(m.group(2)))
-
-            host = "a://{}".format(host)
-
-        if port == 55555:
-            web_token_json = [2, host, token_string]
+                compressed = True
+
+        key_type = cls.config_to_key_type(compressed, use_rest, use_ssl)
+        default_port = cls.config_default_port(use_rest, use_ssl)
+        if port != default_port:
+            host = f"{host}:{port}"
+        if token_string is not None:
+            key_json = [APERTUREDB_KEY_VERSION, key_type, host, token_string]
         else:
-            web_token_json = [2, host, port, token_string]
-        simplified = json.dumps(web_token_json)
+            key_json = [APERTUREDB_KEY_VERSION,
+                        key_type, host, username, password]
+        simplified = json.dumps(key_json, separators=(',', ':'))
         encoded = b64encode(simplified.encode('utf-8')).decode('utf-8')
         return encoded
 
     @classmethod
     def reinflate(cls, encoded_str: list) -> object:
+        name = "default_key"
         try:
             decoded = b64decode(encoded_str.encode('utf-8'))
             as_list = json.loads(decoded.decode('utf-8'))
         except:
             raise Exception(
                 "Unable to make configuration from the provided string")
         version = as_list[0]
-        if version not in (1, 2):
+        if version not in (APERTUREDB_KEY_VERSION,):
             raise ValueError("version identifier of configuration was"
-                             f"g{as_list[0]}, which is not supported")
-        if version == 1:
-            host, port, username, password, name, use_ssl, \
-                use_rest, use_keepalive, retry_interval_seconds, \
-                retry_max_attempts = as_list[1:]
-            if host.startswith(AD_CLOUD_SCHEME):
-                host = host[len(AD_CLOUD_SCHEME):] + APERTURE_CLOUD
-            use_ssl = bool(use_ssl)
-            use_rest = bool(use_rest)
-            use_keepaliave = bool(use_keepalive)
-            c = Configuration(
-                host, port, username, password, name, use_ssl, use_rest, use_keepalive,
-                retry_interval_seconds, retry_max_attempts)
-        elif version == 2:
-            host = as_list[1]
-            if host.startswith("a://"):
-                m = re.match("a://([^.]*)\.(\d+)", host)
-                host = "{}.farm{:04d}.cloud.aperturedata.io".format(
-                    m.group(1), int(m.group(2)))
-
-            cur_arg = 2
-            # second arg is port
-            if isinstance(as_list[2], int):
-                cur_arg = 3
-            else:
-                port = 55555
-
-            name = "default"
-
-            username = None
-            password = None
-            token = None
-            # if 1 argument left, treat as token.
-            if len(as_list) - cur_arg == 1:
-                token = "adbp_" + as_list[cur_arg]
-            else:
-                username = as_list[cur_arg]
-                password = as_list[cur_arg + 1]
-
-            c = Configuration(host, port, username,
-                              password, name, token = token)
+                             f"{version}, which is not supported")
+        is_compressed, use_rest, use_ssl = cls.key_type_to_config(as_list[1])
+        host = as_list[2]
+        token = username = password = None
+        if len(as_list) == 4:
+            token = "adbp_" + as_list[3]
+        elif len(as_list) == 5:
+            username, password = as_list[3:]
+
+        port_match = re.match(".*:(\d+)$", host)
+        if port_match is not None:
+            port = int(port_match.group(1))
+            host = host[:-1 * (len(port_match.group(1)) + 1)]
+        else:
+            port = cls.config_default_port(use_rest, use_ssl)
+
+        if is_compressed:
+            try:
+                first, second = host.split('.')
+                host = "{}.farm{:04d}{}".format(
+                    first, int(second), APERTUREDB_CLOUD)
+            except Exception as e:
+                raise ValueError(
+                    f"Unable to parse compressed host: {host} Error: {e}")
+
+        c = Configuration(
+            host, port, username, password, name, use_ssl, use_rest)
+        if token:
+            c.token = token
         return c

aperturedb/Connector.py

@@ -130,6 +130,10 @@ def __init__(self, host="localhost", port=55555,
 
         if key is not None:
             self.config = Configuration.reinflate(key)
+            self.host = self.config.host
+            self.port = self.config.port
+            self.use_ssl = self.config.use_ssl
+            token = self.config.token
         elif config is None:
             self.host = host
             self.port = port
@@ -143,6 +147,7 @@ def __init__(self, host="localhost", port=55555,
                 username=user,
                 password=password,
                 name="runtime",
+                token=token,
                 use_keepalive=use_keepalive,
                 retry_interval_seconds=retry_interval_seconds,
                 retry_max_attempts=retry_max_attempts
@@ -153,6 +158,7 @@ def __init__(self, host="localhost", port=55555,
             self.port = config.port
             self.use_ssl = config.use_ssl
             self.use_keepalive = config.use_keepalive
+            token = config.token
 
         self.conn = None
 
@@ -530,13 +536,8 @@ def clone(self) -> Connector:
             Connector: Clone of original Connector
         """
         return type(self)(
-            self.host,
-            self.port,
-            self.config.username,
-            self.config.password,
-            self.token,
-            use_ssl=self.use_ssl,
-            shared_data=self.shared_data)
+            shared_data=self.shared_data,
+            config=self.config)
 
     def create_new_connection(self):
         from aperturedb.CommonLibrary import issue_deprecation_warning

aperturedb/ConnectorRest.py

@@ -91,9 +91,7 @@ def __init__(self, host="localhost", port=None,
         self.last_query_time = 0
 
         self.url = ('https' if self.use_ssl else 'http') + \
-            '://' + host + ':' + str(self.port) + '/api/'
-
-        self.token = token
+            '://' + self.host + ':' + str(self.port) + '/api/'
 
     def __del__(self):
         logger.info("Done with connector REST.")

aperturedb/ParallelQuery.py

@@ -241,7 +241,7 @@ def worker(self, thid: int, generator, start: int, end: int, run_event) -> None:
                 self.error_counter += 1
 
             if self.stats:
-                self.pb.update(self.batchsize)
+                self.pb.update(batch_end - batch_start)
         logger.info(f"Worker {thid} executed {total_batches} batches")
 
     def get_objects_existed(self) -> int:

aperturedb/Parallelizer.py

@@ -59,7 +59,7 @@ def batched_run(self, generator, batchsize: int, numthreads: int, stats: bool):
         else:
             self.total_actions = len(generator)
         self.pb = tqdm(total=self.total_actions, desc="Progress",
-                       unit="batches", unit_scale=True, dynamic_ncols=True)
+                       unit="items", unit_scale=True, dynamic_ncols=True)
         start_time = time.time()
 
         if self.total_actions < batchsize:

aperturedb/__init__.py

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

aperturedb/cli/keys.py

@@ -1,28 +1,28 @@
-from enum import Enum
 from typing import Annotated
 
 import typer
 
 from aperturedb.cli.console import console
 from aperturedb.Configuration import Configuration
 from aperturedb.Connector import Connector
-from aperturedb.CommonLibrary import create_connector
-from aperturedb.Utils import Utils
 
 app = typer.Typer()
 
 
 @app.command(help="Create Key for a user")
 def generate(user: Annotated[str, typer.Argument(help="The user to generate a key for")]):
+    from aperturedb.CommonLibrary import create_connector
     conn = create_connector()
     key = generate_user_key(conn, user)
     console.log(f"Key for {user} is", key, highlight=False)
 
 
 def generate_user_key(conn: Connector, user: str):
+    from aperturedb.Utils import Utils
     u = Utils(conn)
     token = u.generate_token()
     u.assign_token(user, token)
-    key = Configuration.create_web_token(
-        conn.config.host, conn.config.port, token)
+    key = Configuration.create_aperturedb_key(
+        conn.config.host, conn.config.port, token, conn.config.use_rest,
+        conn.config.use_ssl)
     return key

test/test_Key.py

@@ -0,0 +1,81 @@
+import logging
+import typer
+
+from aperturedb.Configuration import Configuration
+
+logger = logging.getLogger(__name__)
+
+key_pairs = {
+    "WzEsMSwibG9jYWxob3N0IiwiYWRtaW4iLCJhZG1pbiJd":
+    [1, 1, "localhost", "admin", "admin"],
+        "WzEsMCwiMTI3LjAuMC4xOjU1NTU0IiwiWVFadVZVV2Zab0FkWjJrUUVMeFB5RnptZHJ3WXd0cjBBRGEiXQ==":
+            [1, 0, "127.0.0.1:55554", "YQZuVUWfZoAdZ2kQELxPyFzmdrwYwtr0ADa"],
+        "WzEsNywid29ya2Zsb3ctbG9hZGVkLWZvMWphdTN0LjAiLCJhZG1pbiIsIjEyMzRCVFFMUF8lMnR0Il0=":
+            [1, 7, "workflow-loaded-fo1jau3t.farm0000.cloud.aperturedata.io",
+                "admin", "1234BTQLP_%2tt"],
+        "WzEsNSwidGVzdC0zcWpxdDZrcy40IiwiWVFadVZVV2Zab0FkWjJrUUVMeFB5RnptZHJ3WXd0cjBBRGEiXQ==":
+            [1, 5, "test-3qjqt6ks.farm0004.cloud.aperturedata.io",
+                "YQZuVUWfZoAdZ2kQELxPyFzmdrwYwtr0ADa"],
+        "WzEsMiwiMTkyLjE2OC40LjEyOjU1NTU1IiwiYWRtaW4iLCJhZG1pbiJd":
+        [1, 2, "192.168.4.12:55555", "admin", "admin"],
+        "WzEsMywiYXBlcnR1cmVkYi5iaWdjb3JwLmlvOjE5MTgiLCJZUVp1VlVXZlpvQWRaMmtRRUx4UHlGem1kcndZd3RyMEFEYSJd":
+        [1, 3, "aperturedb.bigcorp.io:1918", "YQZuVUWfZoAdZ2kQELxPyFzmdrwYwtr0ADa"],
+        "WzEsNCwidGNwLTU1N2Vwbm4zLjkwOToxOTE4IiwiYWRtaW4iLCI4OTBFcE1uKyElMiRfIl0=":
+        [1, 4, "tcp-557epnn3.farm0909.cloud.aperturedata.io:1918",
+            "admin", "890EpMn+!%2$_"],
+        "WzEsNiwiaHR0cC05MGpnM3pwcy4xMjo0NDMiLCJZUVp1VlVXZlpvQWRaMmtRRUx4UHlGem1kcndZd3RyMEFEYSJd":
+        [1, 6, "http-90jg3zps.farm0012.cloud.aperturedata.io:443",
+            "YQZuVUWfZoAdZ2kQELxPyFzmdrwYwtr0ADa"]
+}
+
+
+class TestApertureDBKey():
+
+    def test_encode_keys(self):
+        for key, data in key_pairs.items():
+            logger.info(f"Testing encoding of {key}")
+            config_type = data[1]
+            host = data[2]
+            username = password = token = None
+            comp, rest, ssl = Configuration.key_type_to_config(config_type)
+            if host.rfind(':') != -1:
+                port = int(host.split(':')[1])
+                host = host.split(':')[0]
+            else:
+                port = Configuration.config_default_port(rest, ssl)
+            if len(data) == 4:
+                token = data[3]
+            else:
+                username = data[3]
+                password = data[4]
+            c = Configuration(host, port, username, password,
+                              "encoding test", use_rest=rest, use_ssl=ssl, token=token)
+            deflated = c.deflate()
+            assert deflated == key
+
+    def test_decode_keys(self):
+        for key, data in key_pairs.items():
+            logger.info(f"Testing decoding of {key}")
+            config = Configuration.reinflate(key)
+            config_type = data[1]
+            host = data[2]
+            if config_type == 0 or config_type == 4:
+                assert not config.use_rest and not config.use_ssl
+            if config_type == 1 or config_type == 5:
+                assert not config.use_rest and config.use_ssl
+            if config_type == 2 or config_type == 6:
+                assert config.use_rest and not config.use_ssl
+            if config_type == 3 or config_type == 7:
+                assert config.use_rest and config.use_ssl
+
+            if host.rfind(':') != -1:
+                port = int(host.split(':')[1])
+                host = host.split(':')[0]
+                assert config.port == port
+
+            if len(data) == 4:
+                assert config.token == "adbp_" + data[3]
+            else:
+                assert config.username == data[3] and config.password == data[4]
+
+            assert(config.host == host)