async def measure(stats, f, *args, **kws):
async def qmap(q, corof, iterable):
async def consume(q, total=None):
async def map(coro, iterable, total=None):
def run(coro, *args):
def interruptable(iterable, is_active):
def setup_sigint_handling(loop, q, iterable):
    def stop():
def remove_sigint_handler(loop):
def run_many(coro, iterable, concurrency, num_items=None): HTTP_DEFAULT_HDRS = {'Content-Type': 'application/json'}
class CrateJsonEncoder(json.JSONEncoder):
class SqlException(Exception):
class AsyncpgClient:
class HttpClient:
    def default(self, o):
    def __init__(self, message):
def _to_http_uri(s: str) -> str:
def _to_http_hosts(hosts: Union[Iterable[str], str]) -> List[str]:
async def _exec(session, url, data):
def _plain_or_callable(obj):
def _date_or_none(d: str) -> str:
def _to_dsn(hosts):
def _to_boolean(v):
def _verify_ssl_from_first(hosts):
    def __init__(self, hosts, pool_size=25):
    async def _get_pool(self):
    async def execute(self, stmt, args=None):
    async def execute_many(self, stmt, bulk_args):
    async def get_server_version(self):
    async def _close_pool(self):
    def close(self):
    def __enter__(self):
    def __exit__(self, *exs):
def _append_sql(host):
    def __init__(self, hosts, conn_pool_limit=25):
    async def _session(self):
    async def execute(self, stmt, args=None):
    async def execute_many(self, stmt, bulk_args):
    async def get_server_version(self):
    async def _close(self):
    def close(self):
    def __enter__(self):
    def __exit__(self, exc_type, exc_val, exc_tb):
