def test_iteration_concat_in() -> None: gen_in = iteration.concat_in( iteration.given_in([b'1', b'2']), iteration.given_in([b'3']) ) assert list(gen_in) == [b'1', b'2', b'3']
def test_iteration_adhoc_in() -> None: gen_in = iteration.empty_in() assert list(gen_in) == [] gen_in = iteration.given_in([b'1', b'2', b'3']) assert list(gen_in) == [b'1', b'2', b'3'] gen_in = iteration.concat_in( iteration.given_in([b'1', b'2']), iteration.given_in([b'3']) ) assert list(gen_in) == [b'1', b'2', b'3']
def run() -> None: connection.run_process( [clickhouse.binary_file(), 'client'], iteration.given_in([b'select number from numbers(1000000)']), iteration.ignore_out(), iteration.empty_out() )()
def run() -> None: connection.run_http( 'localhost', 8123, '/', iteration.given_in([b'select number from numbers(1000000)']), iteration.ignore_out(), iteration.empty_out() )()
def test_connection_http() -> None: ck.LocalSession(stop=True, start=True) stdout_list: typing.List[bytes] = [] status = connection.run_http('localhost', 8123, '/', {}, iteration.given_in([b'select 1']), iteration.collect_out(stdout_list), iteration.empty_out())() assert stdout_list == [b'1\n'] assert status == 200
def test_connection_process() -> None: ck.LocalSession(stop=True, start=True) stdout_list: typing.List[bytes] = [] status = connection.run_process([clickhouse.binary_file(), 'client'], iteration.given_in([b'select 1']), iteration.collect_out(stdout_list), iteration.empty_out())() assert stdout_list == [b'1\n'] assert status == 0
def test_connection_ssh() -> None: ck.LocalSession() ssh_client = connection.connect_ssh('localhost') stdout_list: typing.List[bytes] = [] status = connection.run_ssh( ssh_client, [clickhouse.binary_file(), 'client'], iteration.given_in([b'select 1']), iteration.collect_out(stdout_list), iteration.empty_out() )() assert stdout_list == [b'1\n'] assert status == 0
def query_async( self, query: str, data: bytes = b'', method: typing.Optional[typing_extensions.Literal['tcp', 'http', 'ssh']] = None, settings: typing.Optional[typing.Dict[str, str]] = None ) -> typing.Callable[[], bytes]: stdout_list: typing.List[bytes] = [] gen_in = iteration.given_in([data]) gen_out = iteration.collect_out(stdout_list) raw_join = self._run(query, gen_in, gen_out, method, settings) def join() -> bytes: raw_join() return b''.join(stdout_list) return join
def query_async( self, query_text: str, method: typing_extensions.Literal['tcp', 'http', 'ssh'] = 'http', gen_in: typing.Optional[typing.Generator[bytes, None, None]] = None, gen_out: typing.Optional[typing.Generator[None, bytes, None]] = None, settings: typing.Optional[typing.Dict[str, str]] = None ) -> typing.Callable[[], typing.Optional[bytes]]: # create connection(s) stdout_list: typing.List[bytes] = [] stderr_list: typing.List[bytes] = [] if gen_in is None: gen_stdin = iteration.given_in([f'{query_text}\n'.encode()]) else: gen_stdin = iteration.concat_in( iteration.given_in([f'{query_text}\n'.encode()]), gen_in) if gen_out is None: gen_stdout = iteration.collect_out(stdout_list) else: gen_stdout = gen_out gen_stderr = iteration.collect_out(stderr_list) if settings is None: full_settings: typing.Dict[str, str] = {} else: full_settings = settings if method == 'tcp': join_raw = connection.run_process([ clickhouse.binary_file(), 'client', f'--host={self._host}', f'--port={self._tcp_port}', *(f'--{key}={value}' for key, value in full_settings.items()), ], gen_stdin, gen_stdout, gen_stderr) good_status = 0 elif method == 'http': join_raw = connection.run_http( self._host, self._http_port, f'/?{urllib.parse.urlencode(full_settings)}', gen_stdin, gen_stdout, gen_stderr) good_status = 200 elif method == 'ssh': self._require_ssh() assert self._ssh_binary_file is not None join_raw = connection.run_ssh(self._ssh_client, [ *self._ssh_command_prefix, self._ssh_binary_file, 'client', f'--port={self._tcp_port}', *(f'--{key}={value}' for key, value in full_settings.items()), ], gen_stdin, gen_stdout, gen_stderr) good_status = 0 # join connection(s) def join() -> typing.Optional[bytes]: if join_raw() != good_status: raise exception.QueryError(self._host, query_text, b''.join(stderr_list).decode()) if gen_out is None: return b''.join(stdout_list) return None return join
def start(self, ping_interval: float = 0.1, ping_retry: int = 50) -> typing.Optional[int]: pid = self.get_pid() if pid is not None: return None config_path = self._path.joinpath('config.xml') pid_path = self._path.joinpath('pid') # create dir stderr_list: typing.List[bytes] = [] if connection.run_ssh(self._ssh_client, [ 'mkdir', '--parents', str(self._path), ], iteration.empty_in(), iteration.empty_out(), iteration.collect_out(stderr_list))(): raise exception.ShellError(self._host, b''.join(stderr_list)) # setup stderr_list = [] if connection.run_ssh( self._ssh_client, [ *self._ssh_command_prefix, 'python3', '-m', 'ck.clickhouse.setup', ], iteration.given_in([ repr({ 'tcp_port': self._tcp_port, 'http_port': self._http_port, 'user': self._user, 'password': self._password, 'data_dir': str(self._path), 'config': self._config, }).encode() ]), iteration.empty_out(), iteration.collect_out(stderr_list))(): raise exception.ShellError(self._host, b''.join(stderr_list)) # run assert self._ssh_binary_file is not None if connection.run_ssh(self._ssh_client, [ *self._ssh_command_prefix, self._ssh_binary_file, 'server', '--daemon', f'--config-file={config_path}', f'--pid-file={pid_path}', ], iteration.empty_in(), iteration.empty_out(), iteration.empty_out())(): raise exception.ServiceError(self._host, 'daemon') # wait for server initialization for _ in range(ping_retry): pid = self.get_pid() if pid is not None: break time.sleep(ping_interval) else: raise exception.ServiceError(self._host, 'pid') while not self.ping(): time.sleep(ping_interval) if self.get_pid() is None: raise exception.ServiceError(self._host, f'pid_{pid}') return pid
def _run( self, query: str, gen_in: typing.Generator[bytes, None, None], gen_out: typing.Generator[None, bytes, None], method: typing.Optional[typing_extensions.Literal['tcp', 'http', 'ssh']], settings: typing.Optional[typing.Dict[str, str]] ) -> typing.Callable[[], None]: self._prepare() # create connection(s) stderr_list: typing.List[bytes] = [] gen_stdin = iteration.concat_in( iteration.given_in([f'{query}\n'.encode()]), gen_in) gen_stdout = gen_out gen_stderr = iteration.collect_out(stderr_list) real_method = method or self._method real_settings = { **self._settings, **(settings or {}), } if real_method == 'tcp': raw_join = connection.run_process([ clickhouse.binary_file(), 'client', f'--host={self._host}', f'--port={self._tcp_port}', f'--user={self._user}', *((f'--password={self._password}', ) if self._password else ()), *(f'--{key}={value}' for key, value in real_settings.items()), ], gen_stdin, gen_stdout, gen_stderr) good_status = 0 elif real_method == 'http': raw_join = connection.run_http( self._host, self._http_port, f'/?{urllib.parse.urlencode(real_settings)}', { 'X-ClickHouse-User': self._user, 'X-ClickHouse-Key': self._password, }, gen_stdin, gen_stdout, gen_stderr) good_status = 200 elif real_method == 'ssh': self._require_ssh() assert self._ssh_binary_file is not None raw_join = connection.run_ssh(self._ssh_client, [ *self._ssh_command_prefix, self._ssh_binary_file, 'client', f'--port={self._tcp_port}', f'--user={self._user}', *((f'--password={self._password}', ) if self._password else ()), *(f'--{key}={value}' for key, value in real_settings.items()), ], gen_stdin, gen_stdout, gen_stderr) good_status = 0 # join connection(s) def join() -> None: if raw_join() != good_status: raise exception.QueryError(self._host, query, b''.join(stderr_list)) return join
def test_iteration_given_in() -> None: gen_in = iteration.given_in([b'1', b'2', b'3']) assert list(gen_in) == [b'1', b'2', b'3']