Skip to content

Commit

Permalink
feat: Add uid and gid columns to users and kernels tables
Browse files Browse the repository at this point in the history
  • Loading branch information
fregataa committed Dec 23, 2024
1 parent 42b4030 commit 4cd9c9a
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
"""add uid and gid columns
Revision ID: f6ca2f2d04c1
Revises: 0bb88d5a46bf
Create Date: 2024-12-20 12:16:07.077845
"""

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "f6ca2f2d04c1"
down_revision = "0bb88d5a46bf"
branch_labels = None
depends_on = None


def upgrade() -> None:
op.add_column(
"kernels", sa.Column("uid", sa.Integer(), server_default=sa.text("NULL"), nullable=True)
)
op.add_column(
"kernels",
sa.Column("gids", sa.ARRAY(sa.Integer()), server_default=sa.text("NULL"), nullable=True),
)
op.add_column(
"users",
sa.Column("container_uid", sa.Integer(), server_default=sa.text("NULL"), nullable=True),
)
op.add_column(
"users",
sa.Column(
"container_gids", sa.ARRAY(sa.Integer()), server_default=sa.text("NULL"), nullable=True
),
)


def downgrade() -> None:
op.drop_column("users", "container_gids")
op.drop_column("users", "container_uid")
op.drop_column("kernels", "gids")
op.drop_column("kernels", "uid")
3 changes: 3 additions & 0 deletions src/ai/backend/manager/models/kernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -416,6 +416,9 @@ class KernelRow(Base):
cluster_hostname = sa.Column(
"cluster_hostname", sa.String(length=64), nullable=False, default=default_hostname
)
uid = sa.Column("uid", sa.Integer, nullable=True, server_default=sa.null())
gids = sa.Column("gids", sa.ARRAY(sa.Integer), nullable=True, server_default=sa.null())

# Resource ownership
scaling_group = sa.Column(
"scaling_group", sa.ForeignKey("scaling_groups.name"), index=True, nullable=True
Expand Down
2 changes: 2 additions & 0 deletions src/ai/backend/manager/models/user.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,8 @@ class UserStatus(enum.StrEnum):
sa.ForeignKey("keypairs.access_key", ondelete="SET NULL"),
nullable=True, # keypairs.user is non-nullable
),
sa.Column("container_uid", sa.Integer, nullable=True, server_default=sa.null()),
sa.Column("container_gids", sa.ARRAY(sa.Integer), nullable=True, server_default=sa.null()),
)


Expand Down

0 comments on commit 4cd9c9a

Please sign in to comment.