drive_downloader.py

import os
import json
import io
import zipfile

from google.oauth2.service_account import Credentials
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseDownload


class ZipDownloader:
    def __init__(self, service_account_json: str):
        self.creds = Credentials.from_service_account_info(json.loads(service_account_json))
        self.service = build("drive", "v3", credentials=self.creds)

    def download_zip_from_drive(self, file_id: str, output_path: str) -> str:
        """
        Downloads a ZIP file from Google Drive and saves it locally.
        Returns the path to the downloaded ZIP.
        """
        request = self.service.files().get_media(fileId=file_id)
        local_zip_path = os.path.join(output_path, "downloaded.zip")

        os.makedirs(output_path, exist_ok=True)
        fh = io.FileIO(local_zip_path, 'wb')
        downloader = MediaIoBaseDownload(fh, request)

        print(f"⬇️ Downloading ZIP file from Drive ID: {file_id}")
        done = False
        while not done:
            status, done = downloader.next_chunk()
            print(f"   ⏬ Progress: {int(status.progress() * 100)}%")

        print(f"✅ ZIP downloaded to: {local_zip_path}")
        return local_zip_path

    def unzip(self, zip_path: str, extract_to: str):
        """
        Unzips the downloaded ZIP file to a specified directory.
        """
        print(f"📂 Extracting ZIP: {zip_path} -> {extract_to}")
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(extract_to)
        print("✅ Extraction complete.")