feat: auditlogs 2 in vfolder

src/ai/backend/client/cli/admin/__init__.py

@@ -11,6 +11,7 @@ def admin():
 from . import (  # noqa
     acl,
     agent,
+    audit_logs,
     domain,
     etcd,
     group,

src/ai/backend/client/cli/admin/audit_logs.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import sys
+
+import click
+
+from ai.backend.client.output.fields import auditlog_fields
+from ai.backend.client.session import Session
+
+from ..extensions import pass_ctx_obj
+from ..types import CLIContext
+from . import admin
+
+
+@admin.group()
+def audit_logs() -> None:
+    """
+    Events audit logs commands.
+    """
+
+
+@audit_logs.command()
+# @click.pass_obj
+@pass_ctx_obj
+@click.option("-u", "--user-id", type=str, default=None, help="User ID to audit.")
+@click.option("--filter", "filter_", default=None, help="Set the query filter expression.")
+@click.option("--order", default=None, help="Set the query ordering expression.")
+@click.option("--offset", default=0, help="The index of the current page start for pagination.")
+@click.option("--limit", default=None, help="The page size for pagination.")
+def list(ctx: CLIContext, user_id, filter_, order, offset, limit) -> None:
+    """
+    List audit logs.
+    (admin privilege required)
+    """
+    fields = [
+        auditlog_fields["user_id"],
+        auditlog_fields["access_key"],
+        auditlog_fields["email"],
+        auditlog_fields["action"],
+        auditlog_fields["target_type"],
+        auditlog_fields["target"],
+        auditlog_fields["data"],
+        auditlog_fields["created_at"],
+    ]
+    try:
+        with Session() as session:
+            fetch_func = lambda pg_offset, pg_size: session.AuditLog.paginated_list(
+                user_id,
+                fields=fields,
+                page_offset=pg_offset,
+                page_size=pg_size,
+                filter=filter_,
+                order=order,
+            )
+            ctx.output.print_paginated_list(
+                fetch_func,
+                initial_page_offset=offset,
+                page_size=limit,
+            )
+    except Exception as e:
+        ctx.output.print_error(e)
+        sys.exit(1)

src/ai/backend/client/func/audit_logs.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import Sequence, Union
+
+from ai.backend.client.output.fields import auditlog_fields
+from ai.backend.client.output.types import FieldSpec, PaginatedResult
+from ai.backend.client.pagination import fetch_paginated_result
+
+from .base import BaseFunction, api_function
+
+__all__ = ("AuditLog",)
+
+
+_default_list_fields = [
+    auditlog_fields["user_id"],
+    auditlog_fields["access_key"],
+    auditlog_fields["email"],
+    auditlog_fields["action"],
+    auditlog_fields["data"],
+    auditlog_fields["target_type"],
+    auditlog_fields["target"],
+    auditlog_fields["created_at"],
+]
+
+
+class AuditLog(BaseFunction):
+    """
+    Provides management of audit logs.
+    """
+
+    @api_function
+    @classmethod
+    async def paginated_list(
+        cls,
+        user_id: Union[str, str] = None,
+        *,
+        fields: Sequence[FieldSpec] = _default_list_fields,
+        page_offset: int = 0,
+        page_size: int = 20,
+        filter: str = None,
+        order: str = None,
+    ) -> PaginatedResult[dict]:
+        """
+        Fetches the list of audit logs.
+        :param user_id: Fetches audit log from a user
+        """
+        variables = {
+            "user_id": (user_id, "String"),  # list by user_id
+            "filter": (filter, "String"),
+            "order": (order, "String"),
+        }
+
+        return await fetch_paginated_result(
+            "auditlog_list",
+            variables,
+            fields,
+            page_offset=page_offset,
+            page_size=page_size,
+        )

src/ai/backend/client/func/keypair.py

@@ -25,6 +25,8 @@
     keypair_fields["secret_key"],
     keypair_fields["is_active"],
     keypair_fields["is_admin"],
+    keypair_fields["rate_limit"],
+    keypair_fields["resource_policy"],
 )
 
 _default_result_fields = (

src/ai/backend/client/func/user.py

@@ -57,6 +57,15 @@
     user_fields["main_access_key"],
 )
 
