graphql-engine/server/bench-wrk/sportsdb_setup.py

import argparse
from zipfile import ZipFile
import os
from contextlib import contextmanager
import threading
import requests
import requests_cache
from colorama import Fore, Style

from port_allocator import PortAllocator
from run_postgres import Postgres
from run_hge import HGE


def _first_true(iterable, default=False, pred=None):
    return next(filter(pred, iterable), default)


class HGETestSetup:

    sportsdb_url='http://www.sportsdb.org/modules/sd/assets/downloads/sportsdb_sample_postgresql.zip'

    default_work_dir = 'test_output'

    previous_work_dir_file = '.previous_work_dir'

    def __init__(self, pg_url, remote_pg_url, pg_docker_image, hge_url, remote_hge_url, hge_docker_image=None, hge_args=[], skip_remote_graphql_setup=False, skip_stack_build=False):
        self.pg_url = pg_url
        self.remote_pg_url = remote_pg_url
        self.pg_docker_image = pg_docker_image
        self.hge_url = hge_url
        self.remote_hge_url = remote_hge_url
        self.hge_docker_image = hge_docker_image
        self.hge_args = hge_args
        self.skip_remote_graphql_setup = skip_remote_graphql_setup
        self.skip_stack_build = skip_stack_build
        self.port_allocator = PortAllocator()
        self.init_work_dir()
        self.init_pgs()
        self.init_hges()
        self.set_previous_work_dir()

    def get_previous_work_dir(self):
        """Get the work directory of the previous error"""
        try:
            with open(self.previous_work_dir_file) as f:
                return f.read()
        except FileNotFoundError:
            return None

    def set_previous_work_dir(self):
        """
        Set the current work directory as previous work directory
        This directory will be used as the default work directory choice
        """
        with open(self.previous_work_dir_file, 'w') as f:
            return f.write(self.work_dir)

    def get_work_dir(self):
        """
        Get the work directory from environmental variable, or from the user input
        """
        default_work_dir = self.get_previous_work_dir() or self.default_work_dir
        return os.environ.get('WORK_DIR') \
            or input(Fore.YELLOW + '(Set WORK_DIR environmental variable to avoid this)\n'
                     + 'Please specify the work directory. (default:{}):'.format(default_work_dir)
                     + Style.RESET_ALL).strip() \
            or default_work_dir

    def init_work_dir(self):
        """
        Get the current work directory from user input or environmental variables
        and create the work directory if it is not present
        """
        self.work_dir = self.get_work_dir()
        print ("WORK_DIR: ", self.work_dir)
        os.makedirs(self.work_dir, exist_ok=True)
        requests_cache.install_cache(self.work_dir + '/sportsdb_cache')

    def init_pgs(self):
        def _init_pg(data_dir, url):
            return Postgres(
                port_allocator=self.port_allocator, docker_image=self.pg_docker_image,
                db_data_dir= self.work_dir + '/' + data_dir, url=url
            )

        self.pg = _init_pg('sportsdb_data', self.pg_url)

        if not self.skip_remote_graphql_setup:
            self.remote_pg = _init_pg('remote_sportsdb_data', self.remote_pg_url)

    def init_hges(self):
        def _init_hge(pg, hge_url, log_file):
            return HGE(
                pg=pg, url=hge_url, port_allocator=self.port_allocator,
                args=self.hge_args, log_file= self.work_dir + '/' + log_file,
                docker_image=self.hge_docker_image
            )

        self.hge = _init_hge(self.pg, self.hge_url, 'hge.log')

        if not self.skip_remote_graphql_setup:
            self.remote_hge = _init_hge(self.remote_pg, self.remote_hge_url, 'remote_hge.log')

    @contextmanager
    def graphql_engines_setup(self):
        try:
            self._setup_graphql_engines()
            yield
        finally:
            self.teardown()

    def _setup_graphql_engines(self):

        if not self.hge_docker_image and not self.skip_stack_build:
            HGE.do_stack_build()

        def run_concurrently(threads):
            for thread in threads:
                thread.start()

            for thread in threads:
                thread.join()

        def run_concurrently_fns(*fns):
            threads = [threading.Thread(target=fn) for fn in fns]
            return run_concurrently(threads)

        def set_hge(hge, schema, hge_type):
            pg = hge.pg
            # Schema and data
            pg.run_sql_from_file(sql_file)
            pg.set_id_as_primary_key_for_tables(schema='public')
            pg.move_tables_to_schema('public', schema)

            # Metadata stuff
            hge.track_all_tables_in_schema(schema)
            hge.create_obj_fk_relationships(schema)
            hge.create_arr_fk_relationships(schema)

        def start_remote_postgres_docker():
            if not self.skip_remote_graphql_setup:
                self.remote_pg.start_postgres_docker()

        run_concurrently_fns(
            self.pg.start_postgres_docker,
            start_remote_postgres_docker)

        print("Postgres url:", self.pg.url)

        if not self.skip_remote_graphql_setup:
            print("Remote Postgres url:", self.remote_pg.url)
            self.remote_hge.run()

        self.hge.run()

        # Skip if the tables are already present
        tables = self.pg.get_all_tables_in_a_schema('hge')
        if len(tables) > 0:
            return

        # Download sportsdb
        zip_file = self.download_sportsdb_zip(self.work_dir+ '/sportsdb.zip')
        sql_file = self.unzip_sql_file(zip_file)

        def set_remote_hge():
            if not self.skip_remote_graphql_setup:
                set_hge(self.remote_hge, 'remote_hge', 'Remote')

        # Create the required tables and move them to required schemas
        hge_thread = threading.Thread(
            target=set_hge, args=(self.hge, 'hge', 'Main'))
        remote_hge_thread = threading.Thread(
            target=set_remote_hge)
        run_concurrently([hge_thread, remote_hge_thread])

        if not self.skip_remote_graphql_setup:
            # Add remote_hge as remote schema
            self.hge.add_remote_schema(
                'remote_hge', self.remote_hge.url + '/v1/graphql',
                self.remote_hge.admin_auth_headers()
            )

        # TODO update the remote schema url if needed
        tables = self.pg.get_all_tables_in_a_schema('hdb_catalog')

	    # Create remote relationships only if it is supported
        if 'hdb_remote_relationship' not in tables:
            return

        # Create remote relationships
        if not self.skip_remote_graphql_setup:
            self.create_remote_relationships()

    def create_remote_relationships(self):
        self.hge.create_remote_obj_rel_to_itself('hge', 'remote_hge', 'remote_hge')
        self.hge.create_remote_arr_fk_ish_relationships('hge', 'remote_hge', 'remote_hge')
        self.hge.create_remote_obj_fk_ish_relationships('hge', 'remote_hge', 'remote_hge')

    def teardown(self):
        for res in [self.hge, self.pg]:
            res.teardown()
        if not self.skip_remote_graphql_setup:
            for res in [self.remote_hge, self.remote_pg]:
                res.teardown()

    def unzip_sql_file(self, zip_file):
        with ZipFile(zip_file, 'r') as zip:
            sql_file = zip.infolist()[0]
            print('DB SQL file:', sql_file.filename)
            zip.extract(sql_file, self.work_dir)
        return self.work_dir + '/' + sql_file.filename

    def download_sportsdb_zip(self, filename, url=sportsdb_url):
        with requests.get(url, stream=True) as r:
            r.raise_for_status()
            total = 0
            print()
            with open(filename, 'wb') as f:
                for chunk in r.iter_content(chunk_size=8192):
                    if chunk:
                        total += len(chunk)
                        print("\rDownloaded: ", int(total/1024)  , 'KB', end='')
                        f.write(chunk)
            print('\nDB Zip File:', filename)
        return filename


