Skip to content

Commit 5cf882d

Browse files
Eden-D-ZhangEden-D-Zhang
committed
Refactor db config
1 parent 0f88d3e commit 5cf882d

File tree

16 files changed

+143
-91
lines changed

16 files changed

+143
-91
lines changed

‎components/clp-package-utils/clp_package_utils/controller.py‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,12 @@
1313
from clp_py_utils.clp_config import (
1414
API_SERVER_COMPONENT_NAME,
1515
AwsAuthType,
16+
CLP_DB_PASS_ENV_VAR_NAME,
17+
CLP_DB_ROOT_PASS_ENV_VAR_NAME,
18+
CLP_DB_ROOT_USER_ENV_VAR_NAME,
19+
CLP_DB_USER_ENV_VAR_NAME,
1620
ClpConfig,
21+
ClpDbUserType,
1722
COMPRESSION_JOBS_TABLE_NAME,
1823
COMPRESSION_SCHEDULER_COMPONENT_NAME,
1924
COMPRESSION_WORKER_COMPONENT_NAME,
@@ -138,16 +143,14 @@ def _set_up_env_for_database(self) -> EnvVarsDict:
138143
}
139144

140145
# Credentials
146+
credentials = self._clp_config.database.credentials
141147
env_vars |= {
142-
"CLP_DB_PASS": self._clp_config.database.password,
143-
"CLP_DB_USER": self._clp_config.database.username,
148+
CLP_DB_ROOT_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.ROOT].password,
149+
CLP_DB_ROOT_USER_ENV_VAR_NAME: credentials[ClpDbUserType.ROOT].username,
150+
CLP_DB_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].password,
151+
CLP_DB_USER_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].username,
144152
}
145153

146-
if self._clp_config.database.has_root_password():
147-
env_vars |= {
148-
"CLP_DB_ROOT_PASS": self._clp_config.database.root_password,
149-
}
150-
151154
# Paths
152155
env_vars |= {
153156
"CLP_DB_CONF_LOGGING_FILE_HOST": str(conf_logging_file),

‎components/clp-package-utils/clp_package_utils/general.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -454,6 +454,7 @@ def generate_credentials_file(credentials_file_path: pathlib.Path):
454454
DB_COMPONENT_NAME: {
455455
"username": "clp-user",
456456
"password": secrets.token_urlsafe(8),
457+
"root_username": "root",
457458
"root_password": secrets.token_urlsafe(8),
458459
},
459460
QUEUE_COMPONENT_NAME: {"username": "clp-user", "password": secrets.token_urlsafe(8)},

‎components/clp-package-utils/clp_package_utils/scripts/archive_manager.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
CLP_DB_USER_ENV_VAR_NAME,
1212
CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
1313
CLP_DEFAULT_DATASET_NAME,
14+
ClpDbUserType,
1415
StorageEngine,
1516
StorageType,
1617
)
@@ -226,9 +227,10 @@ def main(argv: list[str]) -> int:
226227
mounts.logs_dir,
227228
mounts.archives_output_dir,
228229
]
230+
credentials = clp_config.database.credentials
229231
extra_env_vars = {
230-
CLP_DB_USER_ENV_VAR_NAME: clp_config.database.username,
231-
CLP_DB_PASS_ENV_VAR_NAME: clp_config.database.password,
232+
CLP_DB_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].password,
233+
CLP_DB_USER_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].username,
232234
}
233235
container_start_cmd: list[str] = generate_container_start_cmd(
234236
container_name, necessary_mounts, clp_config.container_image_ref, extra_env_vars

‎components/clp-package-utils/clp_package_utils/scripts/compress.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
CLP_DB_USER_ENV_VAR_NAME,
1212
CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
1313
CLP_DEFAULT_DATASET_NAME,
14+
ClpDbUserType,
1415
StorageEngine,
1516
StorageType,
1617
)
@@ -252,9 +253,10 @@ def main(argv):
252253
logger.error("No filesystem paths given for compression.")
253254
return -1
254255

256+
credentials = clp_config.database.credentials
255257
extra_env_vars = {
256-
CLP_DB_USER_ENV_VAR_NAME: clp_config.database.username,
257-
CLP_DB_PASS_ENV_VAR_NAME: clp_config.database.password,
258+
CLP_DB_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].password,
259+
CLP_DB_USER_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].username,
258260
}
259261
container_start_cmd = generate_container_start_cmd(
260262
container_name, necessary_mounts, clp_config.container_image_ref, extra_env_vars

‎components/clp-package-utils/clp_package_utils/scripts/compress_from_s3.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
CLP_DB_USER_ENV_VAR_NAME,
1212
CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
1313
CLP_DEFAULT_DATASET_NAME,
14+
ClpDbUserType,
1415
StorageEngine,
1516
StorageType,
1617
)
@@ -306,9 +307,10 @@ def main(argv):
306307
logger.error("No S3 URLs given for compression.")
307308
return -1
308309

