Start extracting tenant knowledge from pgcon

edb/server/bootstrap.py

@@ -127,7 +127,7 @@ def __init__(
         dbname: Optional[str] = None,
         log_listener: Optional[Callable[[str, str], None]] = None,
     ):
-        self._conn: Optional[pgcon.PGConnection] = None
+        self._conn: Optional[pgcon.PGConnectionRaw] = None
         self._cluster = cluster
         self._dbname = dbname
         self._log_listener = log_listener or _pg_log_listener

edb/server/pgcluster.py

@@ -170,7 +170,7 @@ async def connect(self,
                       source_description: str,
                       apply_init_script: bool = False,
                       **kwargs: Unpack[pgconnparams.CreateParamsKwargs]
-    ) -> pgcon.PGConnection:
+    ) -> pgcon.PGConnectionRaw:
         """Connect to this cluster, with optional overriding parameters. If
         overriding parameters are specified, they are applied to a copy of the
         connection parameters before the connection takes place."""

edb/server/pgcon/__init__.py

@@ -27,7 +27,8 @@
 )
 
 from .pgcon import (
-    PGConnection,
+    PGConnectionRaw,
+    AbstractFrontendConnection,
 )
 from .connect import (
     pg_connect,
@@ -37,13 +38,92 @@
     RESET_STATIC_CFG_SCRIPT,
 )
 
+from edb.server import defines as edbdef
+from typing import Protocol, Optional
+
+
+class PGConnection(Protocol):
+    async def sql_execute(
+        self,
+        sql: bytes | tuple[bytes, ...],
+        *,
+        tx_isolation: edbdef.TxIsolationLevel | None = None,
+    ) -> None: ...
+
+    async def sql_fetch(
+        self,
+        sql: bytes,
+        *,
+        args: tuple[bytes, ...] | list[bytes] = (),
+        use_prep_stmt: bool = False,
+        tx_isolation: edbdef.TxIsolationLevel | None = None,
+        state: Optional[bytes] = None,
+    ) -> list[tuple[bytes, ...]]: ...
+
+    async def sql_fetch_val(
+        self,
+        sql: bytes,
+        *,
+        args: tuple[bytes, ...] | list[bytes] = (),
+        use_prep_stmt: bool = False,
+        tx_isolation: edbdef.TxIsolationLevel | None = None,
+        state: Optional[bytes] = None,
+    ) -> bytes: ...
+
+    async def sql_fetch_col(
+        self,
+        sql: bytes,
+        *,
+        args: tuple[bytes, ...] | list[bytes] = (),
+        use_prep_stmt: bool = False,
+        tx_isolation: edbdef.TxIsolationLevel | None = None,
+        state: Optional[bytes] = None,
+    ) -> list[bytes]: ...
+
+
+class PGConnectionEventListener(Protocol):
+    """Protocol for PGConnection event listeners."""
+
+    def on_pg_conn_closed(
+        self,
+        conn: PGConnectionRaw,
+        expected: bool,
+        exc: Optional[Exception],
+    ) -> None:
+        """Called when a connection is closed normally."""
+        pass
+
+    def on_pg_conn_unavailable(self, conn: PGConnectionRaw, msg: str) -> None:
+        """Set the message to display when PostgreSQL is unavailable."""
+        pass
+
+    def on_pg_conn_parameter_updated(
+        self, conn: PGConnectionRaw, name: str, value: str
+    ) -> None:
+        """Called when a parameter status is updated on a system connection."""
+        pass
+
+    def on_pg_conn_metrics(
+        self,
+        conn: PGConnectionRaw,
+        metric: str,
+        value: int,
+        *args,
+    ) -> None:
+        """Called when a metric is updated."""
+        pass
+
+
 __all__ = (
     'pg_connect',
     'PGConnection',
+    'PGConnectionRaw',
+    'PGConnectionEventListener',
     'BackendError',
     'BackendConnectionError',
     'BackendPrivilegeError',
     'BackendCatalogNameError',
+    'AbstractFrontendConnection',
     'SETUP_TEMP_TABLE_SCRIPT',
     'SETUP_CONFIG_CACHE_SCRIPT',
     'SETUP_DML_DUMMY_TABLE_SCRIPT',

edb/server/pgcon/connect.py

@@ -142,7 +142,7 @@ async def pg_connect(
     backend_params: pg_params.BackendRuntimeParams,
     source_description: str,
     apply_init_script: bool = True,
-) -> pgcon.PGConnection:
+) -> pgcon.PGConnectionRaw:
     global INIT_CON_SCRIPT
 
     if isinstance(dsn_or_connection, str):
