moved pg to gymlib

Author: patrickhyw

dbms/postgres/cli.py

@@ -9,10 +9,8 @@
 from typing import Any, Optional
 
 import click
-import pglast
-import psutil
-import psycopg
 import sqlalchemy
+from gymlib.pg import create_sqlalchemy_conn, sql_file_execute
 from gymlib.shell import subprocess_run
 from gymlib.symlinks_paths import (
     get_dbdata_tgz_symlink_path,
@@ -28,7 +26,7 @@
     is_fully_resolved,
     is_ssd,
 )
-from sqlalchemy import create_engine, text
+from sqlalchemy import text
 
 from benchmark.constants import DEFAULT_SCALE_FACTOR
 from benchmark.job.load_info import JobLoadInfo
@@ -365,91 +363,3 @@ def sqlalchemy_conn_execute(
     conn: sqlalchemy.Connection, sql: str
 ) -> sqlalchemy.engine.CursorResult[Any]:
     return conn.execute(text(sql))
-
-
-def sql_file_queries(dbgym_workspace: DBGymWorkspace, filepath: Path) -> list[str]:
-    with dbgym_workspace.open_and_save(filepath) as f:
-        lines: list[str] = []
-        for line in f:
-            if line.startswith("--"):
-                continue
-            if len(line.strip()) == 0:
-                continue
-            lines.append(line)
-        queries_str = "".join(lines)
-        queries: list[str] = pglast.split(queries_str)
-        return queries
-
-
-def sql_file_execute(
-    dbgym_workspace: DBGymWorkspace, conn: sqlalchemy.Connection, filepath: Path
-) -> None:
-    for sql in sql_file_queries(dbgym_workspace, filepath):
-        sqlalchemy_conn_execute(conn, sql)
-
-
-# The reason pgport is an argument is because when doing agnet HPO, we want to run multiple instances of Postgres
-#   at the same time. In this situation, they need to have different ports
-def get_connstr(pgport: int = DEFAULT_POSTGRES_PORT, use_psycopg: bool = True) -> str:
-    connstr_suffix = f"{DBGYM_POSTGRES_USER}:{DBGYM_POSTGRES_PASS}@localhost:{pgport}/{DBGYM_POSTGRES_DBNAME}"
-    # use_psycopg means whether or not we use the psycopg.connect() function
-    # counterintuively, you *don't* need psycopg in the connection string if you *are*
-    #   using the psycopg.connect() function
-    connstr_prefix = "postgresql" if use_psycopg else "postgresql+psycopg"
-    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)
-    return psycopg_conn
-
-
-def create_sqlalchemy_conn(
-    pgport: int = DEFAULT_POSTGRES_PORT,
-) -> sqlalchemy.Connection:
-    connstr = get_connstr(use_psycopg=False, pgport=pgport)
-    engine: sqlalchemy.Engine = create_engine(
-        connstr,
-        execution_options={"isolation_level": "AUTOCOMMIT"},
-    )
-    return engine.connect()
-
-
-def get_is_postgres_running() -> bool:
-    """
-    This is often used in assertions to ensure that Postgres isn't running before we
-    execute some code.
-
-    I intentionally do not have a function that forcefully *stops* all Postgres instances.
-    This is risky because it could accidentally stop instances it wasn't supposed (e.g.
-    Postgres instances run by other users on the same machine).
-
-    Stopping Postgres instances is thus a responsibility of the human to take care of.
-    """
-    return len(get_running_postgres_ports()) > 0
-
-
-def get_running_postgres_ports() -> list[int]:
-    """
-    Returns a list of all ports on which Postgres is currently running.
-
-    There are ways to check with psycopg/sqlalchemy. However, I chose to check using
-    psutil to keep it as simple as possible and orthogonal to how connections work.
-    """
-    running_ports = []
-
-    for conn in psutil.net_connections(kind="inet"):
-        if conn.status == "LISTEN":
-            try:
-                proc = psutil.Process(conn.pid)
-                if proc.name() == "postgres":
-                    running_ports.append(conn.laddr.port)
-            except (psutil.NoSuchProcess, psutil.AccessDenied):
-                continue
-
-    return running_ports

gymlib_package/gymlib/pg.py

