diff --git a/src/api/organization/project/branch/__init__.py b/src/api/organization/project/branch/__init__.py index 00f4086e..9df87730 100644 --- a/src/api/organization/project/branch/__init__.py +++ b/src/api/organization/project/branch/__init__.py @@ -158,6 +158,21 @@ async def _persist_branch_status(branch_id: Identifier, status: BranchServiceSta await session.commit() +async def _get_branch_db_port(branch: Branch) -> int | None: + if branch.db_port is not None: + return branch.db_port + port = await _resolve_branch_db_port(branch.id) + if port is None: + return None + async with AsyncSessionLocal() as session: + 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: """Best-effort cleanup when provisioning fails so allocations aren't stranded.""" try: @@ -1009,7 +1024,9 @@ 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 = await _get_branch_db_port(branch) + if port is None: + logger.warning("db port not yet available for branch %s", branch.id) # 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) diff --git a/src/models/branch.py b/src/models/branch.py index d049981d..bb62a59b 100644 --- a/src/models/branch.py +++ b/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"),) @@ -351,7 +353,7 @@ def __new__(cls, value: str, description: str): class DatabaseInformation(BaseModel): host: str - port: int + port: int | None username: str name: str encrypted_connection_string: str diff --git a/src/models/migrations/versions/2b4e8f1a6c03_branch_db_port.py b/src/models/migrations/versions/2b4e8f1a6c03_branch_db_port.py new file mode 100644 index 00000000..cbcac6b6 --- /dev/null +++ b/src/models/migrations/versions/2b4e8f1a6c03_branch_db_port.py @@ -0,0 +1,30 @@ +"""Add branch.db_port + +Revision ID: 2b4e8f1a6c03 +Revises: f4f677e4e9b9 +Create Date: 2026-03-19 00:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + +# 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 + + +def upgrade() -> None: + """Upgrade schema.""" + op.add_column( + 'branch', + sa.Column('db_port', sa.Integer(), nullable=True), + ) + + +def downgrade() -> None: + """Downgrade schema.""" + op.drop_column('branch', 'db_port')