Source code for operators.gdrive.gdrive_folder

"""Google Drive Folder Creation Operator."""

from __future__ import annotations

from typing import Sequence

from airflow.utils.context import Context

from operators.gdrive.gdrive_base import GDriveBaseOperator




class GDriveCreateFolderOperator(GDriveBaseOperator):
    """
    Create a folder in Google Drive (or return existing if it already exists).

    :param folder_name: Name of the folder to create
    :param google_drive_id: Parent folder ID (inherited from base class)

    Returns dict with keys: id, name, webViewLink
    """

    template_fields: Sequence[str] = GDriveBaseOperator.template_fields + (
        "folder_name",
    )
    ui_color = "#A6E6A6"

    def __init__(
        self,
        *,
        folder_name: str,
        google_drive_id: str,
        **kwargs,
    ):
        super().__init__(google_drive_id=google_drive_id, **kwargs)
        self.folder_name = folder_name

    def _find_existing_folder(self) -> dict | None:
        """Find existing folder by name in parent."""
        query = (
            f"'{self.google_drive_id}' in parents "
            f"and name = '{self.folder_name}' "
            f"and mimeType = 'application/vnd.google-apps.folder' "
            f"and trashed = false"
        )
        result = (
            self.drive.files()
            .list(
                q=query,
                fields="files(id, name, webViewLink)",
                supportsAllDrives=True,
                includeItemsFromAllDrives=True,
                corpora="allDrives",
            )
            .execute()
        )
        files = result.get("files", [])
        return files[0] if files else None



    def execute(self, context: Context) -> dict:
        existing = self._find_existing_folder()
        if existing:
            self.log.info(
                "Folder '%s' already exists (id=%s)",
                self.folder_name,
                existing["id"],
            )
            return existing

        self.log.info(
            "Creating folder '%s' in parent '%s'",
            self.folder_name,
            self.google_drive_id,
        )
        file_metadata = {
            "name": self.folder_name,
            "mimeType": "application/vnd.google-apps.folder",
            "parents": [self.google_drive_id],
        }
        created = (
            self.drive.files()
            .create(
                body=file_metadata,
                fields="id, name, webViewLink",
                supportsAllDrives=True,
            )
            .execute()
        )
        self.log.info(
            "Created folder '%s' (id=%s) - %s",
            created.get("name"),
            created.get("id"),
            created.get("webViewLink"),
        )
        return created