def client(hosts, concurrency=25): def to_insert(table, d):
    """Generate an insert statement using the given table and dictionary.

    Args:
        table (str): table name
        d (dict): dictionary with column names as keys and values as values.
    Returns:
        tuple of statement and arguments

    >>> to_insert('doc.foobar', {'name': 'Marvin'})
    ('insert into doc.foobar ("name") values ($1)', ['Marvin'])
    """

    columns = []
    args = []
    for key, val in d.items():
        columns.append('"{}"'.format(key))
        args.append(val)
    stmt = 'insert into {table} ({columns}) values ({params})'.format(
        table=table,
        columns=', '.join(columns),
        params=', '.join(f'${i + 1}' for i in range(len(columns)))
    )
    return (stmt, args) def load_spec(spec_file):
    ext = os.path.splitext(spec_file)[1]
    loader = spec_loaders[ext]
    return loader(spec_file) class Runner:
    def __init__(self, hosts, concurrency, sample_mode):
        self.concurrency = concurrency
        self.client = client(hosts, concurrency=concurrency)
        self.sampler = get_sampler(sample_mode)

    def warmup(self, stmt, num_warmup):
        statements = itertools.repeat((stmt,), num_warmup)
        aio.run_many(self.client.execute, statements, 0, num_items=num_warmup)

    def run(self, stmt, *, iterations=None, duration=None, args=None, bulk_args=None):
        if bulk_args:
            args = bulk_args
            f = self.client.execute_many
        else:
            f = self.client.execute
        statements = _generate_statements(stmt, args, iterations, duration)
        return run_and_measure(
            f, statements, self.concurrency, iterations, sampler=self.sampler)

    def __enter__(self):
        return self

    def __exit__(self, *ex):
        self.client.close() class Result:

    def __init__(self,
                 version_info,
                 statement,
                 timed_stats,
                 concurrency,
                 meta=None,
                 bulk_size=None,
                 name=None):
        self.version_info = version_info
        self.statement = str(statement)
        self.meta = meta and DotDict(meta) or None
        self.started = timed_stats.started
        self.ended = timed_stats.ended
        self.runtime_stats = DotDict(timed_stats.stats.get())
        self.concurrency = concurrency
        self.bulk_size = bulk_size
        self.name = name

    def as_dict(self):
        return self.__dict__ def run_and_measure(f, statements, concurrency, num_items=None, sampler=None):
    stats = Stats(sampler)
    measure = partial(aio.measure, stats, f)
    started = int(time() * 1000)
    aio.run_many(measure, statements, concurrency, num_items=num_items)
    ended = int(time() * 1000)
    return TimedStats(started, ended, stats) def eval_fail_if(fail_if: str, result):
    fail_if = fail_if.format(**result.as_dict())
    if eval(fail_if):
        raise FailIf("Expression failed: " + fail_if) def as_bulk_queries(queries, bulk_size):
    """Group a iterable of (stmt, args) by stmt into (stmt, bulk_args).

    bulk_args will be a list of the args grouped by stmt.

    len(bulk_args) will be <= bulk_size
    """
    stmt_dict = defaultdict(list)
    for stmt, args in queries:
        bulk_args = stmt_dict[stmt]
        bulk_args.append(args)
        if len(bulk_args) == bulk_size:
            yield stmt, bulk_args
            del stmt_dict[stmt]
    for stmt, bulk_args in stmt_dict.items():
        yield stmt, bulk_args def as_statements(lines: Iterator[str]) -> Iterator[str]:
    """Create an iterator that transforms lines into sql statements.

    Statements within the lines must end with ";"
    The last statement will be included even if it does not end in ';'

    >>> list(as_statements(['select * from', '-- comments are filtered', 't;']))
    ['select * from t']

    >>> list(as_statements(['a;', 'b', 'c;', 'd', ' ']))
    ['a', 'b c', 'd']
    """
    lines = (l.strip() for l in lines if l)
    lines = (l for l in lines if l and not l.startswith('--'))
    parts = []
    for line in lines:
        parts.append(line.rstrip(';'))
        if line.endswith(';'):
            yield ' '.join(parts)
            parts.clear()
    if parts:
        yield ' '.join(parts) def get_lines(filename: str) -> Iterator[str]:
    """Create an iterator that returns the lines of a utf-8 encoded file."""
    if filename.endswith('.gz'):
        with gzip.open(filename, 'r') as f:
            for line in f:
                yield line.decode('utf-8')
    else:
        with open(filename, 'r', encoding='utf-8') as f:
            for line in f:
                yield line def parse_version(version: str) -> tuple:
    """Parse a string formatted X[.Y.Z] version number into a tuple

    >>> parse_version('10.2.3')
    (10, 2, 3)

    >>> parse_version('12')
    (12, 0, 0)
    """
    parts = version.split('.')
    missing = 3 - len(parts)
    return tuple(map(int, parts + (['0'] * missing))) def try_len(o: Any) -> Optional[int]:
    """ Return len of `o` or None if `o` doesn't support len

    >>> try_len([1, 2])
    2

    >>> try_len(print)
    >>> try_len(None)
    """
    if not o:
        return None
    try:
        return len(o)
    except TypeError:
        return None def dicts_from_lines(lines):
    """ returns a generator producing dicts from json lines

    1 JSON object per line is supported:

        {"name": "n1"}
        {"name": "n2"}

    Or 1 JSON object:

        {
            "name": "n1"
        }

    Or a list of JSON objects:

        [
            {"name": "n1"},
            {"name": "n2"},
        ]

    Or a list of JSON objects in a single line:

        [{"name": "n1"}, {"name": "n2"}]
    """
    lines = iter(lines)
    for line in lines:
        line = line.strip()
        if not line:
            continue  # skip empty lines
        try:
            data = loads(line)
            if isinstance(data, list):
                yield from data
            else:
                yield data
        except ValueError:
            content = line + ''.join(lines)
            dicts = loads(content)
            if isinstance(dicts, list):
                yield from dicts
            else:
                yield dicts class Logger(contextlib.ExitStack):

    def __init__(self,
                 output_fmt='text',
                 logfile_info=None,
                 logfile_result=None):
        super().__init__()
        info_output = self._open(logfile_info)
        result_output = self._open(logfile_result)
        self.info = partial(print, file=info_output)
        presult = partial(print, file=result_output)
        if output_fmt == 'json':
            self.result = lambda r: presult(to_jsonstr(r.as_dict()))
        else:
            self.result = lambda r: presult(_format_short(r.runtime_stats))

    def _open(self, filepath):
        if filepath:
            return self.enter_context(_wopen(filepath))
        return sys.stdout