+_user_info_fields = (
+    user_fields["username"],
+    user_fields["full_name"],
+    user_fields["domain_name"],
+    user_fields["role"],
+    user_fields["status"],
+    user_fields["description"],
+)
+
 
 class UserRole(enum.StrEnum):
     """
@@ -257,6 +266,42 @@ async def detail_by_uuid(
         )
         return data["user_from_uuid"]
 
+    @api_function
+    @classmethod
+    async def detail_by_email(
+        cls,
+        email: str = None,
+        fields: Sequence[FieldSpec] = _user_info_fields,
+    ) -> Sequence[dict]:
+        """
+        Fetch basic information of a user by user's email. If email is not specified,
+        requester's information will be returned.
+        :param email: email of the user to fetch.
+        :param fields: Additional per-user query fields to fetch.
+        """
+        if email is None:
+            query = textwrap.dedent(
+                """\
+                    query {
+                        user {$fields}
+                    }
+                """
+            )
+        else:
+            query = textwrap.dedent(
+                """\
+                    query($email: String!) {
+                        user_from_email(email: $email) {$fields}
+                    }
+                """
+            )
+        query = query.replace("$fields", " ".join(f.field_ref for f in fields))
+        print("query{0}", query)
+        variables = {"email": email}
+        data = await api_session.get().Admin._query(query, variables if email is not None else None)
+        print("data from func", data["user_from_email"])
+        return data["user_from_email"]
+
     @api_function
     @classmethod
     async def create(

src/ai/backend/client/output/fields.py

@@ -326,6 +326,19 @@
 ])
 
 
+auditlog_fields = FieldSet([
+    FieldSpec("type"),
+    FieldSpec("user_id"),
+    FieldSpec("access_key"),
+    FieldSpec("email"),
+    FieldSpec("action"),
+    FieldSpec("data"),
+    FieldSpec("target_type"),
+    FieldSpec("target"),
+    FieldSpec("created_at"),
+])
+
+
 routing_fields = FieldSet([
     FieldSpec("routing_id"),
     FieldSpec("status"),

src/ai/backend/client/session.py

@@ -267,6 +267,7 @@ class BaseSession(metaclass=abc.ABCMeta):
         "VFolder",
         "Dotfile",
         "ServerLog",
+        "AuditLog",
         "Permission",
         "Service",
         "Model",
@@ -294,6 +295,7 @@ def __init__(
         from .func.acl import Permission
         from .func.admin import Admin
         from .func.agent import Agent, AgentWatcher
+        from .func.audit_logs import AuditLog
         from .func.auth import Auth
         from .func.bgtask import BackgroundTask
         from .func.domain import Domain
@@ -339,6 +341,7 @@ def __init__(
         self.VFolder = VFolder
         self.Dotfile = Dotfile
         self.ServerLog = ServerLog
+        self.AuditLog = AuditLog
         self.Permission = Permission
         self.Service = Service
         self.Model = Model

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

@@ -0,0 +1,85 @@
+import contextvars
+import json
+from collections.abc import Mapping, Sequence
+from typing import Any, NamedTuple
+
+from aiohttp import web
+from aiohttp.typedefs import Handler
+
+from .context import RootContext
+
+
+class AuditLogData(NamedTuple):
+    previous: str = json.dumps({})
+    current: str = json.dumps({})
+
+    def to_dict(self) -> Mapping[str, str]:
+        return {
+            "previous": self.previous,
+            "current": self.current,
+        }
+
+
+audit_log_data: contextvars.ContextVar[AuditLogData] = contextvars.ContextVar(
+    "audit_log_data", default=AuditLogData()
+)
+
+audit_log_target: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "audit_log_target", default=None
+)
+
+
+@web.middleware
+async def audit_log_middleware(request: web.Request, handler: Handler) -> web.StreamResponse:
+    from ai.backend.manager.models import AuditLogRow
+
+    root_ctx: RootContext = request.app["_root.context"]
+    success = False
+    exc = None
+
+    try:
+        res = await handler(request)
+        success = True
+        return res
+    except Exception as e:
+        exc = e
+        raise
+    finally:
+        if request.get("audit_log"):
+            async with root_ctx.db.begin_session() as sess:
+                new_log = AuditLogRow(
+                    user_id=request["user"]["uuid"],
+                    access_key=request["keypair"]["access_key"],
+                    email=request["user"]["email"],
+                    action=request["audit_log_action"],
+                    data=audit_log_data.get().to_dict(),
+                    target_type=request["audit_log_target_type"],
+                    success=success,
+                    target=audit_log_target.get(),
+                    rest_resource=f"{request.method} {request.path}",
+                )
+                if exc:
+                    new_log.error = str(exc)
+                sess.add(new_log)
+
+
+def set_target(target: Any) -> None:
+    audit_log_target.set(str(target))
+
+
+def update_previous(
+    data_to_insert: Mapping[str, Any] | Sequence[Any],
+) -> None:
+    prev_audit_log_data = AuditLogData(
+        previous=json.dumps(data_to_insert), current=audit_log_data.get().current
+    )
+    audit_log_data.set(prev_audit_log_data)
+
+
+def update_current(
+    data_to_insert: Mapping[str, Any] | Sequence[Any],
+) -> None:
+    prev_audit_log_data = AuditLogData(
+        previous=audit_log_data.get().previous, current=json.dumps(data_to_insert)
+    )
+    audit_log_data.set(prev_audit_log_data)

src/ai/backend/manager/api/schema.graphql

@@ -12,6 +12,8 @@
   agent(agent_id: String!): Agent
   agent_list(limit: Int!, offset: Int!, filter: String, order: String, scaling_group: String, status: String): AgentList
   agents(scaling_group: String, status: String): [Agent]
+  auditlog_list(limit: Int!, offset: Int!, filter: String, order: String, user_id: String): AuditLogList
+  auditlog(user_id: ID): [AuditLog]
   agent_summary(agent_id: String!): AgentSummary
   agent_summary_list(limit: Int!, offset: Int!, filter: String, order: String, scaling_group: String, status: String): AgentSummaryList
   domain(name: String): Domain
@@ -61,6 +63,7 @@
   user(domain_name: String, email: String): User
   user_from_uuid(domain_name: String, user_id: ID): User
   users(domain_name: String, group_id: UUID, is_active: Boolean, status: String): [User]
+  user_from_email(email: String): User
   user_list(limit: Int!, offset: Int!, filter: String, order: String, domain_name: String, group_id: UUID, is_active: Boolean, status: String): UserList
 
   """Added in 24.03.0."""
@@ -279,6 +282,25 @@
   total_count: Int!
 }
 
+type AuditLogList implements PaginatedList {
+  items: [AuditLog]!
+  total_count: Int!
+}
+
+type AuditLog implements Item {
+  id: ID
+  user_id: String
+  access_key: String
+  email: String
+  action: String
+  data: JSONString
+  target_type: String
+  target: String
+  created_at: DateTime
+  rest_resource: String
+  gql_query: String
+}
+
 """A schema for normal users."""
 type AgentSummary implements Item {
   id: ID