feature: add git service token management

changes/1136.feature.md

@@ -0,0 +1 @@
+Add git services(GitHub & Gitlab) token management

src/ai/backend/manager/api/userconfig.py

@@ -1,8 +1,10 @@
+import json
 import logging
 import re
 from typing import TYPE_CHECKING, Any, Tuple
 
 import aiohttp_cors
+import sqlalchemy as sa
 import trafaret as t
 from aiohttp import web
 
@@ -11,6 +13,8 @@
 
 from ..models import (
     MAXIMUM_DOTFILE_SIZE,
+    git_token,
+    git_tokens,
     keypairs,
     query_accessible_vfolders,
     query_bootstrap_script,
@@ -235,6 +239,47 @@ async def get_bootstrap_script(request: web.Request) -> web.Response:
         return web.json_response(script)
 
 
+@auth_required
+@server_status_required(READ_ALLOWED)
+async def get_git_tokens(request: web.Request) -> web.Response:
+    root_ctx: RootContext = request.app["_root.context"]
+    user_uuid = request["user"]["uuid"]
+
+    async with root_ctx.db.begin_readonly() as conn:
+        query = sa.select([git_tokens.c.domain, git_tokens.c.token]).where(
+            git_tokens.c.user_id == user_uuid
+        )
+        query_result = await conn.execute(query)
+        rows = query_result.fetchall()
+        result = [{"domain": row["domain"], "token": row["token"]} for row in rows]
+    return web.json_response(result, status=200)
+
+
+@auth_required
+@check_api_params(
+    t.Dict(
+        {
+            t.Key("params"): t.String,
+        },
+    )
+)
+async def update_git_tokens(request: web.Request, params: Any) -> web.Response:
+    resp = []
+    root_ctx: RootContext = request.app["_root.context"]
+    user_uuid = request["user"]["uuid"]
+
+    # get token list then
+    token_list = params["params"]
+    async with root_ctx.db.begin() as conn:
+        await git_token.insert_update_git_tokens(conn, user_uuid, token_list)
+
+    token_list_json = json.loads(token_list)
+    for key, value in token_list_json.items():
+        resp.append({"domain": key, "token": value})
+
+    return web.json_response(resp)
+
+
 async def init(app: web.Application) -> None:
     pass
 
@@ -258,5 +303,7 @@ def create_app(
     cors.add(app.router.add_route("DELETE", "/dotfiles", delete))
     cors.add(app.router.add_route("POST", "/bootstrap-script", update_bootstrap_script))
     cors.add(app.router.add_route("GET", "/bootstrap-script", get_bootstrap_script))
+    cors.add(app.router.add_route("GET", "/git-tokens", get_git_tokens))
+    cors.add(app.router.add_route("PATCH", "/git-tokens", update_git_tokens))
 
     return app, []

src/ai/backend/manager/models/__init__.py

@@ -4,6 +4,7 @@
 from . import dotfile as _dotfile
 from . import endpoint as _endpoint
 from . import error_logs as _errorlogs
+from . import git_token as _gtoken
 from . import group as _group
 from . import image as _image
 from . import kernel as _kernel
@@ -42,6 +43,7 @@
     *_sessiontemplate.__all__,
     *_storage.__all__,
     *_errorlogs.__all__,
+    *_gtoken.__all__,
 )
 
 from .acl import *  # noqa
@@ -50,6 +52,7 @@
 from .dotfile import *  # noqa
 from .endpoint import *  # noqa
 from .error_logs import *  # noqa
+from .git_token import *  # noqa
 from .group import *  # noqa
 from .image import *  # noqa
 from .kernel import *  # noqa

src/ai/backend/manager/models/alembic/versions/3d31a989dd0a_add_git_token_table.py

@@ -0,0 +1,46 @@
+"""add git_token table
+
+Revision ID: 3d31a989dd0a
+Revises: 02535458c0b3
+Create Date: 2023-08-31 01:04:48.109184
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+from ai.backend.manager.models.base import GUID
+
+# revision identifiers, used by Alembic.
+revision = "3d31a989dd0a"
+down_revision = "02535458c0b3"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        "git_tokens",
+        sa.Column("user_id", GUID(), nullable=False, index=True),
+        sa.Column("domain", sa.String(length=200), nullable=False, index=True),
+        sa.Column("token", sa.String(length=200)),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
+        sa.Column(
+            "modified_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=True,
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["users.uuid"],
+            name=op.f("fk_git_tokens_user_id_users"),
+            onupdate="CASCADE",
+            ondelete="CASCADE",
+        ),
+        sa.UniqueConstraint("user_id", "domain", name="uq_git_tokens_user_id_domain"),
+    )
+    # Create a default github and gitlab tokens
+
+
+def downgrade():
+    op.drop_table("git_tokens")

src/ai/backend/manager/models/git_token.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import datetime
+import json
+import logging
+import uuid
+from typing import Sequence
+
+import graphene
+import sqlalchemy as sa
+from graphene.types.datetime import DateTime as GQLDateTime
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection
+from sqlalchemy.sql import text
+
+from ai.backend.common.logging import BraceStyleAdapter
+
+from .base import GUID, Base, metadata
+
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))  # type: ignore[name-defined]
+
+
+__all__: Sequence[str] = (
+    "git_tokens",
+    "GitToken",
+    "GitTokenRow",
+    "insert_update_git_tokens",
+)
+
+git_tokens = sa.Table(
+    "git_tokens",
+    metadata,
+    sa.Column(
+        "user_id",
+        GUID,
+        sa.ForeignKey("users.uuid", onupdate="CASCADE", ondelete="CASCADE"),
+        index=True,
+        nullable=False,
+        primary_key=True,
+    ),
+    sa.Column("domain", sa.String(length=200), unique=True, primary_key=True, nullable=False),
+    sa.Column("token", sa.String(length=200)),
+    sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
+    sa.Column(
+        "modified_at",
+        sa.DateTime(timezone=True),
+        server_default=sa.func.now(),
+        onupdate=sa.func.current_timestamp(),
+    ),
+)
+
+
+async def insert_update_git_tokens(
+    conn: SAConnection, user_uuid: uuid.UUID, token_list_str: str
+) -> None:
+    token_list = json.loads(token_list_str)
+    domain_list = []
+    # check domain_list
+    for key, value in token_list.items():
+        domain_list.append(key)
+    # delete some non-exist key on domain list
+    delete_stmt = sa.delete(git_tokens).where(
+        sa.and_(git_tokens.c.user_id == user_uuid), sa.not_(git_tokens.c.domain.in_(domain_list))
+    )
+    await conn.execute(delete_stmt)
+
+    for key, value in token_list.items():
+        data = {
+            "user_id": user_uuid,
+            "domain": key,
+            "token": value,
+            "modified_at": datetime.datetime.now(),
+        }
+        query = insert(git_tokens).values(data)
+        query = query.on_conflict_do_update(
+            index_elements=["user_id", "domain"],
+            set_=dict(token=query.excluded.token, modified_at=text("NOW()")),
+        )
+        await conn.execute(query)
+
+    await conn.commit()
+
+
+class GitToken(graphene.ObjectType):
+    user_id = graphene.UUID()
+    domain = graphene.String()
+    token = graphene.String()
+    created_at = GQLDateTime()
+    modified_at = GQLDateTime()
+
+
+class GitTokenRow(Base):
+    __table__ = git_tokens