Copy_GDrive_to_GCS_Bucket.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Copy GDrive to GCS Bucket",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Provide `gdrive_path` per the following examples\n",
        "\n",
        "Examples:\n",
        "* `//My Drive/` instructs the script to copy everything in your Google Drive\n",
        "* `//DeptName/` instructs the script to copy everything in the DeptName Shared Google Drive\n",
        "* `//My Drive/path/to/files` instructs the script to copy everything in the `/path/to/files` directory of your  Google Drive"
      ],
      "metadata": {
        "id": "4GB722y7sEnr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Basic Google Drive info\n",
        "gdrive_path = '//My Drive/'"
      ],
      "metadata": {
        "id": "KX02mtTVpEdU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "A few words to consider when deciding which Bucket to copy files into...\n",
        "* It is recommended that a separate Bucket be used for each Google Drive\n",
        "* If deemed appropriate, you may use one Bucket for shared Google Drives\n"
      ],
      "metadata": {
        "id": "fE1vhD50rgND"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Basic Bucket information\n",
        "project_id = 'YOUR_GCP_PROJECT_ID' # e.g. bborie-sandbox\n",
        "bucket_name = 'YOUR_BUCKET_NAME_HERE' # e.g. bborie-sandbox"
      ],
      "metadata": {
        "id": "cyxs7tlaq5vq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Basic Script info\n",
        "num_workers = 8\n",
        "dry_run = False"
      ],
      "metadata": {
        "id": "orX4dxsGySF5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Do not edit anything after this Text block. Just execute"
      ],
      "metadata": {
        "id": "WCW3_178swqx"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eEMbFfJfnFH3"
      },
      "outputs": [],
      "source": [
        "from collections import deque\n",
        "import io\n",
        "import logging\n",
        "import os\n",
        "import os.path as osp\n",
        "import queue\n",
        "import tempfile\n",
        "import threading\n",
        "import time\n",
        "from types import SimpleNamespace\n",
        "\n",
        "from googleapiclient.discovery import build\n",
        "from googleapiclient.http import MediaIoBaseDownload\n",
        "import google.auth\n",
        "from google.cloud import storage\n",
        "from google.colab import auth, drive\n",
        "import requests"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Authenticate into GCP to access your Bucket\n",
        "auth.authenticate_user()"
      ],
      "metadata": {
        "id": "riEDW8OWqz2a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "os.environ['GOOGLE_CLOUD_PROJECT'] = project_id\n",
        "\n",
        "CONNECT_TIMEOUT = 10 # in seconds\n",
        "READ_TIMEOUT = 900 # 15 minutes\n",
        "TIMEOUT = (CONNECT_TIMEOUT, READ_TIMEOUT)\n",
        "\n",
        "CHUNK_SIZE = 10 * 1024 * 1024 # 10 MB upload chunks\n",
        "\n",
        "FOLDER_MIME_TYPE = 'application/vnd.google-apps.folder'\n",
        "MIME_TYPE_MAP = {\n",
        "    'application/vnd.google-apps.document': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',\n",
        "    'application/vnd.google-apps.drawing': 'image/png',\n",
        "    'application/vnd.google-apps.presentation': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',\n",
        "    'application/vnd.google-apps.spreadsheet': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',\n",
        "    'application/vnd.google-apps.script': 'application/vnd.google-apps.script+json',\n",
        "}\n",
        "\n",
        "file_queue = queue.Queue()\n",
        "LOCK = threading.Lock()\n",
        "\n",
        "logging.basicConfig()\n",
        "logger = logging.getLogger()\n",
        "logger.setLevel(logging.INFO)"
      ],
      "metadata": {
        "id": "sOop50hivfwS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class UnsupportedMimeType(Exception):\n",
        "    pass\n",
        "\n",
        "def clean_value(v):\n",
        "    return v.replace(\"'\", \"\\\\'\")\n",
        "\n",
        "def iterfiles(service, name=None, is_folder=None, parent=None, drive_id=None, order_by='folder,name,createdTime'):\n",
        "    q = ['trashed=false']\n",
        "    if name is not None:\n",
        "        q.append(f\"name='{clean_value(name)}'\")\n",
        "    if is_folder is not None:\n",
        "        q.append(f\"mimeType {'=' if is_folder else '!='} '{FOLDER_MIME_TYPE}'\")\n",
        "    if parent is not None:\n",
        "        q.append(f\"'{clean_value(parent)}' in parents\")\n",
        "\n",
        "    params = {'pageToken': None, 'orderBy': order_by}\n",
        "    if q:\n",
        "        params['q'] = ' and '.join(q)\n",
        "\n",
        "    if drive_id:\n",
        "        params['driveId'] = drive_id\n",
        "        params['corpora'] = 'drive'\n",
        "        params['includeItemsFromAllDrives'] = True\n",
        "        params['supportsAllDrives'] = True\n",
        "\n",
        "    while True:\n",
        "        response = service.files().list(**params).execute()\n",
        "        for f in response['files']:\n",
        "            yield f\n",
        "        try:\n",
        "            params['pageToken'] = response['nextPageToken']\n",
        "        except KeyError:\n",
        "            return\n",
        "\n",
        "def walk(service, top='root', by_name=False, drive_id=None):\n",
        "    if by_name:\n",
        "        top, = iterfiles(service, name=top, is_folder=True, drive_id=drive_id)\n",
        "    else:\n",
        "        params = dict(fileId=top)\n",
        "        if drive_id:\n",
        "            params['supportsAllDrives'] = True\n",
        "        top = service.files().get(**params).execute()\n",
        "        if top['mimeType'] != FOLDER_MIME_TYPE:\n",
        "            raise ValueError(f'not a folder: {top!r}')\n",
        "\n",
        "    stack = [((top['name'],), top)]\n",
        "    while stack:\n",
        "        path, top = stack.pop()\n",
        "\n",
        "        dirs, files = is_file = [], []\n",
        "        for f in iterfiles(service, parent=top['id'], drive_id=drive_id):\n",
        "            is_file[f['mimeType'] != FOLDER_MIME_TYPE].append(f)\n",
        "\n",
        "        yield path, top, dirs, files\n",
        "\n",
        "        if dirs:\n",
        "            stack.extend((path + (d['name'],), d) for d in reversed(dirs))\n",
        "\n",
        "def get_drive_service():\n",
        "    return build('drive', version='v3', cache_discovery=False)\n",
        "\n",
        "def get_storage_client(project):\n",
        "    return storage.Client(project=project)\n",
        "\n",
        "def init_workers(num_workers=8, thread_kwargs=None):\n",
        "    thread_kwargs = thread_kwargs or {}\n",
        "    file_workers = [\n",
        "        threading.Thread(target=copy_file, daemon=True, kwargs=thread_kwargs)\n",
        "        for x in range(num_workers)\n",
        "    ]\n",
        "    list(map(lambda x: x.start(), file_workers))\n",
        "\n",
        "    return file_workers\n",
        "\n",
        "def copy_file(project, bucket, dry_run=False):\n",
        "    with LOCK:\n",
        "        creds = google.auth.default()[0]\n",
        "        service = get_drive_service()\n",
        "        client = get_storage_client(project)\n",
        "        bucket = client.bucket(bucket)\n",
        "\n",
        "    while True:\n",
        "        try:\n",
        "            f = file_queue.get()\n",
        "            if f is None:\n",
        "                break\n",
        "\n",
        "            _copy_file(f, service, creds, bucket, dry_run=dry_run)\n",
        "        except Exception as e:\n",
        "            log('error', f'Error copying: {f.drivePath} => gs://{bucket.name}/{f.bucketPath}')\n",
        "            log('error', e)\n",
        "        finally:\n",
        "            file_queue.task_done()\n",
        "\n",
        "def _copy_file(f, service, creds, bucket, dry_run=False):\n",
        "    with tempfile.NamedTemporaryFile() as fh:\n",
        "        if dry_run:\n",
        "            log('info', f'Copied: {f.id} {f.drivePath} => gs://{bucket.name}/{f.bucketPath}')\n",
        "            return\n",
        "\n",
        "        #\n",
        "        # download\n",
        "        #\n",
        "\n",
        "        file = io.FileIO(fh.name, 'w+')\n",
        "        if '.google-apps.' in f.mimeType:\n",
        "            mime_type = MIME_TYPE_MAP.get(f.mimeType)\n",
        "            if not mime_type:\n",
        "                raise UnsupportedMimeType()\n",
        "\n",
        "            f.mimeType = mime_type\n",
        "\n",
        "            response = service.files().get(fileId=f.id, supportsAllDrives=True, fields='*').execute()\n",
        "            response = requests.get(\n",
        "                response['exportLinks'][mime_type],\n",
        "                headers={\n",
        "                    'Authorization': 'bearer {creds.access_token}'\n",
        "                },\n",
        "                stream=True\n",
        "            )\n",
        "            response.raise_for_status()\n",
        "            for chunk in response.iter_content(chunk_size=CHUNK_SIZE):\n",
        "                fh.write(chunk)\n",
        "        else:\n",
        "            req = service.files().get_media(fileId=f.id)\n",
        "            downloader = MediaIoBaseDownload(file, req)\n",
        "            done = False\n",
        "            while not done:\n",
        "                _, done = downloader.next_chunk()\n",
        "\n",
        "        #\n",
        "        # upload\n",
        "        #\n",
        "\n",
        "        blob = bucket.blob(f.bucketPath, chunk_size=CHUNK_SIZE)\n",
        "        blob.upload_from_filename(fh.name)\n",
        "        blob.content_type = f.mimeType\n",
        "        blob.patch()\n",
        "\n",
        "        log('info', f'Copied: {f.id} {f.drivePath} => gs://{bucket.name}/{f.bucketPath}')\n",
        "\n",
        "def log(level, msg):\n",
        "    fn = getattr(logger, level)\n",
        "    message = f'({threading.get_ident()}) {msg}'\n",
        "    fn(message)\n",
        "\n",
        "def copy_from_gdrive_to_bucket(project, drive, bucket, workers, dry_run):\n",
        "    if not drive.startswith('//'):\n",
        "        drive = '//' + osp.join('My Drive', (drive if drive[0] != '/' else drive[1:]))\n",
        "        logger.info(f'Assuming \"My Drive\": {drive}')\n",
        "\n",
        "    drive_name, drive_prefix = drive[2:].split('/', 1)\n",
        "    if drive_prefix.endswith('/'):\n",
        "        drive_prefix = drive_prefix[:-1]\n",
        "    drive_parts = drive_prefix.split('/')\n",
        "\n",
        "    drive_id = None\n",
        "    drive_parent = 'root'\n",
        "\n",
        "    service = get_drive_service()\n",
        "    creds = google.auth.default()[0]\n",
        "\n",
        "    file_workers = init_workers(\n",
        "        workers,\n",
        "        thread_kwargs={\n",
        "            'project': project,\n",
        "            'bucket': bucket,\n",
        "            'dry_run': dry_run,\n",
        "        }\n",
        "    )\n",
        "\n",
        "    if drive_name.strip().lower() != 'my drive':\n",
        "        # go look up driveId\n",
        "        response = service.drives().list(q=f\"name='{clean_value(drive_name)}'\").execute()\n",
        "        drive_id = next(\n",
        "            drive_dict['id']\n",
        "            for drive_dict in response['drives']\n",
        "        )\n",
        "        drive_parent = None\n",
        "\n",
        "    for part in drive_parts:\n",
        "        for f in iterfiles(service, name=part, parent=drive_parent, is_folder=True, drive_id=drive_id):\n",
        "            drive_parent = f['id']\n",
        "\n",
        "    dq = deque([], 1000)\n",
        "    for _path, root, dirs, files in walk(service, top=drive_parent, by_name=False, drive_id=drive_id):\n",
        "        for f in files:\n",
        "            f = SimpleNamespace(**f)\n",
        "            f.drivePath = osp.join(drive, *_path[1:], f.name)\n",
        "            bucket_path = f.bucketPath = osp.join(drive[2:], *_path[1:], f.name)\n",
        "            if bucket_path in dq:\n",
        "                count = dq.count(bucket_path)\n",
        "                f.bucketPath += f' ({count})'\n",
        "            dq.append(bucket_path)\n",
        "            file_queue.put(f)\n",
        "\n",
        "    while True:\n",
        "        qsize = file_queue.qsize()\n",
        "        logger.info(f'Estimate to process: {qsize}')\n",
        "        if qsize < 5:\n",
        "            break\n",
        "        time.sleep(10)\n",
        "\n",
        "    # gracefully shut down workers\n",
        "    for x in range(len(file_workers)):\n",
        "        file_queue.put(None)\n",
        "\n",
        "    file_queue.join()\n",
        "\n",
        "    logger.info('All done! Bye...')"
      ],
      "metadata": {
        "id": "xtmxuyN8ykfa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "copy_from_gdrive_to_bucket(project_id, gdrive_path, bucket_name, num_workers, dry_run)"
      ],
      "metadata": {
        "id": "zs8a_Wb60LV0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "cZNZ20qD8eO1"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}