From 6983d4d665786f6496287cac57f0063dad122ae3 Mon Sep 17 00:00:00 2001
From: Sanghun Lee <sanghun@lablup.com>
Date: Fri, 20 Dec 2024 12:27:56 +0900
Subject: [PATCH] feat: Add uid and gid columns to users and kernels tables

---
 .../f6ca2f2d04c1_add_uid_and_gid_columns.py   | 43 +++++++++++++++++++
 src/ai/backend/manager/models/kernel.py       |  3 ++
 src/ai/backend/manager/models/user.py         |  2 +
 3 files changed, 48 insertions(+)
 create mode 100644 src/ai/backend/manager/models/alembic/versions/f6ca2f2d04c1_add_uid_and_gid_columns.py

diff --git a/src/ai/backend/manager/models/alembic/versions/f6ca2f2d04c1_add_uid_and_gid_columns.py b/src/ai/backend/manager/models/alembic/versions/f6ca2f2d04c1_add_uid_and_gid_columns.py
new file mode 100644
index 00000000000..f1ae8015eaa
--- /dev/null
+++ b/src/ai/backend/manager/models/alembic/versions/f6ca2f2d04c1_add_uid_and_gid_columns.py
@@ -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")
diff --git a/src/ai/backend/manager/models/kernel.py b/src/ai/backend/manager/models/kernel.py
index e54f4827fd2..42f09579f98 100644
--- a/src/ai/backend/manager/models/kernel.py
+++ b/src/ai/backend/manager/models/kernel.py
@@ -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
diff --git a/src/ai/backend/manager/models/user.py b/src/ai/backend/manager/models/user.py
index 2c7b207f52e..7acc8a14b11 100644
--- a/src/ai/backend/manager/models/user.py
+++ b/src/ai/backend/manager/models/user.py
@@ -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()),
 )