ClickHouse/dbms/tests/integration/helpers/client.py

import errno
import subprocess as sp
from threading import Timer
import tempfile


class Client:
    def __init__(self, host, port=9000, command='/usr/bin/clickhouse-client'):
        self.host = host
        self.port = port
        self.command = [command, '--host', self.host, '--port', str(self.port), '--stacktrace']


    def query(self, sql, stdin=None, timeout=None, settings=None):
        return self.get_query_request(sql, stdin=stdin, timeout=timeout, settings=settings).get_answer()


    def get_query_request(self, sql, stdin=None, timeout=None, settings=None):
        command = self.command[:]

        if stdin is None:
            command += ['--multiquery']
            stdin = sql
        else:
            command += ['--query', sql]

        if settings is not None:
            for setting, value in settings.iteritems():
                command += ['--' + setting, str(value)]

        return CommandRequest(command, stdin, timeout)


class QueryTimeoutExceedException(Exception):
    pass


class QueryRuntimeException(Exception):
    pass


class CommandRequest:
    def __init__(self, command, stdin=None, timeout=None):
        # Write data to tmp file to avoid PIPEs and execution blocking
        stdin_file = tempfile.TemporaryFile()
        stdin_file.write(stdin)
        stdin_file.seek(0)
        self.stdout_file = tempfile.TemporaryFile()
        self.stderr_file = tempfile.TemporaryFile()

        #print " ".join(command)

        self.process = sp.Popen(command, stdin=stdin_file, stdout=self.stdout_file, stderr=self.stderr_file)

        self.timer = None
        self.process_finished_before_timeout = True
        if timeout is not None:
            def kill_process():
                if self.process.poll() is None:
                    self.process_finished_before_timeout = False
                    self.process.kill()

            self.timer = Timer(timeout, kill_process)
            self.timer.start()


    def get_answer(self):
        self.process.wait()
        self.stdout_file.seek(0)
        self.stderr_file.seek(0)

        stdout = self.stdout_file.read()
        stderr = self.stderr_file.read()

        if self.timer is not None and not self.process_finished_before_timeout:
            raise QueryTimeoutExceedException('Client timed out!')

        if self.process.returncode != 0 or stderr:
            raise QueryRuntimeException('Client failed! Return code: {}, stderr: {}'.format(self.process.returncode, stderr))

        return stdout