310+
credentials = clp_config.database.credentials
309311
extra_env_vars = {
310-
CLP_DB_USER_ENV_VAR_NAME: clp_config.database.username,
311-
CLP_DB_PASS_ENV_VAR_NAME: clp_config.database.password,
312+
CLP_DB_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].password,
313+
CLP_DB_USER_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].username,
312314
}
313315
container_start_cmd = generate_container_start_cmd(
314316
container_name, necessary_mounts, clp_config.container_image_ref, extra_env_vars

‎components/clp-package-utils/clp_package_utils/scripts/dataset_manager.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
CLP_DB_PASS_ENV_VAR_NAME,
1212
CLP_DB_USER_ENV_VAR_NAME,
1313
CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
14+
ClpDbUserType,
1415
StorageEngine,
1516
StorageType,
1617
)
@@ -155,9 +156,10 @@ def main(argv: list[str]) -> int:
155156
if aws_mount:
156157
necessary_mounts.append(mounts.aws_config_dir)
157158

159+
credentials = clp_config.database.credentials
158160
extra_env_vars = {
159-
CLP_DB_USER_ENV_VAR_NAME: clp_config.database.username,
160-
CLP_DB_PASS_ENV_VAR_NAME: clp_config.database.password,
161+
CLP_DB_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].password,
162+
CLP_DB_USER_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].username,
161163
}
162164
container_start_cmd = generate_container_start_cmd(
163165
container_name, necessary_mounts, clp_config.container_image_ref, extra_env_vars

‎components/clp-package-utils/clp_package_utils/scripts/decompress.py‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
1212
CLP_DEFAULT_DATASET_NAME,
1313
ClpConfig,
14+
ClpDbUserType,
1415
StorageEngine,
1516
StorageType,
1617
)
@@ -132,9 +133,10 @@ def handle_extract_file_cmd(
132133
)
133134
)
134135

136+
credentials = clp_config.database.credentials
135137
extra_env_vars = {
136-
CLP_DB_USER_ENV_VAR_NAME: clp_config.database.username,
137-
CLP_DB_PASS_ENV_VAR_NAME: clp_config.database.password,
138+
CLP_DB_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].password,
139+
CLP_DB_USER_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].username,
138140
}
139141
container_start_cmd = generate_container_start_cmd(
140142
container_name, necessary_mounts, clp_config.container_image_ref, extra_env_vars
@@ -214,9 +216,10 @@ def handle_extract_stream_cmd(
214216
container_clp_config, clp_config, get_container_config_filename(container_name)
215217
)
216218
necessary_mounts = [mounts.logs_dir]
219+
credentials = clp_config.database.credentials
217220
extra_env_vars = {
218-
CLP_DB_USER_ENV_VAR_NAME: clp_config.database.username,
219-
CLP_DB_PASS_ENV_VAR_NAME: clp_config.database.password,
221+
CLP_DB_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].password,
222+
CLP_DB_USER_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].username,
220223
}
221224
container_start_cmd = generate_container_start_cmd(
222225
container_name, necessary_mounts, clp_config.container_image_ref, extra_env_vars

‎components/clp-package-utils/clp_package_utils/scripts/native/decompress.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
CLP_DB_USER_ENV_VAR_NAME,
1414
CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
1515
ClpConfig,
16+
ClpDbUserType,
1617
Database,
1718
)
1819
from clp_py_utils.clp_metadata_db_utils import get_files_table_name
@@ -245,10 +246,11 @@ def handle_extract_file_cmd(
245246
"--db-table-prefix", clp_db_connection_params["table_prefix"],
246247
]
247248
# fmt: on
249+
credentials = clp_db_connection_params["credentials"]
248250
extract_env = {
249251
**os.environ,
250-
CLP_DB_USER_ENV_VAR_NAME: clp_db_connection_params["username"],
251-
CLP_DB_PASS_ENV_VAR_NAME: clp_db_connection_params["password"],
252+
CLP_DB_USER_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].username,
253+
CLP_DB_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].password,
252254
}
253255

254256
files_to_extract_list_path = None

