store the port in DB as part of _persist_branch_db_port instead of migration

Author: boddumanohar

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

@@ -158,18 +158,19 @@ 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)
+async def _get_branch_db_port(branch: Branch) -> int:
+    if branch.db_port is not None:
+        return branch.db_port
+    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
+        raise VelaKubernetesError(f"Could not resolve db port for branch {branch.id}")
     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()
+        db_branch = await session.get(Branch, branch.id)
+        if db_branch is not None:
+            db_branch.db_port = port
+            await session.commit()
+    branch.db_port = port
+    return port
 
 
 async def _cleanup_failed_branch_deployment(branch_id: Identifier) -> None:
@@ -729,7 +730,6 @@ async def _deploy_branch_environment_task(
             branch_slug,
         )
         return
-    await _persist_branch_db_port(branch_id)
     await _persist_branch_status(branch_id, BranchServiceStatus.STARTING)
 
 
@@ -803,7 +803,6 @@ async def _clone_branch_environment_task(
             branch_slug,
         )
         return
-    await _persist_branch_db_port(branch_id)
     await _persist_branch_status(branch_id, BranchServiceStatus.STARTING)
 
 
@@ -884,7 +883,6 @@ async def _restore_branch_environment_task(
             branch_slug,
         )
         return
-    await _persist_branch_db_port(branch_id)
     await _persist_branch_status(branch_id, BranchServiceStatus.STARTING)
 
 
@@ -1026,7 +1024,7 @@ async def _public(branch: Branch) -> BranchPublic:
     project = await branch.awaitable_attrs.project
 
     db_host = _resolve_db_host(branch) or ""
-    port = branch.db_port or (await _resolve_branch_db_port(branch.id)) or 0
+    port = await _get_branch_db_port(branch)
 
     # 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/migrations/versions/2b4e8f1a6c03_branch_db_port.py

@@ -1,87 +1,21 @@
-"""Add and backfill branch.db_port
+"""Add 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."""
@@ -90,38 +24,6 @@ def upgrade() -> None:
         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."""