cache external postgres port in Branch model

Author: boddumanohar

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

@@ -169,6 +169,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:
@@ -846,6 +860,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)
 
 
@@ -917,6 +932,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)
 
 
@@ -995,6 +1011,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)
 
 
@@ -1136,7 +1153,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, ConfigDict, ValidationError, model_validator
 from pydantic import Field as PydanticField
 from sqlalchemy import BigInteger, Boolean, Column, String, Text, UniqueConstraint, text
@@ -115,6 +116,7 @@ class Branch(AsyncAttrs, Model, table=True):
         sa_column=Column(JSONB, nullable=False, server_default=text("'{}'::jsonb")),
     )
     pitr_enabled: bool = Field(default=False, sa_column=Column(Boolean, nullable=False, server_default=text("false")))
+    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,109 @@
+"""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 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
+
+logger = logging.getLogger(__name__)
+
+# revision identifiers, used by Alembic.
+revision: str = '2b4e8f1a6c03'
+down_revision: Union[str, Sequence[str], None] = 'f4f677e4e9b9'
+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 _deployment_namespace(branch_id: str) -> str:
+    prefix = os.environ.get("VELA_DEPLOYMENT_NAMESPACE_PREFIX", "vela")
+    return f"{prefix}-{branch_id.lower()}" if prefix else branch_id.lower()
+
+
+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_node_port(branch_id: str) -> int | None:
+    core_v1 = kubernetes_client.CoreV1Api()
+    namespace = _deployment_namespace(branch_id)
+    for service_name in (_branch_service_name("pgbouncer"), _branch_service_name("db")):
+        try:
+            svc = core_v1.read_namespaced_service(name=service_name, namespace=namespace)
+            ports = svc.spec.ports
+            if ports and ports[0].node_port:
+                return ports[0].node_port
+        except kubernetes_client.exceptions.ApiException as exc:
+            if exc.status == 404:
+                continue
+            raise
+    return None
+
+
+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()
+
+    updates = []
+    for row in branch_rows:
+        port = _get_node_port(row["id"])
+        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')