class HGETestSetupArgs:

    default_pg_docker_image = 'circleci/postgres:11.5-alpine-postgis'

    def __init__(self):
        self.set_arg_parse_options()
        self.parse_args()

    def set_arg_parse_options(self):
        self.arg_parser = argparse.ArgumentParser()
        self.set_pg_options()
        self.set_hge_options()

    def parse_args(self):
        self.parsed_args = self.arg_parser.parse_args()
        self.set_pg_confs()
        self.set_hge_confs()

    def set_pg_confs(self):
        self.pg_url, self.remote_pg_url, self.pg_docker_image = self.get_params(
            ['pg_url', 'remote_pg_url', 'pg_docker_image']
        )
        if not self.pg_url:
            self.pg_docker_image = self.pg_docker_image or self.default_pg_docker_image

    def set_hge_confs(self):
        self.hge_docker_image, self.hge_url, self.remote_hge_url = self.get_params(
            ['hge_docker_image', 'hge_url', 'remote_hge_url']
        )
        self.skip_stack_build = self.parsed_args.skip_stack_build
        self.skip_remote_graphql_setup = self.parsed_args.skip_remote_graphql_setup
        self.hge_args =  self.parsed_args.hge_args[1:]

    def set_pg_options(self):
        pg_opts = self.arg_parser.add_argument_group('Postgres').add_mutually_exclusive_group()
        pg_opts.add_argument('--pg-url', metavar='HASURA_BENCH_PG_URLS', help='Postgres database url to be used for tests', required=False)
        pg_opts.add_argument('--remote-pg-url', metavar='HASURA_BENCH_REMOTE_PG_URLS', help='Url of Postgres database which is attached/has to be attached, with remote graphql-engine', required=False)
        pg_opts.add_argument('--pg-docker-image', metavar='HASURA_BENCH_PG_DOCKER_IMAGE', help='Postgres docker image to be used for tests', required=False)

    def set_hge_options(self):
        hge_opts = self.arg_parser.add_argument_group('Hasura GraphQL Engine')
        hge_opts.add_argument('--hge-url', metavar='HASURA_BENCH_HGE_URL', help='Url of Hasura graphql-engine')
        hge_opts.add_argument('--remote-hge-url', metavar='HASURA_BENCH_REMOTE_HGE_URL', help='Url of remote Hasura graphql-engine')
        hge_opts.add_argument('--hge-docker-image', metavar='HASURA_BENCH_HGE_DOCKER_IMAGE', help='GraphQl engine docker image to be used for tests', required=False)
        hge_opts.add_argument('--skip-stack-build', help='Skip stack build if this option is set', action='store_true', required=False)
        hge_opts.add_argument('--skip-remote-graphql-setup', help='Skip setting up of remote graphql engine', action='store_true', required=False)
        self.arg_parser.add_argument('hge_args', nargs=argparse.REMAINDER)

    def default_env(self, attr):
        return 'HASURA_BENCH_' + attr.upper()

    def get_param(self, attr, env_key=None):
        if not env_key:
            env_key = self.default_env(attr)
        return _first_true([getattr(self.parsed_args, attr), os.getenv(env_key)])

    def get_params(self, params_loc):
        params_out = []
        for param_loc in params_loc:
            # You can specify just the attribute name for the parameter
            if isinstance(param_loc, str):
                attr = param_loc
                env = self.default_env(attr)
            # Or you can specify attribute and environmental variables as a tuple
            else:
                (attr, env) = param_loc
            param = self.get_param(attr, env)
            params_out.append(param)
        return params_out


class HGETestSetupWithArgs(HGETestSetup, HGETestSetupArgs):

    def __init__(self):
        HGETestSetupArgs.__init__(self)
        HGETestSetup.__init__(
            self,
            pg_url = self.pg_url,
            remote_pg_url = self.remote_pg_url,
            pg_docker_image = self.pg_docker_image,
            hge_url = self.hge_url,
            remote_hge_url = self.remote_hge_url,
            hge_docker_image = self.hge_docker_image,
            skip_stack_build = self.skip_stack_build,
            skip_remote_graphql_setup = self.skip_remote_graphql_setup,
            hge_args = self.hge_args
        )

if __name__ == "__main__":
    test_setup = HGETestSetupWithArgs()
    with test_setup.graphql_engines_setup():
        print("Hasura GraphQL engine is running on URL:",test_setup.hge.url+ '/v1/graphql')
        input(Fore.BLUE+'Press Enter to stop GraphQL engine' + Style.RESET_ALL)