connstr -> kv_connstr in pg.util

scripts/run_protox_e2e_test.py

@@ -1,5 +1,4 @@
 import os
-import re
 import shutil
 import subprocess
 import sys
@@ -19,7 +18,6 @@
     default_traindata_path,
     default_tuning_steps_dpath,
     default_workload_path,
-    get_latest_run_path_from_workspace_path,
     workload_name_fn,
 )
 

tune/env/env_integtests_dbgym_config.yaml

@@ -1 +1 @@
-dbgym_workspace_path: ../env_integtest_workspace/
+dbgym_workspace_path: ../dbgym_env_integtest_workspace/

tune/env/integtest_pg_conn.py

@@ -1,4 +1,3 @@
-import os
 import subprocess
 import unittest
 from pathlib import Path
@@ -13,8 +12,6 @@
     default_dbdata_parent_dpath,
     default_pgbin_path,
     default_pristine_dbdata_snapshot_path,
-    get_symlinks_path_from_workspace_path,
-    get_tmp_path_from_workspace_path,
 )
 
 ENV_INTEGTESTS_DBGYM_CONFIG_FPATH = Path("tune/env/env_integtests_dbgym_config.yaml")

tune/env/pg_conn.py

@@ -22,15 +22,9 @@
 from plumbum import local
 from psycopg.errors import ProgramLimitExceeded, QueryCanceled
 
-# TODO: remove this import
 from util.log import DBGYM_LOGGER_NAME
-from util.pg import (
-    DBGYM_POSTGRES_DBNAME,
-    DBGYM_POSTGRES_PASS,
-    DBGYM_POSTGRES_USER,
-    SHARED_PRELOAD_LIBRARIES,
-)
-from util.workspace import DBGymConfig, link_result, open_and_save, parent_dpath_of_path
+from util.pg import DBGYM_POSTGRES_DBNAME, SHARED_PRELOAD_LIBRARIES, get_kv_connstr
+from util.workspace import DBGymConfig, open_and_save, parent_dpath_of_path
 
 DEFAULT_CONNECT_TIMEOUT = 300
 
@@ -77,13 +71,13 @@ def __init__(
 
         self._conn: Optional[psycopg.Connection[Any]] = None
 
-    def get_connstr(self) -> str:
-        return f"host=localhost port={self.pgport} user={DBGYM_POSTGRES_USER} password={DBGYM_POSTGRES_PASS} dbname={DBGYM_POSTGRES_DBNAME}"
+    def get_kv_connstr(self) -> str:
+        return get_kv_connstr(self.pgport)
 
     def conn(self) -> psycopg.Connection[Any]:
         if self._conn is None:
             self._conn = psycopg.connect(
-                self.get_connstr(), autocommit=True, prepare_threshold=None
+                self.get_kv_connstr(), autocommit=True, prepare_threshold=None
             )
         return self._conn
 
@@ -329,7 +323,7 @@ def cancel_fn(conn_str: str) -> None:
         conn.execute("SET statement_timeout = 300000")
 
         try:
-            timer = threading.Timer(300.0, cancel_fn, args=(self.get_connstr(),))
+            timer = threading.Timer(300.0, cancel_fn, args=(self.get_kv_connstr(),))
             timer.start()
 
             conn.execute(sql)

tune/protox/env/pg_env.py

@@ -1,9 +1,7 @@
 import copy
-import json
 import logging
 import time
-from pathlib import Path
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Optional
 
 import gymnasium as gym
 import psycopg
@@ -430,7 +428,7 @@ def attempt_checkpoint(conn_str: str) -> None:
                         # We've killed the index operation.
                         or "operational" in stderr
                     )
-                    attempt_checkpoint(self.pg_conn.get_connstr())
+                    attempt_checkpoint(self.pg_conn.get_kv_connstr())
                 return False
 
             assert ret == 0, stderr

util/pg.py

@@ -1,5 +1,9 @@
+"""
+There are multiple parts of the codebase which interact with Postgres. This file contains helpers common to all those parts.
+"""
+
 from pathlib import Path
-from typing import Any, List, NewType, Union
+from typing import Any
 
 import pglast
 import psutil
@@ -55,6 +59,10 @@ def get_connstr(pgport: int = DEFAULT_POSTGRES_PORT, use_psycopg: bool = True) -
     return connstr_prefix + "://" + connstr_suffix
 
 
+def get_kv_connstr(pgport: int = DEFAULT_POSTGRES_PORT) -> str:
+    return f"host=localhost port={pgport} user={DBGYM_POSTGRES_USER} password={DBGYM_POSTGRES_PASS} dbname={DBGYM_POSTGRES_DBNAME}"
+
+
 def create_psycopg_conn(pgport: int = DEFAULT_POSTGRES_PORT) -> psycopg.Connection[Any]:
     connstr = get_connstr(use_psycopg=True, pgport=pgport)
     psycopg_conn = psycopg.connect(connstr, autocommit=True, prepare_threshold=None)