fmt and mypy

cmu-db · Nov 4, 2024 · 013e059 · 013e059
1 parent 548f11e
commit 013e059
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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: