Feat/drive workflow

.env.example

@@ -30,7 +30,7 @@ jwt_algorithm=HS256
 encryption_key=super_secret_encryption_key
 totp_issuer=MultiAI
 
-GOOGLE_CLIENT_ID=
-GOOGLE_CLIENT_SECRET=
+GOOGLE_CLIENT_ID=817881186650-a2id6qrap9q70idfhjro7ch54j5qv98r.apps.googleusercontent.com
+GOOGLE_CLIENT_SECRET=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
 GOOGLE_REDIRECT_URI=http://localhost:8000/staff/drive/callback
 GOOGLE_OAUTH_SCOPES=https://www.googleapis.com/auth/drive.readonly openid email profile

app/infra/google_drive.py

@@ -5,6 +5,7 @@
 import urllib.request
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
+from typing import BinaryIO
 
 from app.core.exceptions import AppException
 from app.core.config import settings
@@ -13,6 +14,8 @@
 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"
+GOOGLE_DRIVE_FILES_URL = "https://www.googleapis.com/drive/v3/files"
+GOOGLE_DRIVE_DEFAULT_LIST_FIELDS = "nextPageToken,files(id,name,mimeType,thumbnailLink,iconLink)"
 
 
 @dataclass
@@ -87,20 +90,19 @@ async def exchange_code(code: str) -> GoogleTokenResponse:
             "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 GoogleDriveClient._build_token_response(data)
 
-        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 refresh_token(refresh_token: str) -> GoogleTokenResponse:
+        GoogleDriveClient.validate_settings()
+        payload = {
+            "refresh_token": refresh_token,
+            "client_id": settings.GOOGLE_CLIENT_ID,
+            "client_secret": settings.GOOGLE_CLIENT_SECRET,
+            "grant_type": "refresh_token",
+        }
+        data = await GoogleDriveClient._post_form(GOOGLE_TOKEN_URL, payload)
+        return GoogleDriveClient._build_token_response(data)
 
     @staticmethod
     async def get_user_info(access_token: str) -> GoogleUserInfo:
@@ -114,6 +116,52 @@ async def get_user_info(access_token: str) -> GoogleUserInfo:
             verified_email=bool(data.get("verified_email", False)),
         )
 
+    @staticmethod
+    async def list_files(
+        access_token: str,
+        query: str,
+        page_size: int = 30,
+        page_token: str | None = None,
+        fields: str | None = None,
+    ) -> dict[str, object]:
+        params: dict[str, str] = {
+            "q": query,
+            "pageSize": str(max(1, min(page_size, 1000))),
+            "fields": fields or GOOGLE_DRIVE_DEFAULT_LIST_FIELDS,
+        }
+
+        if page_token:
+            params["pageToken"] = page_token
+
+        return await GoogleDriveClient._get_json(
+            GOOGLE_DRIVE_FILES_URL,
+            params=params,
+            headers=GoogleDriveClient._auth_headers(access_token),
+        )
+
+    @staticmethod
+    async def download_file(access_token: str, file_id: str) -> BinaryIO:
+        def _request() -> BinaryIO:
+            url = f"{GOOGLE_DRIVE_FILES_URL}/{file_id}?alt=media"
+            request = urllib.request.Request(
+                url,
+                headers=GoogleDriveClient._auth_headers(access_token),
+                method="GET",
+            )
+            try:
+                return urllib.request.urlopen(request, timeout=30)
+            except urllib.error.HTTPError as exc:
+                details = exc.read().decode("utf-8", errors="ignore")
+                raise AppException.bad_request(
+                    f"Drive file download failed: {details or exc.reason}"
+                ) from exc
+            except urllib.error.URLError as exc:
+                raise AppException.internal_error(
+                    "Unable to reach Google Drive API"
+                ) from exc
+
+        return await asyncio.to_thread(_request)
+
     @staticmethod
     async def _post_form(url: str, payload: dict[str, str]) -> dict[str, object]:
         encoded = urllib.parse.urlencode(payload).encode("utf-8")
@@ -143,21 +191,47 @@ def _request() -> dict[str, object]:
     @staticmethod
     async def _get_json(
         url: str,
+        *,
+        params: dict[str, str] | None = None,
         headers: dict[str, str] | None = None,
     ) -> dict[str, object]:
         def _request() -> dict[str, object]:
-            request = urllib.request.Request(url, headers=headers or {}, method="GET")
+            target = url
+            if params:
+                target = f"{url}?{urllib.parse.urlencode(params, safe=',()')}"
+            request = urllib.request.Request(target, 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}"
+                    f"Google API 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)
+
+    @staticmethod
+    def _auth_headers(access_token: str) -> dict[str, str]:
+        return {"Authorization": f"Bearer {access_token}"}
+
+    @staticmethod
+    def _build_token_response(data: dict[str, object]) -> GoogleTokenResponse:
+        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",
+        )

