google_drive_create_update_files_through_streamlit.py

%%writefile app.py
from __future__ import print_function

import streamlit as st
from googleapiclient.errors import HttpError
from googleapiclient.http import MediaFileUpload
from googleapiclient.discovery import build
from oauth2client.service_account import ServiceAccountCredentials
import os


class GoogleDriveService:
    def __init__(self):
        self._SCOPES=['https://www.googleapis.com/auth/drive']

        _base_path = os.path.dirname('/content/')
        _credential_path=os.path.join(_base_path, 'credential.json')
        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = _credential_path

    def build(self):
        creds = ServiceAccountCredentials.from_json_keyfile_name(os.getenv("GOOGLE_APPLICATION_CREDENTIALS"), self._SCOPES)
        service = build('drive', 'v3', credentials=creds)

        return service

def create_basic(fid, path):
    'fid: the folder id of drive under which you would create a new file.'
    try:
        file_metadata = {'name':'global_model.h5', 'parents': [fid]}
        media = MediaFileUpload(path,
                                mimetype='application/x-hdf5')
        # pylint: disable=maybe-no-member
        file = g_drive_service.files().create(body=file_metadata, media_body=media,
                                      fields='id').execute()
        print(F'File ID: {file.get("id")}')

    except HttpError as error:
        print(F'An error occurred: {error}')
        file = None

    return file.get('id')

def update_basic(fid, path):
    'fid: the id of the drive file you want to update'
    try:
        media = MediaFileUpload(path, mimetype='application/x-hdf5')
        # pylint: disable=maybe-no-member
        file_metadata = {'name':'global_model.h5'}
        file = g_drive_service.files().update(fileId=fid,body=file_metadata, media_body=media).execute()
        print(F'File ID: {file.get("id")}')

    except HttpError as error:
        print(F'An error occurred: {error}')
        file = None

    return file.get('id')

selected_fields="files(id,name,webViewLink)"
g_drive_service=GoogleDriveService().build()
list_file=g_drive_service.files().list(fields=selected_fields).execute()
print(list_file)

folder_id='1pKC2WaM0UxeBIS_d05kP5TC2wKavuIo9'

test_create = st.button('Create file')


if test_create:
    path = "/content/global_model.h5" # upload an example .h5 file and find its path in colab directory
    st.write('testing upload from streamlit') #displayed when the button is clicked
    update_id=create_basic(folder_id, path)
    st.write(update_id)
    st.write('testing updating from streamlit')
    update_basic(update_id, path)