add staff Google Drive OAuth integration

.env.example

@@ -28,4 +28,8 @@ PGADMIN_PORT=5050
 jwt_secret=super_secret_jwt_key
 jwt_algorithm=HS256
 encryption_key=super_secret_encryption_key
-totp_issuer=MultiAI
+totp_issuer=MultiAI
+GOOGLE_CLIENT_ID=
+GOOGLE_CLIENT_SECRET=
+GOOGLE_REDIRECT_URI=http://localhost:8000/staff/drive/callback
+GOOGLE_OAUTH_SCOPES=https://www.googleapis.com/auth/drive.readonly openid email profile

app/container.py

@@ -5,9 +5,12 @@
 from db.generated import user as user_queries
 from db.generated import session as session_queries
 from db.generated import devices as device_queries
+from db.generated import stuff_user as staff_user_queries
+from db.generated import staff_drive_connections as staff_drive_queries
 from app.service.users import AuthService
 from app.service.session import SessionService
 from app.service.device import DeviceService
+from app.service.staff_drive import StaffDriveService
 
 
 
@@ -21,6 +24,8 @@ def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
         self.user_querier = user_queries.AsyncQuerier(conn)
         self.session_querier = session_queries.AsyncQuerier(conn)
         self.device_querier = device_queries.AsyncQuerier(conn)
+        self.staff_user_querier = staff_user_queries.AsyncQuerier(conn)
+        self.staff_drive_querier = staff_drive_queries.AsyncQuerier(conn)
 
         # services
         self.session_service = SessionService()
@@ -40,6 +45,12 @@ def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
             session_querier=self.session_querier,
         )
 
+        self.staff_drive_service = StaffDriveService(
+            staff_user_querier=self.staff_user_querier,
+            drive_connection_querier=self.staff_drive_querier,
+            redis=self.redis,
+        )
+
 
 
 async def get_container(

app/core/config.py

@@ -36,6 +36,14 @@ class Settings(BaseSettings):
     encryption_key: str
     totp_issuer: str = "multAI"
 
+    # Google Drive OAuth
+    GOOGLE_CLIENT_ID: str = ""
+    GOOGLE_CLIENT_SECRET: str = ""
+    GOOGLE_REDIRECT_URI: str = ""
+    GOOGLE_OAUTH_SCOPES: str = (
+        "https://www.googleapis.com/auth/drive.readonly openid email profile"
+    )
+
     class Config:
         env_file = ".env"
         extra = "ignore"

app/deps/staff_auth.py

@@ -0,0 +1,24 @@
+from typing import Annotated
+import uuid
+
+from fastapi import Depends, Header
+
+from app.container import Container, get_container
+from app.core.exceptions import AppException
+from db.generated.models import StaffUser
+
+
+async def get_current_staff_user(
+    x_staff_user_id: Annotated[str, Header(alias="X-Staff-User-Id")],
+    container: Annotated[Container, Depends(get_container)],
+) -> StaffUser:
+    try:
+        staff_user_id = uuid.UUID(x_staff_user_id)
+    except ValueError as exc:
+        raise AppException.bad_request("Invalid X-Staff-User-Id header") from exc
+
+    staff_user = await container.staff_user_querier.get_staff_user_by_id(id=staff_user_id)
+    if staff_user is None:
+        raise AppException.not_found("Staff user not found")
+
+    return staff_user

app/infra/google_drive.py

@@ -0,0 +1,163 @@
+import asyncio
+import json
+import urllib.error
+import urllib.parse
+import urllib.request
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+
+from app.core.exceptions import AppException
+from app.core.config import settings
+
+
+GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
+GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
+GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v2/userinfo"
+
+
+@dataclass
+class GoogleTokenResponse:
+    access_token: str
+    refresh_token: str | None
+    expires_at: datetime | None
+    scope: str
+    token_type: str
+
+
+@dataclass
+class GoogleUserInfo:
+    id: str
+    email: str
+    verified_email: bool
+
+
+class GoogleDriveClient:
+    @staticmethod
+    def _require_str(data: dict[str, object], key: str) -> str:
+        value = data.get(key)
+        if not isinstance(value, str) or not value:
+            raise AppException.bad_request(f"Google response missing '{key}'")
+        return value
+
+    @staticmethod
+    def _optional_str(data: dict[str, object], key: str) -> str | None:
+        value = data.get(key)
+        if value is None:
+            return None
+        if not isinstance(value, str):
+            raise AppException.bad_request(f"Google response field '{key}' is invalid")
+        return value
+
+    @staticmethod
+    def validate_settings() -> None:
+        if (
+            not settings.GOOGLE_CLIENT_ID
+            or not settings.GOOGLE_CLIENT_SECRET
+            or not settings.GOOGLE_REDIRECT_URI
+        ):
+            raise AppException.bad_request(
+                "Google Drive OAuth is not configured in environment variables"
+            )
+
+    @staticmethod
+    def build_consent_url(state: str) -> str:
+        GoogleDriveClient.validate_settings()
+        query = urllib.parse.urlencode(
+            {
+                "client_id": settings.GOOGLE_CLIENT_ID,
+                "redirect_uri": settings.GOOGLE_REDIRECT_URI,
+                "response_type": "code",
+                "scope": settings.GOOGLE_OAUTH_SCOPES,
+                "access_type": "offline",
+                "include_granted_scopes": "true",
+                "prompt": "consent",
+                "state": state,
+            }
+        )
+        return f"{GOOGLE_AUTH_URL}?{query}"
+
+    @staticmethod
+    async def exchange_code(code: str) -> GoogleTokenResponse:
+        GoogleDriveClient.validate_settings()
+        payload = {
+            "code": code,
+            "client_id": settings.GOOGLE_CLIENT_ID,
+            "client_secret": settings.GOOGLE_CLIENT_SECRET,
+            "redirect_uri": settings.GOOGLE_REDIRECT_URI,
+            "grant_type": "authorization_code",
+        }
+        data = await GoogleDriveClient._post_form(GOOGLE_TOKEN_URL, payload)
+        expires_at = None
+        expires_in = data.get("expires_in")
+        if isinstance(expires_in, int):
+            expires_at = datetime.now(timezone.utc) + timedelta(seconds=expires_in)
+
+        return GoogleTokenResponse(
+            access_token=GoogleDriveClient._require_str(data, "access_token"),
+            refresh_token=GoogleDriveClient._optional_str(data, "refresh_token"),
+            expires_at=expires_at,
+            scope=GoogleDriveClient._optional_str(data, "scope")
+            or settings.GOOGLE_OAUTH_SCOPES,
+            token_type=GoogleDriveClient._optional_str(data, "token_type")
+            or "Bearer",
+        )
+
+    @staticmethod
+    async def get_user_info(access_token: str) -> GoogleUserInfo:
+        data = await GoogleDriveClient._get_json(
+            GOOGLE_USERINFO_URL,
+            headers={"Authorization": f"Bearer {access_token}"},
+        )
+        return GoogleUserInfo(
+            id=GoogleDriveClient._require_str(data, "id"),
+            email=GoogleDriveClient._require_str(data, "email"),
+            verified_email=bool(data.get("verified_email", False)),
+        )
+
+    @staticmethod
+    async def _post_form(url: str, payload: dict[str, str]) -> dict[str, object]:
+        encoded = urllib.parse.urlencode(payload).encode("utf-8")
+
+        def _request() -> dict[str, object]:
+            request = urllib.request.Request(
+                url,
+                data=encoded,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+                method="POST",
+            )
+            try:
+                with urllib.request.urlopen(request, timeout=15) as response:
+                    return json.loads(response.read().decode("utf-8"))
+            except urllib.error.HTTPError as exc:
+                details = exc.read().decode("utf-8", errors="ignore")
+                raise AppException.bad_request(
+                    f"Google token exchange failed: {details or exc.reason}"
+                ) from exc
+            except urllib.error.URLError as exc:
+                raise AppException.internal_error(
+                    "Unable to reach Google OAuth endpoints"
+                ) from exc
+
+        return await asyncio.to_thread(_request)
+
+    @staticmethod
+    async def _get_json(
+        url: str,
+        headers: dict[str, str] | None = None,
+    ) -> dict[str, object]:
+        def _request() -> dict[str, object]:
+            request = urllib.request.Request(url, headers=headers or {}, method="GET")
+            try:
+                with urllib.request.urlopen(request, timeout=15) as response:
+                    return json.loads(response.read().decode("utf-8"))
+            except urllib.error.HTTPError as exc:
+                details = exc.read().decode("utf-8", errors="ignore")
+                raise AppException.bad_request(
+                    f"Google user info request failed: {details or exc.reason}"
+                ) from exc
+            except urllib.error.URLError as exc:
+                raise AppException.internal_error(
+                    "Unable to reach Google APIs"
+                ) from exc
+
+        return await asyncio.to_thread(_request)

app/main.py

@@ -12,6 +12,7 @@
 from app.infra.nats import NatsClient
 from app.infra.redis import RedisClient
 from app.router.mobile.auth import router as mobile_auth_router
+from app.router.staff.drive import router as staff_drive_router
 
 
 
@@ -115,3 +116,4 @@ def health_check() -> dict[str, str]:
 
 
 app.include_router(mobile_auth_router, prefix="/mobile")
+app.include_router(staff_drive_router, prefix="/staff")

app/router/staff/__init__.py

@@ -0,0 +1 @@
+

app/router/staff/drive.py

@@ -0,0 +1,70 @@
+from fastapi import APIRouter, Depends, Query
+
+from app.container import Container, get_container
+from app.core.exceptions import AppException
+from app.deps.staff_auth import get_current_staff_user
+from app.schema.response.staff.drive import (
+    GoogleDriveCallbackResponse,
+    GoogleDriveConnectResponse,
+    GoogleDriveConnectionStatusResponse,
+    GoogleDriveDisconnectResponse,
+)
+from db.generated.models import StaffUser
+
+
+router = APIRouter(prefix="/drive", tags=["staff-drive"])
+
+
+@router.get("/connect", response_model=GoogleDriveConnectResponse)
+async def connect_google_drive(
+    current_staff_user: StaffUser = Depends(get_current_staff_user),
+    container: Container = Depends(get_container),
+) -> GoogleDriveConnectResponse:
+    authorization_url, state = await container.staff_drive_service.create_connect_url(
+        current_staff_user
+    )
+    return GoogleDriveConnectResponse(authorization_url=authorization_url, state=state)
+
+
+@router.get("/callback", response_model=GoogleDriveCallbackResponse)
+async def google_drive_callback(
+    code: str = Query(...),
+    state: str = Query(...),
+    error: str | None = Query(default=None),
+    container: Container = Depends(get_container),
+) -> GoogleDriveCallbackResponse:
+    if error is not None:
+        raise AppException.bad_request(f"Google OAuth error: {error}")
+
+    connection = await container.staff_drive_service.handle_callback(code, state)
+    return GoogleDriveCallbackResponse(
+        message="Google Drive connected successfully",
+        google_email=connection.google_email,
+    )
+
+
+@router.get("/status", response_model=GoogleDriveConnectionStatusResponse)
+async def google_drive_status(
+    current_staff_user: StaffUser = Depends(get_current_staff_user),
+    container: Container = Depends(get_container),
+) -> GoogleDriveConnectionStatusResponse:
+    connection = await container.staff_drive_service.get_status(current_staff_user.id)
+    if connection is None:
+        return GoogleDriveConnectionStatusResponse(connected=False)
+
+    return GoogleDriveConnectionStatusResponse(
+        connected=True,
+        google_email=connection.google_email,
+        scopes=[scope for scope in connection.scopes.split(" ") if scope],
+        connected_at=connection.connected_at,
+        token_expires_at=connection.token_expires_at,
+    )
+
+
+@router.post("/disconnect", response_model=GoogleDriveDisconnectResponse)
+async def disconnect_google_drive(
+    current_staff_user: StaffUser = Depends(get_current_staff_user),
+    container: Container = Depends(get_container),
+) -> GoogleDriveDisconnectResponse:
+    await container.staff_drive_service.disconnect(current_staff_user.id)
+    return GoogleDriveDisconnectResponse(message="Google Drive disconnected successfully")

app/schema/response/staff/drive.py

@@ -0,0 +1,25 @@
+from datetime import datetime
+
+from pydantic import BaseModel, Field
+
+
+class GoogleDriveConnectResponse(BaseModel):
+    authorization_url: str
+    state: str
+
+
+class GoogleDriveConnectionStatusResponse(BaseModel):
+    connected: bool
+    google_email: str | None = None
+    scopes: list[str] = Field(default_factory=list)
+    connected_at: datetime | None = None
+    token_expires_at: datetime | None = None
+
+
+class GoogleDriveCallbackResponse(BaseModel):
+    message: str
+    google_email: str
+
+
+class GoogleDriveDisconnectResponse(BaseModel):
+    message: str