def _require_ssh(self) -> None:
# connect
if self._ssh_client is None:
self._ssh_client = connection.connect_ssh(self._host,
self._ssh_port,
self._ssh_username,
self._ssh_password,
self._ssh_public_key)
# lookup
stdout_list: typing.List[bytes] = []
stderr_list: typing.List[bytes] = []
if connection.run_ssh(self._ssh_client, [
*self._ssh_command_prefix,
'python3',
'-m',
'ck.clickhouse.lookup',
], iteration.empty_in(), iteration.collect_out(stdout_list),
iteration.collect_out(stderr_list))():
raise exception.ShellError(self._host,
b''.join(stderr_list).decode())
(
self._ssh_default_data_dir,
self._ssh_binary_file,
) = b''.join(stdout_list).decode().splitlines()
def test_iteration_adhoc_out() -> None:
gen_out = iteration.empty_out()
next(gen_out)
gen_out.send(b'')
assert not list(gen_out)
gen_out = iteration.empty_out()
next(gen_out)
catched_error = False
try:
gen_out.send(b'1')
except RuntimeError:
catched_error = True
assert catched_error
gen_out = iteration.ignore_out()
next(gen_out)
gen_out.send(b'1')
gen_out.send(b'2')
gen_out.send(b'3')
gen_out.send(b'')
assert not list(gen_out)
stdout_list: typing.List[bytes] = []
gen_out = iteration.collect_out(stdout_list)
next(gen_out)
gen_out.send(b'1')
gen_out.send(b'2')
gen_out.send(b'3')
gen_out.send(b'')
assert stdout_list == [b'1', b'2', b'3']
assert not list(gen_out)
def get_pid(self) -> typing.Optional[int]:
pid_path = self._path.joinpath('pid')
# get pid
stdout_list: typing.List[bytes] = []
if connection.run_ssh(self._ssh_client, [
'cat',
str(pid_path),
], iteration.empty_in(), iteration.collect_out(stdout_list),
iteration.ignore_out())():
return None
pid = int(b''.join(stdout_list).decode().strip())
# find process
if connection.run_ssh(self._ssh_client, [
'kill',
'-0',
str(pid),
], iteration.empty_in(), iteration.empty_out(),
iteration.ignore_out())():
return None
return pid
def stop(self,
ping_interval: float = 0.1,
ping_retry: int = 50) -> typing.Optional[int]:
pid = self.get_pid()
if pid is None:
return None
# kill process
stderr_list: typing.List[bytes] = []
assert self._ssh_client is not None
if connection.run_ssh(self._ssh_client, [
'kill',
'-15',
str(pid),
], iteration.empty_in(), iteration.empty_out(),
iteration.collect_out(stderr_list))():
raise exception.ShellError(self._host, b''.join(stderr_list))
for _ in range(ping_retry):
if self.get_pid() is None:
break
time.sleep(ping_interval)
else:
stderr_list = []
if connection.run_ssh(self._ssh_client, [
'kill',
'-9',
str(pid),
], iteration.empty_in(), iteration.empty_out(),
iteration.collect_out(stderr_list))():
raise exception.ShellError(self._host, b''.join(stderr_list))
while self.get_pid() is not None:
time.sleep(ping_interval)
return pid
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_iteration_collect_out() -> None:
data_list: typing.List[bytes] = []
gen_out = iteration.collect_out(data_list)
next(gen_out)
gen_out.send(b'1')
gen_out.send(b'2')
gen_out.send(b'3')
gen_out.send(b'')
assert data_list == [b'1', b'2', b'3']
assert not list(gen_out)
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
