Enabled remote dataset fetching from url (#63)

* [add] Enabled to specify remote rdb links

* [add] Enabled remote dataset fetching from url

Author: filipecosta90

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "redisbench-admin"
-version = "0.1.54"
+version = "0.1.55"
 description = "Redis benchmark run helper. A wrapper around Redis and Redis Modules benchmark tools ( ftsb_redisearch, memtier_benchmark, redis-benchmark, aibench, etc... )."
 authors = ["filipecosta90 <[email protected]>"]
 readme = "README.md"
@@ -29,6 +29,7 @@ python_terraform = "^0.10.1"
 GitPython = "^3.1.12"
 PyYAML = "^5.4.0"
 numpy = "^1.15.4"
+wget = "^3.2"
 
 [tool.poetry.dev-dependencies]
 pytest = "^4.6"

pyproject.toml

@@ -6,6 +6,7 @@
 import shutil
 import subprocess
 import sys
+import tempfile
 
 import yaml
 
@@ -16,7 +17,7 @@
 from redisbench_admin.utils.local import (
     spinUpLocalRedis,
     getLocalRunFullFilename,
-    isProcessAlive, )
+    isProcessAlive, checkDatasetLocalRequirements, )
 from redisbench_admin.utils.remote import (
     extract_git_vars,
     validateResultExpectations,
@@ -85,11 +86,22 @@ def run_local_command_logic(args):
             # after we've spinned Redis, even on error we should always teardown
             # in case of some unexpected error we fail the test
             try:
+                dirname = ".",
                 # setup Redis
+                # copy the rdb to DB machine
+                dataset = None
+                temporary_dir = tempfile.mkdtemp()
+                logging.info(
+                    "Using local temporary dir to spin up Redis Instance. Path: {}".format(
+                        temporary_dir
+                    )
+                )
+                checkDatasetLocalRequirements(benchmark_config, temporary_dir, dirname)
+
                 redis_process = spinUpLocalRedis(
-                    benchmark_config,
+                    temporary_dir,
                     args.port,
-                    local_module_file, dirname,
+                    local_module_file
                 )
                 if isProcessAlive(redis_process) is False:
                     raise Exception("Redis process is not alive. Failing test.")

redisbench_admin/run_local/run_local.py

@@ -4,17 +4,39 @@
 import tempfile
 import time
 from shutil import copyfile
-
+import wget
 import redis
 
 
-def checkDatasetLocalRequirements(benchmark_config, redis_tmp_dir, dirname="."):
-    for k in benchmark_config["dbconfig"]:
-        if "dataset" in k:
-            dataset = k["dataset"]
-    if dataset is not None:
-        logging.info("Copying rdb {}/{} to {}/dump.rdb".format(dirname, dataset, redis_tmp_dir))
-        copyfile("{}/{}".format(dirname, dataset), "{}/dump.rdb".format(redis_tmp_dir))
+def checkDatasetLocalRequirements(benchmark_config, redis_dbdir, dirname=None, datasets_localtemp_dir="./datasets",dbconfig_keyname="dbconfig"):
+    dataset = None
+    full_path = None
+    tmp_path = None
+    if dbconfig_keyname in benchmark_config:
+        for k in benchmark_config[dbconfig_keyname]:
+            if "dataset" in k:
+                dataset = k["dataset"]
+        if dataset is not None:
+            if dataset.startswith("http"):
+                if not os.path.isdir(datasets_localtemp_dir):
+                    os.mkdir(datasets_localtemp_dir)
+                filename = dataset.split("/")[-1]
+                full_path = '{}/{}'.format(datasets_localtemp_dir, filename)
+                if not os.path.exists(full_path):
+                    logging.info("Retrieving remote file from {} to {}. Using the dir {} as a cache for next time.".format(dataset,full_path, datasets_localtemp_dir))
+                    wget.download(dataset,full_path)
+                else:
+                    logging.info(
+                        "Reusing cached remote file (located at {} ).".format(
+                            full_path))
+            else:
+                full_path = dataset
+                if dirname is not None:
+                    full_path = "{}/{}".format(dirname,full_path)
+                logging.info("Copying rdb from {} to {}/dump.rdb".format(full_path, redis_dbdir))
+            tmp_path = "{}/dump.rdb".format(redis_dbdir)
+            copyfile(full_path,tmp_path )
+    return dataset,full_path,tmp_path
 
 
 def waitForConn(conn, retries=20, command="PING", shouldBe=True):
@@ -40,21 +62,25 @@ def waitForConn(conn, retries=20, command="PING", shouldBe=True):
 
 
 def spinUpLocalRedis(
-        benchmark_config,
+        dbdir,
         port,
         local_module_file,
-        dirname=".",
 ):
-    # copy the rdb to DB machine
-    dataset = None
-    temporary_dir = tempfile.mkdtemp()
+    command = generateStandaloneRedisServerArgs(dbdir, local_module_file, port)
+
     logging.info(
-        "Using local temporary dir to spin up Redis Instance. Path: {}".format(
-            temporary_dir
+        "Running local redis-server with the following args: {}".format(
+            " ".join(command)
         )
     )
-    checkDatasetLocalRequirements(benchmark_config, temporary_dir, dirname)
+    redis_process = subprocess.Popen(command)
+    result = waitForConn(redis.StrictRedis())
+    if result is True:
+        logging.info("Redis available")
+    return redis_process
 
+
+def generateStandaloneRedisServerArgs(dbdir, local_module_file, port):
     # start redis-server
     command = [
         "redis-server",
@@ -63,20 +89,12 @@ def spinUpLocalRedis(
         "--port",
         "{}".format(port),
         "--dir",
-        temporary_dir,
-        "--loadmodule",
-        os.path.abspath(local_module_file),
-    ]
-    logging.info(
-        "Running local redis-server with the following args: {}".format(
-            " ".join(command)
-        )
-    )
-    redis_process = subprocess.Popen(command)
-    result = waitForConn(redis.StrictRedis())
-    if result is True:
-        logging.info("Redis available")
-    return redis_process
+        dbdir]
+    if local_module_file is not None:
+        command.extend(["--loadmodule",
+                        os.path.abspath(local_module_file),
+                        ])
+    return command
 
 
 def isProcessAlive(process):

redisbench_admin/utils/local.py

@@ -14,6 +14,8 @@
 from python_terraform import Terraform
 from tqdm import tqdm
 
+from redisbench_admin.utils.local import checkDatasetLocalRequirements
+
 
 def get_git_root(path):
     git_repo = git.Repo(path, search_parent_directories=True)
@@ -105,18 +107,16 @@ def checkDatasetRemoteRequirements(
         benchmark_config, server_public_ip, username, private_key, remote_dataset_file, dirname
 ):
     res = True
-    for k in benchmark_config["dbconfig"]:
-        if "dataset" in k:
-            dataset = k["dataset"]
+    dataset, fullpath, tmppath = checkDatasetLocalRequirements(benchmark_config, ".", dirname)
     if dataset is not None:
         logging.info('Detected dataset config. Will copy file to remote setup... "{}"'.format(dataset))
         res = copyFileToRemoteSetup(
             server_public_ip,
             username,
             private_key,
-            dataset,
+            fullpath,
             remote_dataset_file,
-            dirname,
+            None,
         )
     return res
 

redisbench_admin/utils/remote.py

@@ -0,0 +1,39 @@
+import os
+import shutil
+import tempfile
+
+from redisbench_admin.utils.local import checkDatasetLocalRequirements
+
+
+def test_check_dataset_local_requirements():
+    url = "https://s3.amazonaws.com/benchmarks.redislabs/redistimeseries/tsbs/datasets/devops/functional/scale-100-redistimeseries_data.rdb"
+    # no db config
+    checkDatasetLocalRequirements({}, ".", ".")
+    # dbconfig with no dataset key
+    checkDatasetLocalRequirements({"dbconfig": {}}, ".", ".")
+
+    # dbconfig with local filename
+    directory_from = tempfile.mkdtemp()
+    directory_to = tempfile.mkdtemp()
+    temp_file1 = tempfile.NamedTemporaryFile(dir=directory_from)
+    checkDatasetLocalRequirements({"dbconfig": [{"dataset": temp_file1.name}]}, directory_to.__str__(), None)
+    assert os.path.exists("{}/{}".format(directory_from.__str__(), temp_file1.name.split("/")[-1]))
+    assert os.path.exists("{}/{}".format(directory_to.__str__(), "dump.rdb"))
+    shutil.rmtree(directory_from)
+    shutil.rmtree(directory_to)
+
+    # dbconfig with remote filename
+    tests_remote_tmp_datasets = "./tests/temp-datasets"
+    if os.path.isdir(tests_remote_tmp_datasets):
+        shutil.rmtree(tests_remote_tmp_datasets)
+
+    directory_to = tempfile.mkdtemp()
+    checkDatasetLocalRequirements({"dbconfig": [{"dataset": url}]}, directory_to.__str__(), None, tests_remote_tmp_datasets)
+    assert os.path.exists("{}/{}".format(directory_to.__str__(), "dump.rdb"))
+    assert os.path.exists("{}/{}".format(tests_remote_tmp_datasets, "scale-100-redistimeseries_data.rdb"))
+    checkDatasetLocalRequirements({"dbconfig": [{"dataset": url}]}, directory_to.__str__(),  None, tests_remote_tmp_datasets)
+    assert os.path.exists("{}/{}".format(directory_to.__str__(), "dump.rdb"))
+    shutil.rmtree(directory_to)
+    if os.path.isdir(tests_remote_tmp_datasets):
+        shutil.rmtree(tests_remote_tmp_datasets)
+