@@ -0,0 +1,114 @@
+"""
+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
+
+import pglast
+import psutil
+import psycopg
+import sqlalchemy
+from gymlib.workspace import DBGymWorkspace
+from sqlalchemy import create_engine, text
+
+DBGYM_POSTGRES_USER = "dbgym_user"
+DBGYM_POSTGRES_PASS = "dbgym_pass"
+DBGYM_POSTGRES_DBNAME = "dbgym"
+DEFAULT_POSTGRES_DBNAME = "postgres"
+DEFAULT_POSTGRES_PORT = 5432
+SHARED_PRELOAD_LIBRARIES = "boot,pg_hint_plan,pg_prewarm"
+
+
+def sqlalchemy_conn_execute(
+    conn: sqlalchemy.Connection, sql: str
+) -> sqlalchemy.engine.CursorResult[Any]:
+    return conn.execute(text(sql))
+
+
+def sql_file_queries(dbgym_workspace: DBGymWorkspace, filepath: Path) -> list[str]:
+    with dbgym_workspace.open_and_save(filepath) as f:
+        lines: list[str] = []
+        for line in f:
+            if line.startswith("--"):
+                continue
+            if len(line.strip()) == 0:
+                continue
+            lines.append(line)
+        queries_str = "".join(lines)
+        queries: list[str] = pglast.split(queries_str)
+        return queries
+
+
+def sql_file_execute(
+    dbgym_workspace: DBGymWorkspace, conn: sqlalchemy.Connection, filepath: Path
+) -> None:
+    for sql in sql_file_queries(dbgym_workspace, filepath):
+        sqlalchemy_conn_execute(conn, sql)
+
+
+# The reason pgport is an argument is because when doing agnet HPO, we want to run multiple instances of Postgres
+#   at the same time. In this situation, they need to have different ports
+def get_connstr(pgport: int = DEFAULT_POSTGRES_PORT, use_psycopg: bool = True) -> str:
+    connstr_suffix = f"{DBGYM_POSTGRES_USER}:{DBGYM_POSTGRES_PASS}@localhost:{pgport}/{DBGYM_POSTGRES_DBNAME}"
+    # use_psycopg means whether or not we use the psycopg.connect() function
+    # counterintuively, you *don't* need psycopg in the connection string if you *are*
+    #   using the psycopg.connect() function
+    connstr_prefix = "postgresql" if use_psycopg else "postgresql+psycopg"
+    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)
+    return psycopg_conn
+
+
+def create_sqlalchemy_conn(
+    pgport: int = DEFAULT_POSTGRES_PORT,
+) -> sqlalchemy.Connection:
+    connstr = get_connstr(use_psycopg=False, pgport=pgport)
+    engine: sqlalchemy.Engine = create_engine(
+        connstr,
+        execution_options={"isolation_level": "AUTOCOMMIT"},
+    )
+    return engine.connect()
+
+
+def get_is_postgres_running() -> bool:
+    """
+    This is often used in assertions to ensure that Postgres isn't running before we
+    execute some code.
+
+    I intentionally do not have a function that forcefully *stops* all Postgres instances.
+    This is risky because it could accidentally stop instances it wasn't supposed (e.g.
+    Postgres instances run by other users on the same machine).
+
+    Stopping Postgres instances is thus a responsibility of the human to take care of.
+    """
+    return len(get_running_postgres_ports()) > 0
+
+
+def get_running_postgres_ports() -> list[int]:
+    """
+    Returns a list of all ports on which Postgres is currently running.
+
+    There are ways to check with psycopg/sqlalchemy. However, I chose to check using
+    psutil to keep it as simple as possible and orthogonal to how connections work.
+    """
+    running_ports = []
+
+    for conn in psutil.net_connections(kind="inet"):
+        if conn.status == "LISTEN":
+            try:
+                proc = psutil.Process(conn.pid)
+                if proc.name() == "postgres":
+                    running_ports.append(conn.laddr.port)
+            except (psutil.NoSuchProcess, psutil.AccessDenied):
+                continue
+
+    return running_ports

gymlib_package/gymlib/pg_conn.py

@@ -19,16 +19,11 @@
 import psutil
 import psycopg
 import yaml
+from gymlib.pg import DBGYM_POSTGRES_DBNAME, SHARED_PRELOAD_LIBRARIES, get_kv_connstr
 from gymlib.workspace import DBGymWorkspace, parent_path_of_path
 from plumbum import local
 from psycopg.errors import ProgramLimitExceeded, QueryCanceled
 
-from dbms.postgres.cli import (
-    DBGYM_POSTGRES_DBNAME,
-    SHARED_PRELOAD_LIBRARIES,
-    get_kv_connstr,
-)
-
 CONNECT_TIMEOUT = 300
 
 

gymlib_package/gymlib/tests/integtest_pg_conn.py

@@ -2,15 +2,14 @@
 import unittest
 
 import psycopg
-from gymlib.pg_conn import PostgresConn
-from gymlib.tests.gymlib_integtest_util import GymlibIntegtestManager
-from gymlib.workspace import DBGymWorkspace
-
-from dbms.postgres.cli import (
+from gymlib.pg import (
     DEFAULT_POSTGRES_PORT,
     get_is_postgres_running,
     get_running_postgres_ports,
 )
+from gymlib.pg_conn import PostgresConn
+from gymlib.tests.gymlib_integtest_util import GymlibIntegtestManager
+from gymlib.workspace import DBGymWorkspace
 
 
 class PostgresConnTests(unittest.TestCase):

orchestrate/replay.py

@@ -1,13 +1,12 @@
 from collections import defaultdict
 from pathlib import Path
 
+from gymlib.pg import DEFAULT_POSTGRES_PORT
 from gymlib.pg_conn import PostgresConn
 from gymlib.tuning_artifacts import TuningArtifactsReader
 from gymlib.workload import Workload
 from gymlib.workspace import DBGymWorkspace
 
-from dbms.postgres.cli import DEFAULT_POSTGRES_PORT
-
 
 def replay(
     dbgym_workspace: DBGymWorkspace, tuning_artifacts_path: Path