@@ -155,7 +155,7 @@ async def pg_connect(
     # disabled.
     pgrawcon, pgconn = await rust_transport.create_postgres_connection(
         connection,
-        lambda: pgcon.PGConnection(dbname=connection.database),
+        lambda: pgcon.PGConnectionRaw(dbname=connection.database),
         source_description=source_description,
     )
 

edb/server/pgcon/pgcon.pxd

@@ -58,8 +58,14 @@ cdef enum PGAuthenticationState:
     PGAUTH_SASL_FINAL = 12
 
 
+cdef class AbstractFrontendConnection:
+
+    cdef write(self, WriteBuffer buf)
+    cdef flush(self)
+
+
 @cython.final
-cdef class PGConnection:
+cdef class PGConnectionRaw:
 
     cdef:
         ReadBuffer buffer
@@ -91,8 +97,7 @@ cdef class PGConnection:
 
         public object connection
         public object addr
-        object server
-        object tenant
+        public object listener
         bint is_system_db
         bint close_requested
 
@@ -103,10 +108,10 @@ cdef class PGConnection:
         bint _is_ssl
 
         public object pinned_by
+        public object data
 
-        object last_state
+        public object last_state
         bint state_reset_needs_commit
-        public object last_init_con_data
 
         str last_indirect_return
 
@@ -131,7 +136,7 @@ cdef class PGConnection:
     cdef send_sync(self)
 
     cdef make_clean_stmt_message(self, bytes stmt_name)
-    cdef send_query_unit_group(
+    cpdef send_query_unit_group(
         self, object query_unit_group, bint sync,
         object bind_datas, bytes state,
         ssize_t start, ssize_t end, int dbver, object parse_array,
@@ -158,7 +163,6 @@ cdef class PGConnection:
         dict type_id_map,
     )
 
-    cdef inline str get_tenant_label(self)
     cpdef set_stmt_cache_size(self, int maxsize)
 
 cdef setting_to_sql(name, setting)

edb/server/pgcon/pgcon.pyi

@@ -29,6 +29,8 @@ import asyncio
 from edb.server import defines as edbdef
 from edb.server import pgconnparams
 
+from . import PGConnectionEventListener
+
 class BackendError(Exception):
     def get_field(self, field: str) -> str | None: ...
     def code_is(self, code: str) -> bool: ...
@@ -37,16 +39,19 @@ class BackendConnectionError(BackendError): ...
 class BackendPrivilegeError(BackendError): ...
 class BackendCatalogNameError(BackendError): ...
 
-class PGConnection(asyncio.Protocol):
+class AbstractFrontendConnection:
+    pass
 
+class PGConnectionRaw(asyncio.Protocol):
     idle: bool
     backend_pid: int
     connection: pgconnparams.ConnectionParams
     addr: tuple[str, int]
     parameter_status: dict[str, str]
     backend_secret: int
     is_ssl: bool
-    last_init_con_data: object
+    listener: PGConnectionEventListener | None
+    data: Any
 
     def __init__(self, dbname): ...
     async def close(self): ...
@@ -96,9 +101,6 @@ class PGConnection(asyncio.Protocol):
     async def signal_sysevent(self, event: str, *, dbname: str) -> None: ...
     def abort(self) -> None: ...
     def is_healthy(self) -> bool: ...
-    async def listen_for_sysevent(self) -> None: ...
-    def mark_as_system_db(self) -> None: ...
-    def set_tenant(self, tenant: Any) -> None: ...
     def is_cancelling(self) -> bool: ...
     def start_pg_cancellation(self) -> None: ...
     def finish_pg_cancellation(self) -> None: ...