app/infra/minio.py

@@ -11,8 +11,8 @@
 
 IMAGES_BUCKET_NAME = "images"
 DOCUMENTS_BUCKET_NAME = "documents"
-WA_SIM_BUCKET_NAME = "wa-sim"
-
+IMAGE_REQUEST_BUCKET_NAME = "image-request"
+HOT_BUCKET_NAME = "hot"
 async def init_minio_client(
     minio_host: str, minio_port: int, minio_root_user: str, minio_root_password: str
 ) -> None:
@@ -23,7 +23,7 @@ async def init_minio_client(
         secure=False,
     )
 
-    for bucket_name in [IMAGES_BUCKET_NAME, DOCUMENTS_BUCKET_NAME, WA_SIM_BUCKET_NAME]:
+    for bucket_name in [IMAGES_BUCKET_NAME, DOCUMENTS_BUCKET_NAME, IMAGE_REQUEST_BUCKET_NAME, HOT_BUCKET_NAME]:
         if not await Bucket.client.bucket_exists(bucket_name):
             await Bucket.client.make_bucket(bucket_name)
 

app/router/staff/drive.py

@@ -3,12 +3,16 @@
 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.request.staff.drive import GoogleDriveImportRequest
 from app.schema.response.staff.drive import (
     GoogleDriveCallbackResponse,
     GoogleDriveConnectResponse,
     GoogleDriveConnectionStatusResponse,
     GoogleDriveDisconnectResponse,
+    GoogleDriveFileListResponse,
+    GoogleDriveImportResponse,
 )
+from app.service.staff_drive import SelectedDriveFile
 from db.generated.models import StaffUser
 
 
@@ -61,10 +65,49 @@ async def google_drive_status(
     )
 
 
+@router.get("/files", response_model=GoogleDriveFileListResponse)
+async def list_google_drive_files(
+    folder_id: str | None = Query(default=None, alias="folderId"),
+    page_token: str | None = Query(default=None, alias="pageToken"),
+    page_size: int = Query(default=30, alias="pageSize", ge=1, le=1000),
+    # current_staff_user: StaffUser = Depends(get_current_staff_user),
+    container: Container = Depends(get_container),
+) -> GoogleDriveFileListResponse:
+    payload = await container.staff_drive_service.list_drive_files(
+        staff_user=current_staff_user,
+        folder_id=folder_id,
+        page_token=page_token,
+        page_size=page_size,
+    )
+    return GoogleDriveFileListResponse(**payload)
+
+
 @router.post("/disconnect", response_model=GoogleDriveDisconnectResponse)
 async def disconnect_google_drive(
-    current_staff_user: StaffUser = Depends(get_current_staff_user),
+    # 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")
+
+
+@router.post("/files/import", response_model=GoogleDriveImportResponse)
+async def import_drive_files(
+    request: GoogleDriveImportRequest,
+    # current_staff_user: StaffUser = Depends(get_current_staff_user),
+    container: Container = Depends(get_container),
+) -> GoogleDriveImportResponse:
+    selections = [
+        SelectedDriveFile(
+            id=file.id,
+            name=file.name,
+            mime_type=file.mime_type,
+        )
+        for file in request.files
+    ]
+
+    result = await container.staff_drive_service.import_images_from_drive(
+        staff_user=current_staff_user,
+        selected_files=selections,
+    )
+    return GoogleDriveImportResponse(files=result)

app/schema/request/staff/drive.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel, Field
+
+
+class GoogleDriveImportFileSelection(BaseModel):
+    id: str
+    name: str
+    mime_type: str | None = Field(default=None, alias="mimeType")
+
+
+class GoogleDriveImportRequest(BaseModel):
+    files: list[GoogleDriveImportFileSelection] = Field(default_factory=list)

app/schema/response/staff/drive.py

@@ -23,3 +23,28 @@ class GoogleDriveCallbackResponse(BaseModel):
 
 class GoogleDriveDisconnectResponse(BaseModel):
     message: str
+
+
+class GoogleDriveFileItem(BaseModel):
+    id: str
+    name: str
+    mime_type: str | None = Field(default=None, alias="mimeType")
+    thumbnail_link: str | None = Field(default=None, alias="thumbnailLink")
+    icon_link: str | None = Field(default=None, alias="iconLink")
+
+
+class GoogleDriveFileListResponse(BaseModel):
+    files: list[GoogleDriveFileItem]
+    next_page_token: str | None = Field(default=None, alias="nextPageToken")
+
+
+class GoogleDriveImportFileResult(BaseModel):
+    drive_file_id: str
+    original_file_name: str
+    minio_bucket: str
+    minio_object_name: str
+    minio_object_path: str
+
+
+class GoogleDriveImportResponse(BaseModel):
+    files: list[GoogleDriveImportFileResult]