‎components/clp-package-utils/clp_package_utils/scripts/search.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
CLP_DB_USER_ENV_VAR_NAME,
1111
CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
1212
CLP_DEFAULT_DATASET_NAME,
13+
ClpDbUserType,
1314
StorageEngine,
1415
StorageType,
1516
)
@@ -134,9 +135,10 @@ def main(argv):
134135
container_clp_config, clp_config, get_container_config_filename(container_name)
135136
)
136137
necessary_mounts = [mounts.logs_dir]
138+
credentials = clp_config.database.credentials
137139
extra_env_vars = {
138-
CLP_DB_USER_ENV_VAR_NAME: clp_config.database.username,
139-
CLP_DB_PASS_ENV_VAR_NAME: clp_config.database.password,
140+
CLP_DB_PASS_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].password,
141+
CLP_DB_USER_ENV_VAR_NAME: credentials[ClpDbUserType.CLP].username,
140142
}
141143
container_start_cmd = generate_container_start_cmd(
142144
container_name, necessary_mounts, clp_config.container_image_ref, extra_env_vars

‎components/clp-py-utils/clp_py_utils/clp_config.py‎

Lines changed: 69 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@
6969
CLP_VERSION_FILE_PATH = pathlib.Path("VERSION")
7070

7171
# Environment variable names
72+
CLP_DB_ROOT_USER_ENV_VAR_NAME = "CLP_DB_ROOT_USER"
7273
CLP_DB_ROOT_PASS_ENV_VAR_NAME = "CLP_DB_ROOT_PASS"
7374
CLP_DB_USER_ENV_VAR_NAME = "CLP_DB_USER"
7475
CLP_DB_PASS_ENV_VAR_NAME = "CLP_DB_PASS"
@@ -162,6 +163,16 @@ def validate_query_engine_package_compatibility(self):
162163
return self
163164

164165

166+
class ClpDbUserType(KebabCaseStrEnum):
167+
CLP = auto()
168+
ROOT = auto()
169+
170+
171+
class DbUserCredentials(BaseModel):
172+
username: NonEmptyStr | None = None
173+
password: NonEmptyStr | None = None
174+
175+
165176
class Database(BaseModel):
166177
DEFAULT_PORT: ClassVar[int] = 3306
167178

@@ -173,17 +184,24 @@ class Database(BaseModel):
173184
auto_commit: bool = False
174185
compress: bool = True
175186

176-
username: NonEmptyStr | None = None
177-
password: NonEmptyStr | None = None
178-
179-
root_password: NonEmptyStr | None = None
187+
credentials: dict[ClpDbUserType, DbUserCredentials] = {
188+
ClpDbUserType.CLP: DbUserCredentials(),
189+
ClpDbUserType.ROOT: DbUserCredentials(),
190+
}
180191

181-
def ensure_credentials_loaded(self):
182-
if self.username is None or self.password is None:
192+
def ensure_credentials_loaded(self, user_type: ClpDbUserType):
193+
if (
194+
self.credentials[user_type].username is None
195+
or self.credentials[user_type].password is None
196+
):
183197
raise ValueError("Credentials not loaded.")
184198

185-
def get_mysql_connection_params(self, disable_localhost_socket_connection: bool = False):
186-
self.ensure_credentials_loaded()
199+
def get_mysql_connection_params(
200+
self,
201+
disable_localhost_socket_connection: bool = False,
202+
user_type: ClpDbUserType = ClpDbUserType.CLP,
203+
):
204+
self.ensure_credentials_loaded(user_type)
187205

188206
host = self.host
189207
if disable_localhost_socket_connection and "localhost" == self.host:
@@ -193,8 +211,8 @@ def get_mysql_connection_params(self, disable_localhost_socket_connection: bool
193211
connection_params = {
194212
"host": host,
195213
"port": self.port,
196-
"user": self.username,
197-
"password": self.password,
214+
"user": self.credentials[user_type].username,
215+
"password": self.credentials[user_type].password,
198216
"database": self.name,
199217
"compress": self.compress,
200218
"autocommit": self.auto_commit,
@@ -203,71 +221,74 @@ def get_mysql_connection_params(self, disable_localhost_socket_connection: bool
203221
connection_params["ssl_cert"] = self.ssl_cert
204222
return connection_params
205223

206-
def get_clp_connection_params_and_type(self, disable_localhost_socket_connection: bool = False):
207-
self.ensure_credentials_loaded()
224+
def get_clp_connection_params_and_type(
225+
self,
226+
disable_localhost_socket_connection: bool = False,
227+
user_type: ClpDbUserType = ClpDbUserType.CLP,
228+
):
229+
self.ensure_credentials_loaded(user_type)
208230

209231
host = self.host
210232
if disable_localhost_socket_connection and "localhost" == self.host:
211233
host = "127.0.0.1"
212234

213-
connection_params_and_type = {
235+
return {
214236
# NOTE: clp-core does not distinguish between mysql and mariadb
215-
"type": DatabaseEngine.MYSQL.value,
216237
"host": host,
217238
"port": self.port,
218-
"username": self.username,
219-
"password": self.password,
220-
"name": self.name,
221-
"table_prefix": CLP_METADATA_TABLE_PREFIX,
239+
"credentials": {user_type: self.credentials[user_type].model_dump()},
240+
"database": self.name,
222241
"compress": self.compress,
223242
"autocommit": self.auto_commit,
243+
"name": self.name,
244+
"type": DatabaseEngine.MYSQL.value,
245+
"table_prefix": CLP_METADATA_TABLE_PREFIX,
224246
}
225-
if self.ssl_cert:
226-
connection_params_and_type["ssl_cert"] = self.ssl_cert
227-
return connection_params_and_type
228247

229248
def dump_to_primitive_dict(self):
230-
d = self.model_dump(
231-
exclude={"root_password", "username", "password"}
232-
)
249+
d = self.model_dump(exclude={"credentials"})
233250
return d
234251

235-
def has_root_password(self) -> bool:
236-
"""
237-
Checks if root password is configured.
238-
239-
:return: True if root password is set.
240-
"""
241-
return self.root_password is not None
242-
243252
def load_credentials_from_file(self, credentials_file_path: pathlib.Path):
244253
config = read_yaml_config_file(credentials_file_path)
245254
if config is None:
246255
raise ValueError(f"Credentials file '{credentials_file_path}' is empty.")
247256
try:
248-
self.username = get_config_value(config, f"{DB_COMPONENT_NAME}.username")
249-
self.password = get_config_value(config, f"{DB_COMPONENT_NAME}.password")
257+
self.credentials[ClpDbUserType.CLP].username = get_config_value(
258+
config, f"{DB_COMPONENT_NAME}.username"
259+
)
260+
self.credentials[ClpDbUserType.CLP].password = get_config_value(
261+
config, f"{DB_COMPONENT_NAME}.password"
262+
)
263+
self.credentials[ClpDbUserType.ROOT].username = get_config_value(
264+
config, f"{DB_COMPONENT_NAME}.root_username"
265+
)
266+
self.credentials[ClpDbUserType.ROOT].password = get_config_value(
267+
config, f"{DB_COMPONENT_NAME}.root_password"
268+
)
250269
except KeyError as ex:
251270
raise ValueError(
252271
f"Credentials file '{credentials_file_path}' does not contain key '{ex}'."
253272
)
254273

255-
try:
256-
self.root_password = get_config_value(config, f"{DB_COMPONENT_NAME}.root_password")
257-
except KeyError:
258-
pass
259-
260-
def load_credentials_from_env(self):
261-
"""
262-
:raise ValueError: if any expected environment variable is not set.
274+
def load_credentials_from_env(self, user_type: ClpDbUserType = ClpDbUserType.CLP):
263275
"""
264-
self.username = _get_env_var(CLP_DB_USER_ENV_VAR_NAME)
265-
self.password = _get_env_var(CLP_DB_PASS_ENV_VAR_NAME)
276+
Loads database credentials from environment variables.
266277
267-
try:
268-
self.root_password = _get_env_var(CLP_DB_ROOT_PASS_ENV_VAR_NAME)
269-
except ValueError:
270-
pass
278+
:param user_type: User type whose credentials are to be loaded.
279+
280+
:raise ValueError: If any expected environment variable is not set.
281+
"""
282+
match user_type:
283+
case ClpDbUserType.CLP:
284+
user_env_var = CLP_DB_USER_ENV_VAR_NAME
285+
pass_env_var = CLP_DB_PASS_ENV_VAR_NAME
286+
case ClpDbUserType.ROOT:
287+
user_env_var = CLP_DB_ROOT_USER_ENV_VAR_NAME
288+
pass_env_var = CLP_DB_ROOT_PASS_ENV_VAR_NAME
289+
290+
self.credentials[user_type].username = _get_env_var(user_env_var)
291+
self.credentials[user_type].password = _get_env_var(pass_env_var)
271292

272293
def transform_for_container(self):
273294
self.host = DB_COMPONENT_NAME

0 commit comments

Comments
 (0)