import argh
import os
import itertools
import subprocess
from functools import partial
from typing import Iterable
from cr8 import aio, clients
from cr8.insert_json import to_insert
from cr8.bench_spec import load_spec
from cr8.engine import Runner, Result, run_and_measure, eval_fail_if
from cr8.misc import (
    as_bulk_queries,
    as_statements,
    get_lines,
    parse_version,
    try_len
)
from cr8.cli import dicts_from_lines
from cr8.log import Logger
                 sample_mode):
        self.benchmark_hosts = benchmark_hosts
        self.sample_mode = sample_mode
        self.spec_dir = spec_dir
        self.client = clients.client(benchmark_hosts)
        self.result_client = clients.client(result_hosts)
        self.server_version_info = aio.run(self.client.get_server_version)
        self.server_version = parse_version(self.server_version_info['number'])
        self.log = log
        self.create_result = partial(
            Result,
            version_info=self.server_version_info
        )
        if fail_if:
            self.fail_if = partial(eval_fail_if, fail_if)
        else:
            self.fail_if = lambda x: None
        if result_hosts:
            self.process_result = _result_to_crate(self.log, self.result_client)
        else:
            self.process_result = log.result

    def _to_inserts(self, data_spec):
        target = data_spec['target']
        source = os.path.join(self.spec_dir, data_spec['source'])
        dicts = dicts_from_lines(get_lines(source))
        return (to_insert(target, d) for d in dicts)

    def exec_instructions(self, instructions):
        filenames = instructions.statement_files
        filenames = (os.path.join(self.spec_dir, i) for i in filenames)
        lines = (line for fn in filenames for line in get_lines(fn))
        statements = itertools.chain(as_statements(lines), instructions.statements)
        for stmt in statements:
            aio.run(self.client.execute, stmt)

        for data_file in instructions.data_files:
            inserts = as_bulk_queries(self._to_inserts(data_file),
                                      data_file.get('bulk_size', 5000))
            concurrency = data_file.get('concurrency', 25)
            aio.run_many(self.client.execute_many, inserts, concurrency=concurrency)
            if self.client.is_cratedb:
                aio.run(self.client.execute, f"refresh table {data_file['target']}")

        for data_cmd in instructions.data_cmds:
            process = subprocess.Popen(
                data_cmd['cmd'],
                stdout=subprocess.PIPE,
                universal_newlines=True
            )
            target = data_cmd['target']
            dicts = dicts_from_lines(process.stdout)
            inserts = as_bulk_queries(
                (to_insert(target, d) for d in dicts),
                data_cmd.get('bulk_size', 5000)
            )
            concurrency = data_cmd.get('concurrency', 25)
            aio.run_many(self.client.execute_many, inserts, concurrency=concurrency)
            if self.client.is_cratedb:
                aio.run(self.client.execute, f"refresh table {target}")

    def update_server_stats(self):
        """Triggers ANALYZE on the server to update statistics."""
        try:
            aio.run(self.client.execute, 'ANALYZE')
        except Exception:
            pass  # swallow; CrateDB 4.1.0+ is required to run ANALYZE

    def run_load_data(self, data_spec, meta=None):
        inserts = self._to_inserts(data_spec)
        statement = next(iter(inserts))[0]
        bulk_size = data_spec.get('bulk_size', 5000)
        inserts = as_bulk_queries(self._to_inserts(data_spec), bulk_size)
        concurrency = data_spec.get('concurrency', 25)
        num_records = data_spec.get('num_records')
        if num_records:
            num_records = max(1, int(num_records / bulk_size))
        timed_stats = run_and_measure(
            self.client.execute_many, inserts, concurrency, num_records)
        self.process_result(self.create_result(
            statement=statement,
            meta=meta,
            timed_stats=timed_stats,
            concurrency=concurrency,
            bulk_size=bulk_size,
        ))

    def _skip_message(self, min_version, stmt):
        msg = ('## Skipping (Version {server_version} instead of {min_version}):\n'
               '   Statement: {statement:.70}')
        msg = msg.format(
            statement=stmt,
            min_version='.'.join((str(x) for x in min_version)),
            server_version='.'.join((str(x) for x in self.server_version)))
        return msg

    def run_queries(self, queries: Iterable[dict], meta=None):
        for query in queries:
            stmt = query['statement']
            iterations = query.get('iterations', 1)
            duration = query.get('duration')
            name = query.get('name')
            concurrency = query.get('concurrency', 1)
            args = query.get('args')
            bulk_args = query.get('bulk_args')
            _min_version = query.get('min_version')
            min_version = _min_version and parse_version(_min_version)
            if min_version and min_version > self.server_version:
                self.log.info(self._skip_message(min_version, stmt))
                continue
            mode_desc = 'Duration' if duration else 'Iterations'
            name_line = name and f'   Name: {name}\n' or ''
            self.log.info(
                (f'\n## Running Query:\n'
                 f'{name_line}'
                 f'   Statement:\n'
                 f'     {stmt}\n'
                 f'   Concurrency: {concurrency}\n'
                 f'   {mode_desc}: {duration or iterations}')
            )

            with Runner(self.benchmark_hosts, concurrency, self.sample_mode) as runner: