"""
Google Drive Upload Operator.
Uploads files to Google Drive folders with support for both local files
and in-memory content strings.
"""
from __future__ import annotations
import os
import tempfile
from typing import Sequence
from airflow.utils.context import Context
from googleapiclient.http import MediaFileUpload
from operators.gdrive.gdrive_base import GDriveBaseOperator
[docs]
class GDriveUploadFileOperator(GDriveBaseOperator):
"""
Upload a file to a Google Drive folder.
Supports uploading from either a local file path or in-memory string content.
Can optionally overwrite existing files with the same name.
:param google_drive_id: Target Google Drive folder ID
:param file_path: Local file path to upload (mutually exclusive with file_content)
:param file_name: Name for the uploaded file in Google Drive
:param file_content: String content to upload (mutually exclusive with file_path)
:param mimetype: MIME type of the file (default: auto-detect or text/plain)
:param overwrite: If True, replace existing file with same name (default: True)
Returns dict with keys: id, name, webViewLink
"""
template_fields: Sequence[str] = GDriveBaseOperator.template_fields + (
"file_path",
"file_name",
"file_content",
)
ui_color = "#A6E6A6"
def __init__(
self,
*,
google_drive_id: str,
file_name: str,
file_path: str | None = None,
file_content: str | None = None,
mimetype: str | None = None,
overwrite: bool = True,
**kwargs,
):
super().__init__(google_drive_id=google_drive_id, **kwargs)
self.file_path = file_path
self.file_name = file_name
self.file_content = file_content
self.mimetype = mimetype
self.overwrite = overwrite
def _find_existing_file(self, filename: str) -> str | None:
"""Find existing file by name in the target folder."""
query = (
f"'{self.google_drive_id}' in parents "
f"and name = '{filename}' "
f"and trashed = false"
)
result = (
self.drive.files()
.list(
q=query,
fields="files(id, name)",
supportsAllDrives=True,
includeItemsFromAllDrives=True,
corpora="allDrives",
)
.execute()
)
files = result.get("files", [])
return files[0]["id"] if files else None
def _determine_mimetype(self, file_path: str) -> str:
"""Determine MIME type based on file extension."""
ext = os.path.splitext(file_path)[1].lower()
mime_map = {
".txt": "text/plain",
".md": "text/markdown",
".csv": "text/csv",
".json": "application/json",
".xml": "application/xml",
".pdf": "application/pdf",
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
".zip": "application/zip",
}
return mime_map.get(ext, "application/octet-stream")
def _upload_file(self, local_path: str) -> dict:
"""Upload file to Google Drive."""
mimetype = self.mimetype or self._determine_mimetype(local_path)
media = MediaFileUpload(local_path, mimetype=mimetype, resumable=True)
existing_file_id = None
if self.overwrite:
existing_file_id = self._find_existing_file(self.file_name)
if existing_file_id:
self.log.info(
"Overwriting existing file '%s' (id=%s)",
self.file_name,
existing_file_id,
)
uploaded = (
self.drive.files()
.update(
fileId=existing_file_id,
media_body=media,
fields="id, name, webViewLink",
supportsAllDrives=True,
)
.execute()
)
else:
self.log.info("Creating new file '%s'", self.file_name)
file_metadata = {
"name": self.file_name,
"parents": [self.google_drive_id],
}
uploaded = (
self.drive.files()
.create(
body=file_metadata,
media_body=media,
fields="id, name, webViewLink",
supportsAllDrives=True,
)
.execute()
)
self.log.info(
"Uploaded %s (id=%s) - %s",
uploaded.get("name"),
uploaded.get("id"),
uploaded.get("webViewLink"),
)
return uploaded
[docs]
def execute(self, context: Context) -> dict:
if not self.file_path and not self.file_content:
raise ValueError("Either file_path or file_content must be provided")
if self.file_path and self.file_content:
raise ValueError("Only one of file_path or file_content should be provided")
if self.file_content:
with tempfile.NamedTemporaryFile(
mode="w",
suffix=os.path.splitext(self.file_name)[1],
delete=False,
encoding="utf-8",
) as tmp:
tmp.write(self.file_content)
tmp_path = tmp.name
try:
return self._upload_file(tmp_path)
finally:
os.unlink(tmp_path)
else:
return self._upload_file(self.file_path)