cache external postgres port in Branch model

Author: boddumanohar

src/api/organization/project/branch/__init__.py

@@ -158,6 +158,20 @@ async def _persist_branch_status(branch_id: Identifier, status: BranchServiceSta
         await session.commit()
 
 
+async def _persist_branch_db_port(branch_id: Identifier) -> None:
+    port = await _resolve_branch_db_port(branch_id)
+    if port is None:
+        logger.warning("Could not resolve db port for branch %s", branch_id)
+        return
+    async with AsyncSessionLocal() as session:
+        branch = await session.get(Branch, branch_id)
+        if branch is None:
+            logger.warning("Branch %s missing while persisting db_port", branch_id)
+            return
+        branch.db_port = port
+        await session.commit()
+
+
 async def _cleanup_failed_branch_deployment(branch_id: Identifier) -> None:
     """Best-effort cleanup when provisioning fails so allocations aren't stranded."""
     try:
@@ -715,6 +729,7 @@ async def _deploy_branch_environment_task(
             branch_slug,
         )
         return
+    await _persist_branch_db_port(branch_id)
     await _persist_branch_status(branch_id, BranchServiceStatus.STARTING)
 
 
@@ -788,6 +803,7 @@ async def _clone_branch_environment_task(
             branch_slug,
         )
         return
+    await _persist_branch_db_port(branch_id)
     await _persist_branch_status(branch_id, BranchServiceStatus.STARTING)
 
 
@@ -868,6 +884,7 @@ async def _restore_branch_environment_task(
             branch_slug,
         )
         return
+    await _persist_branch_db_port(branch_id)
     await _persist_branch_status(branch_id, BranchServiceStatus.STARTING)
 
 
@@ -1009,7 +1026,7 @@ async def _public(branch: Branch) -> BranchPublic:
     project = await branch.awaitable_attrs.project
 
     db_host = _resolve_db_host(branch) or ""
-    port = (await _resolve_branch_db_port(branch.id)) or 0
+    port = branch.db_port or (await _resolve_branch_db_port(branch.id)) or 0
 
     # pg-meta and pg are in the same network. So password is not required in connection string.
     connection_string = _build_connection_string("vela", "postgres", 5432)

src/models/branch.py

@@ -3,6 +3,7 @@
 from enum import StrEnum
 from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, Optional
 
+import sqlalchemy as sa
 from pydantic import BaseModel, model_validator
 from pydantic import Field as PydanticField
 from sqlalchemy import BigInteger, Boolean, Column, String, Text, UniqueConstraint, text
@@ -108,6 +109,7 @@ class Branch(AsyncAttrs, Model, table=True):
     )
     pitr_enabled: bool = Field(default=False, sa_column=Column(Boolean, nullable=False, server_default=text("false")))
     resize_task_id: uuid.UUID | None = Field(default=None, nullable=True)
+    db_port: int | None = Field(default=None, sa_column=Column(sa.Integer, nullable=True))
 
     __table_args__ = (UniqueConstraint("project_id", "name", name="unique_branch_name_per_project"),)
 

src/models/migrations/versions/2b4e8f1a6c03_branch_db_port.py

@@ -0,0 +1,128 @@
+"""Add and backfill branch.db_port
+
+Revision ID: 2b4e8f1a6c03
+Revises: f4f677e4e9b9
+Create Date: 2026-03-19 00:00:00.000000
+
+"""
+import logging
+import os
+from typing import Sequence, Union
+from uuid import UUID
+
+from alembic import op
+import sqlalchemy as sa
+from kubernetes import client as kubernetes_client
+from kubernetes import config as kubernetes_config
+from kubernetes.config.config_exception import ConfigException
+from ulid import ULID
+
+logger = logging.getLogger(__name__)
+
+# revision identifiers, used by Alembic.
+revision: str = '2b4e8f1a6c03'
+down_revision: Union[str, Sequence[str], None] = 'a1b2c3d4e5f6'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+_CHART_NAME = "vela"
+
+
+def _ensure_kube_client_config() -> None:
+    try:
+        kubernetes_config.load_incluster_config()
+    except ConfigException:
+        try:
+            kubernetes_config.load_kube_config()
+        except ConfigException as exc:
+            raise RuntimeError("Kubernetes client not configured. Mount kubeconfig or run in-cluster.") from exc
+
+
+def _uuid_to_ulid_str(uuid_str: str) -> str:
+    """Convert a UUID string (as stored in the DB) back to a ULID string."""
+    return str(ULID.from_bytes(UUID(uuid_str).bytes))
+
+
+def _deployment_namespace(branch_id: str) -> str:
+    prefix = os.environ.get("VELA_DEPLOYMENT_NAMESPACE_PREFIX", "vela")
+    ulid_str = _uuid_to_ulid_str(branch_id).lower()
+    return f"{prefix}-{ulid_str}" if prefix else ulid_str
+
+
+def _release_fullname() -> str:
+    release = os.environ.get("VELA_DEPLOYMENT_RELEASE_NAME", "vela")
+    return release if _CHART_NAME in release else f"{release}-{_CHART_NAME}"
+
+
+def _branch_service_name(component: str) -> str:
+    return f"{_release_fullname()}-{component}"
+
+
+def _get_all_node_ports() -> dict[str, int]:
+    """Fetch node ports for all pgbouncer/db services in a single API call.
+
+    Returns a mapping of namespace -> nodePort, preferring the pgbouncer service
+    over the db service when both exist in the same namespace.
+    """
+    core_v1 = kubernetes_client.CoreV1Api()
+    target_names = {_branch_service_name("pgbouncer"), _branch_service_name("db")}
+
+    svc_list = core_v1.list_service_for_all_namespaces()
+
+    ports_by_namespace: dict[str, int] = {}
+    for svc in svc_list.items:
+        if svc.metadata.name not in target_names:
+            continue
+        ns = svc.metadata.namespace
+        if ns in ports_by_namespace:
+            continue  # already resolved for this namespace
+        ports = svc.spec.ports
+        if ports and ports[0].node_port:
+            ports_by_namespace[ns] = ports[0].node_port
+
+    return ports_by_namespace
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    op.add_column(
+        'branch',
+        sa.Column('db_port', sa.Integer(), nullable=True),
+    )
+
+    bind = op.get_bind()
+    branch_rows = bind.execute(
+        sa.text("SELECT id::text AS id FROM branch")
+    ).mappings().all()
+
+    if not branch_rows:
+        return
+
+    _ensure_kube_client_config()
+    ports_by_namespace = _get_all_node_ports()
+
+    updates = []
+    for row in branch_rows:
+        namespace = _deployment_namespace(row["id"])
+        port = ports_by_namespace.get(namespace)
+        if port is None:
+            logger.warning("Could not resolve db_port for branch %s — skipping", row["id"])
+            continue
+        updates.append({"id": row["id"], "db_port": port})
+
+    if updates:
+        bind.execute(
+            sa.text(
+                """
+                UPDATE branch
+                SET db_port = :db_port
+                WHERE id = CAST(:id AS uuid)
+                """
+            ),
+            updates,
+        )
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    op.drop_column('branch', 'db_port')