From 013e0598046ea7d973889b6da68e92e85d7be49b Mon Sep 17 00:00:00 2001 From: Patrick Wang Date: Sun, 3 Nov 2024 19:43:07 -0500 Subject: [PATCH] fmt and mypy --- benchmark/job/cli.py | 9 ++++++--- benchmark/job/load_info.py | 5 +++-- dbms/load_info_base_class.py | 2 +- dbms/postgres/cli.py | 6 +++++- 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/benchmark/job/cli.py b/benchmark/job/cli.py index 322a5e42..c4cc9ed3 100644 --- a/benchmark/job/cli.py +++ b/benchmark/job/cli.py @@ -5,7 +5,7 @@ from benchmark.job.load_info import JobLoadInfo from util.log import DBGYM_LOGGER_NAME from util.shell import subprocess_run -from util.workspace import DBGymConfig, link_result, default_tables_dname +from util.workspace import DBGymConfig, default_tables_dname, link_result JOB_TABLES_URL = "https://homepages.cwi.nl/~boncz/job/imdb.tgz" @@ -32,7 +32,8 @@ def tpch_workload(dbgym_cfg: DBGymConfig) -> None: def _download_job_data(dbgym_cfg: DBGymConfig) -> None: expected_symlink_dpath = ( - dbgym_cfg.cur_symlinks_data_path(mkdir=True) / f"{default_tables_dname(JobLoadInfo.JOB_SCALE_FACTOR)}.link" + dbgym_cfg.cur_symlinks_data_path(mkdir=True) + / f"{default_tables_dname(JobLoadInfo.JOB_SCALE_FACTOR)}.link" ) if expected_symlink_dpath.exists(): logging.getLogger(DBGYM_LOGGER_NAME).info( @@ -43,7 +44,9 @@ def _download_job_data(dbgym_cfg: DBGymConfig) -> None: logging.getLogger(DBGYM_LOGGER_NAME).info(f"Downloading: {expected_symlink_dpath}") real_data_path = dbgym_cfg.cur_task_runs_data_path(mkdir=True) subprocess_run(f"curl -O {JOB_TABLES_URL}", cwd=real_data_path) - job_data_dpath = dbgym_cfg.cur_task_runs_data_path(default_tables_dname(JobLoadInfo.JOB_SCALE_FACTOR), mkdir=True) + job_data_dpath = dbgym_cfg.cur_task_runs_data_path( + default_tables_dname(JobLoadInfo.JOB_SCALE_FACTOR), mkdir=True + ) subprocess_run("tar -zxvf ../imdb.tgz", cwd=job_data_dpath) subprocess_run(f"rm imdb.tgz", cwd=real_data_path) symlink_dpath = link_result(dbgym_cfg, job_data_dpath) diff --git a/benchmark/job/load_info.py b/benchmark/job/load_info.py index 857f0533..eb335967 100644 --- a/benchmark/job/load_info.py +++ b/benchmark/job/load_info.py @@ -32,7 +32,7 @@ class JobLoadInfo(LoadInfoBaseClass): "name", "person_info", "role_type", - "title" + "title", ] def __init__(self, dbgym_cfg: DBGymConfig): @@ -52,7 +52,8 @@ def __init__(self, dbgym_cfg: DBGymConfig): dbgym_cfg.dbgym_symlinks_path / JobLoadInfo.CODEBASE_DNAME / "data" ) tables_symlink_dpath = ( - data_root_dpath / f"{default_tables_dname(JobLoadInfo.JOB_SCALE_FACTOR)}.link" + data_root_dpath + / f"{default_tables_dname(JobLoadInfo.JOB_SCALE_FACTOR)}.link" ) tables_dpath = tables_symlink_dpath.resolve() assert ( diff --git a/dbms/load_info_base_class.py b/dbms/load_info_base_class.py index f9f750d8..6df647da 100644 --- a/dbms/load_info_base_class.py +++ b/dbms/load_info_base_class.py @@ -14,7 +14,7 @@ def get_schema_fpath(self) -> Path: def get_tables_and_fpaths(self) -> list[tuple[str, Path]]: raise NotImplemented - + # We assume the table file has a "csv-like" format where values are separated by a delimiter. def get_table_file_delimiter(self) -> str: raise NotImplemented diff --git a/dbms/postgres/cli.py b/dbms/postgres/cli.py index 31200111..d28915f4 100644 --- a/dbms/postgres/cli.py +++ b/dbms/postgres/cli.py @@ -254,6 +254,8 @@ def _generic_dbdata_setup(dbgym_cfg: DBGymConfig) -> None: def _load_benchmark_into_dbdata( dbgym_cfg: DBGymConfig, benchmark_name: str, scale_factor: float ) -> None: + load_info: LoadInfoBaseClass + with create_sqlalchemy_conn() as conn: if benchmark_name == "tpch": load_info = TpchLoadInfo(dbgym_cfg, scale_factor) @@ -281,7 +283,9 @@ def _load_into_dbdata( assert conn.connection.dbapi_connection is not None cur = conn.connection.dbapi_connection.cursor() try: - with cur.copy(f"COPY {table} FROM STDIN CSV DELIMITER '{load_info.get_table_file_delimiter()}' ESCAPE '\\'") as copy: + with cur.copy( + f"COPY {table} FROM STDIN CSV DELIMITER '{load_info.get_table_file_delimiter()}' ESCAPE '\\'" + ) as copy: while data := table_csv.read(8192): copy.write(data) finally: