from __future__ import annotations
from typing import List, Dict, Any
import re
from airflow.utils.context import Context
from airflow.exceptions import AirflowException
from operators.gdrive.gdrive_base import GDriveBaseOperator
[docs]
class GDriveListFilesOperator(GDriveBaseOperator):
"""
Recursively list *files* (not sub-folders) under a Drive folder.
Returns → List[dict] (id, name, mimeType)
"""
ui_color = "#FACADE"
template_fields = GDriveBaseOperator.template_fields + ("drive_link",)
_FOLDER_MIME = "application/vnd.google-apps.folder"
def __init__(
self,
*,
drive_link: str | None = None,
google_drive_id: str | None = None,
include_trashed: bool = False,
page_size: int = 1000,
**kwargs,
):
if not (drive_link or google_drive_id):
raise ValueError("Provide either drive_link or google_drive_id")
super().__init__(google_drive_id=google_drive_id, **kwargs)
self.drive_link = drive_link
self.include_trashed = include_trashed
self.page_size = page_size
# ------------- helpers ----------------
@staticmethod
def _extract_folder_id(link: str) -> str:
"""Extract the folder ID from a Google Drive folder URL."""
m = re.search(r"(?:folders/|id=)([a-zA-Z0-9_-]{10,})", link)
if not m:
raise ValueError(f"Cannot parse folder ID from link: {link}")
return m.group(1)
def _walk(self, folder_id: str, sink: list[Dict[str, Any]]):
"""Depth-first traversal, to get each file object"""
q = [
f"'{folder_id}' in parents",
"trashed = false" if not self.include_trashed else None,
]
query = " and ".join([p for p in q if p])
token = None
while True:
resp = (
self.drive.files()
.list(
q=query,
fields="nextPageToken, files(id, name, mimeType)",
pageSize=self.page_size,
pageToken=token,
supportsAllDrives=True,
includeItemsFromAllDrives=True,
corpora="allDrives",
)
.execute()
)
for obj in resp.get("files", []):
if obj["mimeType"] == self._FOLDER_MIME:
self._walk(obj["id"], sink)
else:
sink.append(obj)
token = resp.get("nextPageToken")
if not token:
break
# ------------- operator core ----------
[docs]
def execute(self, context: Context) -> List[Dict[str, str]]:
folder_id = self.google_drive_id
if not folder_id:
folder_id = self._extract_folder_id(self.drive_link)
files: list[dict] = []
self._walk(folder_id, files)
self.log.info("Found %s file(s) under folder %s", len(files), folder_id)
return files # → XCom
# ---------------------------------------------------------------------------#
# GDriveGetFileIdOperator - gdrive file url to file id w/ validation
# ---------------------------------------------------------------------------#
[docs]
class GDriveGetFileIdOperator(GDriveBaseOperator):
"""
Accept a public/share URL *or* a raw file-ID, verify the object exists,
and return a small metadata dict via XCom.
Example
-------
get_file_id = GDriveGetFileIdOperator(
task_id="get_file_id",
file_url="https://drive.google.com/file/d/123456789asdfce/view?usp=sharing",
# gcp_conn_id defaults to "google_cloud_default"
)
"""
ui_color = "#F1BBED"
template_fields = GDriveBaseOperator.template_fields + ("file_url", "file_id")
_FILE_ID_RE = re.compile(r"/d/([a-zA-Z0-9_-]{10,})|id=([a-zA-Z0-9_-]{10,})")
def __init__(
self,
*,
file_url: str | None = None,
file_id: str | None = None,
raise_if_missing: bool = True,
**kwargs,
):
if not (file_url or file_id):
raise ValueError("Provide either 'file_url' or 'file_id'")
super().__init__(**kwargs)
self.file_url = file_url
self.file_id = file_id
self.raise_if_missing = raise_if_missing
# ------------------------------------------------------------------#
# helpers
# ------------------------------------------------------------------#
@classmethod
def _extract_file_id(cls, url: str) -> str:
"""
Extract the file-ID from common Google-Drive share URLs.
"""
m = cls._FILE_ID_RE.search(url)
if not m:
raise ValueError(f"Cannot parse file ID from link: {url}")
# m.groups() returns a tuple like ('AbC123...', None) or (None, 'AbC123…')
try:
return next(g for g in m.groups() if g)
except StopIteration:
# Extremely defensive – regex matched but no capture.
raise ValueError(f"Regex matched, but no ID captured for link: {url}")
# ------------------------------------------------------------------#
# operator core
# ------------------------------------------------------------------#
[docs]
def execute(self, context: Context) -> Dict[str, Any]:
# 1. Resolve the file ID
file_id = self.file_id or self._extract_file_id(self.file_url)
# 2. Query Drive metadata
self.log.info("Validating existence of Drive file %s", file_id)
try:
metadata = (
self.drive.files()
.get(
fileId=file_id,
fields="id, name, mimeType, size, webViewLink",
supportsAllDrives=True,
)
.execute()
)
self.log.info(
"File found: %s (%s)",
metadata.get("name"),
metadata.get("mimeType"),
)
return metadata # → XCom
except Exception as err:
msg = f"File {file_id} not found or inaccessible: {err}"
if self.raise_if_missing:
raise AirflowException(msg) from err
self.log.warning(msg)
return {} # still JSON-serialisable