Skip to content

Connection Utils

ConnWithSetup

Bases: connection

A subclass of psycopg2 connection that sets some session parameters on connect.

Source code in src/autoslo/workload_execution/conn_utils.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class ConnWithSetup(_Conn):
    """
    A subclass of psycopg2 connection that sets some session parameters on connect.
    """

    ONE_HOUR_MS = 60 * 60 * 1000  # One hour in milliseconds

    def __init__(
        self,
        *args,
        autocommit: bool = True,
        search_path: Optional[str] = None,
        statement_timeout: int = ONE_HOUR_MS,
        **kwargs,
    ):
        """
        Initialize the connection with custom session parameters.

        Parameters:
            autocommit: Whether to set autocommit mode.
            search_path: Schema search path to set.
            statement_timeout: Statement timeout in milliseconds.
        """

        super().__init__(*args, **kwargs)
        self.autocommit = autocommit
        with self.cursor() as cur:
            cur.execute("SET enable_result_cache_for_session TO OFF;")
            cur.execute(f"SET statement_timeout TO {statement_timeout};")
            if search_path:
                cur.execute(f"SET search_path TO {search_path};")

__init__(*args, autocommit=True, search_path=None, statement_timeout=ONE_HOUR_MS, **kwargs)

Initialize the connection with custom session parameters.

Parameters:

Name Type Description Default
autocommit bool

Whether to set autocommit mode.

 True
search_path Optional[str]

Schema search path to set.

 None
statement_timeout int

Statement timeout in milliseconds.

 ONE_HOUR_MS
Source code in src/autoslo/workload_execution/conn_utils.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def __init__(
    self,
    *args,
    autocommit: bool = True,
    search_path: Optional[str] = None,
    statement_timeout: int = ONE_HOUR_MS,
    **kwargs,
):
    """
    Initialize the connection with custom session parameters.

    Parameters:
        autocommit: Whether to set autocommit mode.
        search_path: Schema search path to set.
        statement_timeout: Statement timeout in milliseconds.
    """

    super().__init__(*args, **kwargs)
    self.autocommit = autocommit
    with self.cursor() as cur:
        cur.execute("SET enable_result_cache_for_session TO OFF;")
        cur.execute(f"SET statement_timeout TO {statement_timeout};")
        if search_path:
            cur.execute(f"SET search_path TO {search_path};")