Source code for operators.gdrive.gdrive_base

from __future__ import annotations

from airflow.models.baseoperator import BaseOperator
from airflow.hooks.base import BaseHook
from airflow.utils.context import Context
from typing import Sequence
from config import Config
import json

from google.oauth2 import service_account
from googleapiclient.discovery import build




class GDriveBaseOperator(BaseOperator):
    """
    Parent class for every Google-Drive-related operator.
    """

    template_fields: Sequence[str] = (
        "google_drive_id",
        "gcp_conn_id",
    )
    ui_color = "#A6E6A6"

    def __init__(
        self,
        *,
        google_drive_id: str | None = None,
        gcp_conn_id: str = Config.GDRIVE_SERVICE_ACCOUNT,
        scopes: Sequence[str] | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.google_drive_id = google_drive_id
        self.gcp_conn_id = gcp_conn_id
        self.scopes = scopes or [
            "https://www.googleapis.com/auth/drive",
            "https://www.googleapis.com/auth/spreadsheets",
        ]
        self._drive_service = None  # lazy-loaded

    # ------------------------------------------------------------------#
    # Helpers available to subclasses
    # ------------------------------------------------------------------#
    def _get_credentials(self):
        conn = BaseHook.get_connection(self.gcp_conn_id)
        cred_dict = json.loads(conn.extra)  # service-account JSON
        return service_account.Credentials.from_service_account_info(
            cred_dict, scopes=self.scopes
        )

    @property
    def drive(self):
        """googleapiclient drive-v3 service, cached for this task-run."""
        if self._drive_service is None:
            self._drive_service = build(
                "drive", "v3", credentials=self._get_credentials()
            )
        return self._drive_service

    # The base class does **nothing** in execute; concrete operators override.


    def execute(self, context: Context):
        raise NotImplementedError("Subclasses must implement execute()")