Implement support for downloading videos from Google Drive

backend/google_api/admin.py

@@ -16,6 +16,10 @@
 class GoogleCloudTokenInline(admin.StackedInline):
     model = GoogleCloudToken
     autocomplete_fields = ("admin_user",)
+    extra = 0
+
+    def has_add_permission(self, request, obj=None):
+        return False
 
 
 @admin.register(GoogleCloudOAuthCredential)

backend/google_api/migrations/0005_googlecloudoauthcredential_quota_limit_for_drive.py

@@ -0,0 +1,17 @@
+# Generated by Django 5.2.8 on 2026-08-07 10:36
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("google_api", "0004_rename_credential_usedrequestquota_credentials"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="googlecloudoauthcredential",
+            name="quota_limit_for_drive",
+            field=models.IntegerField(default=10000),
+        ),
+    ]

backend/google_api/models.py

@@ -7,6 +7,7 @@
 from django.utils import timezone
 
 DAILY_YOUTUBE_QUOTA = 10_000
+DAILY_DRIVE_QUOTA = 10_000
 
 
 class GoogleCloudOAuthCredentialQuerySet(models.QuerySet):
@@ -52,6 +53,7 @@ class GoogleCloudOAuthCredential(models.Model):
     auth_provider_x509_cert_url = models.URLField()
 
     quota_limit_for_youtube = models.IntegerField(default=DAILY_YOUTUBE_QUOTA)
+    quota_limit_for_drive = models.IntegerField(default=DAILY_DRIVE_QUOTA)
 
     objects = GoogleCloudOAuthCredentialQuerySet.as_manager()
 

backend/google_api/sdk.py

@@ -1,12 +1,22 @@
 import inspect
+import requests
 from google_api.exceptions import NoGoogleCloudQuotaLeftError
 from google_api.models import GoogleCloudOAuthCredential, UsedRequestQuota
 from googleapiclient.discovery import build
 from apiclient.http import MediaFileUpload
+from google.auth.transport.requests import Request
 from google.oauth2.credentials import Credentials
+import logging
 
+logger = logging.getLogger(__name__)
 
-GOOGLE_CLOUD_SCOPES = ["https://www.googleapis.com/auth/youtube"]
+GOOGLE_CLOUD_SCOPES = [
+    "https://www.googleapis.com/auth/youtube",
+    "https://www.googleapis.com/auth/drive.readonly",
+]
+
+DRIVE_API_URL = "https://www.googleapis.com/drive/v3"
+DRIVE_LIST_PAGE_SIZE = 1000
 
 
 def get_available_credentials(service, min_quota):
@@ -42,10 +52,15 @@ def _add_quota(credentials):
         )
 
     def wrapper(func):
+        # Callers that make several calls for one piece of work pass the
+        # credentials they already hold, so every call runs as the same Google
+        # account. Left out, each call picks its own account and the work can
+        # straddle two of them. The quota is charged either way.
         if inspect.isgeneratorfunction(func):
 
-            def wrapped(*args, **kwargs):
-                credentials = get_available_credentials(service, quota)
+            def wrapped(*args, credentials=None, **kwargs):
+                if credentials is None:
+                    credentials = get_available_credentials(service, quota)
                 try:
                     for value in func(*args, credentials=credentials, **kwargs):
                         yield value
@@ -54,8 +69,9 @@ def wrapped(*args, **kwargs):
 
         else:
 
-            def wrapped(*args, **kwargs):
-                credentials = get_available_credentials(service, quota)
+            def wrapped(*args, credentials=None, **kwargs):
+                if credentials is None:
+                    credentials = get_available_credentials(service, quota)
                 try:
                     ret_value = func(*args, credentials=credentials, **kwargs)
                 finally:
@@ -67,6 +83,70 @@ def wrapped(*args, **kwargs):
     return wrapper
 
 
