85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
from typing import TYPE_CHECKING, BinaryIO
|
|
|
|
from clickhouse_connect.driver.httpclient import HttpClient as ClickHouseHttpClient # type: ignore
|
|
import redis
|
|
|
|
if TYPE_CHECKING:
|
|
import tomli as tomllib
|
|
else:
|
|
try:
|
|
import tomllib
|
|
except ImportError:
|
|
import tomli as tomllib
|
|
|
|
|
|
class ControllerConfig:
|
|
__slots__ = ('redis_read', 'redis_write', )
|
|
|
|
redis_read: 'redis.Redis[bytes]' # Set decode_responses to False
|
|
redis_write: 'redis.Redis[bytes]' # Set decode_responses to False
|
|
|
|
def __init__(
|
|
self,
|
|
redis_read: 'redis.Redis[bytes]',
|
|
redis_write: 'redis.Redis[bytes]',
|
|
) -> None:
|
|
self.redis_read = redis_read
|
|
self.redis_write = redis_write
|
|
|
|
|
|
class WorkerConfig:
|
|
__slots__ = ('clickhouse_write', )
|
|
|
|
clickhouse_write: ClickHouseHttpClient
|
|
|
|
def __init__(
|
|
self,
|
|
clickhouse_write: ClickHouseHttpClient,
|
|
) -> None:
|
|
self.clickhouse_write = clickhouse_write
|
|
|
|
|
|
class Config:
|
|
__slots__ = ('controller_config', 'worker_config')
|
|
|
|
controller_config: ControllerConfig
|
|
worker_config: WorkerConfig
|
|
|
|
def __init__(self, controller_config: ControllerConfig, worker_config: WorkerConfig) -> None:
|
|
self.controller_config = controller_config
|
|
self.worker_config = worker_config
|
|
|
|
|
|
def from_toml(toml: BinaryIO) -> Config:
|
|
"""
|
|
Loads the given toml and builds a config instance out of it
|
|
"""
|
|
toml_dict = tomllib.load(toml)
|
|
|
|
redis_read = redis.Redis(
|
|
host=toml_dict['controller']['redis']['read']['host'],
|
|
port=toml_dict['controller']['redis']['read']['port'],
|
|
decode_responses=False,
|
|
)
|
|
|
|
redis_write = redis.Redis(
|
|
host=toml_dict['controller']['redis']['write']['host'],
|
|
port=toml_dict['controller']['redis']['write']['port'],
|
|
decode_responses=False,
|
|
)
|
|
|
|
controller_config = ControllerConfig(redis_read, redis_write)
|
|
|
|
clickhouse_write = ClickHouseHttpClient(
|
|
interface='http',
|
|
host=toml_dict['worker']['clickhouse']['write']['host'],
|
|
port=toml_dict['worker']['clickhouse']['write']['port'],
|
|
username=toml_dict['worker']['clickhouse']['write']['username'],
|
|
password=toml_dict['worker']['clickhouse']['write']['password'],
|
|
database=toml_dict['worker']['clickhouse']['write']['database'],
|
|
)
|
|
|
|
worker_config = WorkerConfig(clickhouse_write)
|
|
|
|
return Config(controller_config, worker_config)
|