+def refreshed(credentials: Credentials) -> Credentials:
+    """Credentials rebuilt from a stored token carry no expiry, so google-auth
+    treats them as expired and credentials.token is the stale stored value.
+    googleapiclient refreshes lazily on its own; direct requests calls must not.
+    """
+    if not credentials.valid:
+        credentials.refresh(Request())
+
+    return credentials
+
+
+def drive_headers(credentials: Credentials) -> dict:
+    return {"Authorization": f"Bearer {refreshed(credentials).token}"}
+
+
+@count_quota("drive", 1)
+def get_drive_credentials(*, credentials: Credentials) -> Credentials:
+    """Refreshed credentials for callers that drive their own Drive requests."""
+    return refreshed(credentials)
+
+
+@count_quota("drive", 1)
+def drive_file_metadata(*, file_id: str, credentials: Credentials) -> dict:
+    response = requests.get(
+        f"{DRIVE_API_URL}/files/{file_id}",
+        params={
+            "fields": "id,name,size,mimeType",
+            "supportsAllDrives": "true",
+        },
+        headers=drive_headers(credentials),
+    )
+    response.raise_for_status()
+    return response.json()
+
+
+@count_quota("drive", 1)
+def drive_list_files_in_folder(*, folder_id: str, credentials: Credentials):
+    headers = drive_headers(credentials)
+    params = {
+        "q": f"'{folder_id}' in parents and trashed = false",
+        "fields": "nextPageToken,files(id,name,mimeType,size)",
+        "pageSize": DRIVE_LIST_PAGE_SIZE,
+        "supportsAllDrives": "true",
+        "includeItemsFromAllDrives": "true",
+    }
+
+    while True:
+        response = requests.get(
+            f"{DRIVE_API_URL}/files", params=params, headers=headers
+        )
+        response.raise_for_status()
+        payload = response.json()
+
+        logger.info("List files in folder %s", folder_id)
+
+        yield from payload.get("files", [])
+
+        page_token = payload.get("nextPageToken")
+        if not page_token:
+            return
+
+        params["pageToken"] = page_token
+
+
 @count_quota("youtube", 1600)
 def youtube_videos_insert(
     *,

backend/google_api/tests/test_models.py

@@ -1,5 +1,9 @@
 import time_machine
-from google_api.models import GoogleCloudOAuthCredential, UsedRequestQuota
+from google_api.models import (
+    GoogleCloudOAuthCredential,
+    GoogleCloudToken,
+    UsedRequestQuota,
+)
 import pytest
 
 
@@ -29,6 +33,39 @@ def test_with_quota_left():
         assert result.youtube_quota_left == 10_000
 
 
+def test_with_quota_left_for_drive():
+    credential = GoogleCloudOAuthCredential.objects.create()
+
+    with time_machine.travel("2023-10-10 00:00:00", tick=False):
+        result = GoogleCloudOAuthCredential.objects.with_quota_left("drive").get()
+        assert result.drive_quota_left == credential.quota_limit_for_drive
+
+        UsedRequestQuota.objects.create(
+            credentials=credential,
+            cost=1,
+            service="drive",
+        )
+
+        result = GoogleCloudOAuthCredential.objects.with_quota_left("drive").get()
+        assert result.drive_quota_left == credential.quota_limit_for_drive - 1
+
+
+def test_get_available_credentials_token_for_drive(admin_user):
+    credential = GoogleCloudOAuthCredential.objects.create()
+    GoogleCloudToken.objects.create(
+        oauth_credential=credential,
+        token="token",
+        admin_user=admin_user,
+    )
+
+    token = GoogleCloudOAuthCredential.get_available_credentials_token(
+        service="drive", min_quota=1
+    )
+
+    assert token is not None
+    assert token.oauth_credential_id == credential.id
+
+
 def test_get_by_client_id():
     credential = GoogleCloudOAuthCredential.objects.create(client_id="test123")
 

backend/google_api/tests/test_sdk.py

@@ -9,14 +9,127 @@
 from google_api.exceptions import NoGoogleCloudQuotaLeftError
 from google_api.sdk import (
     count_quota,
+    drive_file_metadata,
+    drive_list_files_in_folder,
     get_available_credentials,
+    get_drive_credentials,
     youtube_videos_insert,
     youtube_videos_set_thumbnail,
 )
 import pytest
 
 pytestmark = pytest.mark.django_db
 
+GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
+
+
+@pytest.fixture
+def drive_credential(admin_user):
+    stored_credential = GoogleCloudOAuthCredential.objects.create()
+    GoogleCloudToken.objects.create(
+        oauth_credential=stored_credential,
+        token="stale-token",
+        refresh_token="refresh-token",
+        admin_user=admin_user,
+    )
+    return stored_credential
+
+
+def mock_token_refresh(requests_mock):
+    return requests_mock.post(
+        GOOGLE_TOKEN_URL,
+        json={
+            "access_token": "refreshed-token",
+            "expires_in": 3600,
+            "token_type": "Bearer",
+        },
+    )
+
+
+def test_get_drive_credentials_refreshes_the_stored_token(
+    requests_mock, drive_credential
+):
+    mock_token_refresh(requests_mock)
+
+    credentials = get_drive_credentials()
+
+    assert credentials.token == "refreshed-token"
+
+
+def test_drive_file_metadata_sends_the_refreshed_bearer_token(
+    requests_mock, drive_credential
+):
+    mock_token_refresh(requests_mock)
+    metadata_mock = requests_mock.get(
+        "https://www.googleapis.com/drive/v3/files/file123",
+        json={
+            "id": "file123",
+            "name": "talk.mp4",
+            "size": "1024",
+            "mimeType": "video/mp4",
+        },
+    )
+
+    metadata = drive_file_metadata(file_id="file123")
+
+    assert metadata["name"] == "talk.mp4"
+    assert (
+        metadata_mock.last_request.headers["Authorization"] == "Bearer refreshed-token"
+    )
+    assert metadata_mock.last_request.qs["supportsalldrives"] == ["true"]
+    assert drive_credential.usedrequestquota_set.filter(service="drive").count() == 1
+
+
+def test_drive_list_files_in_folder_follows_pagination(requests_mock, drive_credential):
+    mock_token_refresh(requests_mock)
+    requests_mock.get(
+        "https://www.googleapis.com/drive/v3/files",
+        [
+            {
+                "json": {
+                    "files": [{"id": "1", "name": "a.mp4", "mimeType": "video/mp4"}],
+                    "nextPageToken": "page-2",
+                }
+            },
+            {
+                "json": {
+                    "files": [{"id": "2", "name": "b.mp4", "mimeType": "video/mp4"}]
+                }
+            },
+        ],
+    )
+
+    files = list(drive_list_files_in_folder(folder_id="folder123"))
+
+    assert [file["id"] for file in files] == ["1", "2"]
+
+
+def test_drive_list_files_in_folder_queries_only_the_folder_children(
+    requests_mock, drive_credential
+):
+    mock_token_refresh(requests_mock)
+    listing_mock = requests_mock.get(
+        "https://www.googleapis.com/drive/v3/files", json={"files": []}
+    )
+
+    list(drive_list_files_in_folder(folder_id="folder123"))
+
+    query = listing_mock.last_request.qs["q"][0]
+    assert "'folder123' in parents" in query
+    assert "trashed = false" in query
+
+
+def test_drive_helpers_fail_when_no_credential_has_drive_quota(admin_user):
+    stored_credential = GoogleCloudOAuthCredential.objects.create(
+        quota_limit_for_drive=0
+    )
+    GoogleCloudToken.objects.create(
+        oauth_credential=stored_credential, token="token", admin_user=admin_user
+    )
+
+    with pytest.raises(NoGoogleCloudQuotaLeftError):
+        drive_file_metadata(file_id="file123")
+
 
 def test_get_available_credentials(admin_user):
     stored_credential = GoogleCloudOAuthCredential.objects.create()
@@ -63,6 +176,42 @@ def test_function(*, credentials):
         assert used_quota.used_at == datetime.datetime.now(tz=datetime.timezone.utc)
 
 
+def test_count_quota_reuses_the_credentials_it_is_given(admin_user):
+    """A caller making several calls for one job keeps them on one account."""
+    already_in_use = GoogleCloudOAuthCredential.objects.create(client_id="in-use")
+    GoogleCloudToken.objects.create(
+        oauth_credential=already_in_use,
+        client_id="in-use",
+        token="token-in-use",
+        admin_user=admin_user,
+    )
+    # left with more quota, so it would win the lookup if one happened
+    idle = GoogleCloudOAuthCredential.objects.create(client_id="idle")
+    GoogleCloudToken.objects.create(
+        oauth_credential=idle,
+        client_id="idle",
+        token="token-idle",
+        admin_user=admin_user,
+    )
+    UsedRequestQuota.objects.create(credentials=idle, cost=1, service="youtube")
+
+    @count_quota("youtube", 1000)
+    def test_function(*, credentials):
+        return credentials
+
+    passed_in = get_available_credentials("youtube", 1000)
+    returned = test_function(credentials=passed_in)
+
+    assert returned is passed_in
+    # the quota is still charged, to the account that actually did the work
+    assert (
+        GoogleCloudOAuthCredential.objects.get_by_client_id(passed_in.client_id)
+        .usedrequestquota_set.filter(cost=1000)
+        .count()
+        == 1
+    )
+
+
 def test_count_quota_with_generator_function(admin_user):
     stored_credential = GoogleCloudOAuthCredential.objects.create()
     GoogleCloudToken.objects.create(