diff --git a/abr-testing/abr_testing/automation/google_drive_tool.py b/abr-testing/abr_testing/automation/google_drive_tool.py index d7a55ef8d22..5fefe2aaf5b 100644 --- a/abr-testing/abr_testing/automation/google_drive_tool.py +++ b/abr-testing/abr_testing/automation/google_drive_tool.py @@ -35,15 +35,23 @@ def __init__(self, credentials: Any, folder_name: str, email: str) -> None: print("Error! Get file: https://console.cloud.google.com/apis/credentials") sys.exit() - def list_folder(self, delete: Any = False) -> Set[str]: + def list_folder(self, delete: Any = False, folder: bool = False) -> Set[str]: """List folders and files in Google Drive.""" file_names = set() page_token: str = "" + basic_query = f"'{self.parent_folder}' in parents and trashed=false" + folder_query = ( + basic_query + " and mimeType='application/vnd.google-apps.folder'" + ) + if folder is True: + query = folder_query + else: + query = basic_query while True: results = ( self.drive_service.files() .list( - q=f"'{self.parent_folder}' in parents and trashed=false" + q=query if self.parent_folder else "" # type: ignore if self.parent_folder @@ -68,20 +76,29 @@ def list_folder(self, delete: Any = False) -> Set[str]: if not file_names: print("No folders or files found in Google Drive.") return file_names - + def create_folder(self, new_folder_name: str) -> str: """Create folder within defined folder.""" file_metadata = { "name": new_folder_name, "mimeType": "application/vnd.google-apps.folder", - "parents": self.parent_folder - } - file = self.drive_service.files().create(body=file_metadata, fields="id", supportsAllDrives="true").execute() - folder_id = file.get("id") - self.share_permissions(folder_id) - # SHARE FOLDER WITH EMAIL - print(f'Folder ID: "{file.get("id")}".') - return file.get("id") + "parents": [self.parent_folder], + } + list_of_current_folders = self.list_folder(folder=True) + if new_folder_name in list_of_current_folders: + print("Folder name already exists. Try new folder name or delete folder.") + return "" + else: + file = ( + self.drive_service.files() + .create(body=file_metadata, fields="id") # type: ignore + .execute() + ) + folder_id = file.get("id", "") + self.share_permissions(folder_id) + # SHARE FOLDER WITH EMAIL + print(f'Folder ID: "{file.get("id")}".') + return file.get("id", "") def delete_files(self, file_or_folder_id: str) -> None: """Delete a file or folder in Google Drive by ID.""" @@ -92,12 +109,12 @@ def delete_files(self, file_or_folder_id: str) -> None: print(f"Error deleting file/folder with ID: {file_or_folder_id}") print(f"Error details: {str(e)}") - def upload_file(self, file_path: str) -> str: + def upload_file(self, file_path: str, folder_name: str) -> str: """Upload file to Google Drive.""" file_metadata = { "name": os.path.basename(file_path), "mimeType": str(mimetypes.guess_type(file_path)[0]), - "parents": [self.parent_folder], + "parents": [folder_name], } media = MediaFileUpload(file_path, resumable=True) @@ -106,7 +123,12 @@ def upload_file(self, file_path: str) -> str: .create(body=file_metadata, media_body=media, fields="id") # type: ignore .execute() ) - return uploaded_file["id"] + uploaded_file_id = uploaded_file["id"] + try: + self.share_permissions(uploaded_file_id) + except googleapiclient.errors.HttpError: + print(f"File '{uploaded_file_id}' was not found after uploading.") + return uploaded_file_id def upload_missing_files(self, storage_directory: str) -> None: """Upload missing files to Google Drive.""" @@ -124,7 +146,7 @@ def upload_missing_files(self, storage_directory: str) -> None: uploaded_files = [] for file in missing_files: file_path = os.path.join(storage_directory, file) - uploaded_file_id = google_drive.upload_file(self, file_path) + uploaded_file_id = self.upload_file(file_path, self.parent_folder) uploaded_files.append( {"name": os.path.basename(file_path), "id": uploaded_file_id} ) diff --git a/abr-testing/abr_testing/data_collection/abr_robot_error.py b/abr-testing/abr_testing/data_collection/abr_robot_error.py index f5b70a0d271..a03bea034be 100644 --- a/abr-testing/abr_testing/data_collection/abr_robot_error.py +++ b/abr-testing/abr_testing/data_collection/abr_robot_error.py @@ -194,7 +194,7 @@ def get_error_info_from_robot( ) # WRITE ERRORED RUN TO GOOGLE SHEET error_run_log = os.path.join(error_folder_path, os.path.basename(run_log_file_path)) - google_drive.upload_file(error_run_log) + google_drive.upload_file(error_run_log, "1Cvej0eadFOTZr9ILRXJ0Wg65ymOtxL4m") run_id = os.path.basename(error_run_log).split("_")[1].split(".")[0] ( runs_and_robots, diff --git a/abr-testing/abr_testing/tools/test_gd_folder_upload.py b/abr-testing/abr_testing/tools/test_gd_folder_upload.py deleted file mode 100644 index 3faa3e560cd..00000000000 --- a/abr-testing/abr_testing/tools/test_gd_folder_upload.py +++ /dev/null @@ -1,36 +0,0 @@ -from abr_testing.automation import google_drive_tool -import argparse -import sys -import os - -if __name__ == "__main__": - """Get run logs.""" - parser = argparse.ArgumentParser(description="Pulls run logs from ABR robots.") - parser.add_argument( - "storage_directory", - metavar="STORAGE_DIRECTORY", - type=str, - nargs=1, - help="Path to long term storage directory for run logs.", - ) - parser.add_argument( - "folder_name", - metavar="FOLDER_NAME", - type=str, - nargs=1, - help="Google Drive folder name. Open desired folder and copy string after drive/folders/.", - ) - parser.add_argument( - "email", metavar="EMAIL", type=str, nargs=1, help="opentrons gmail." - ) - args = parser.parse_args() - storage_directory = args.storage_directory[0] - folder_name = args.folder_name[0] - email = args.email[0] - try: - credentials_path = os.path.join(storage_directory, "credentials.json") - except FileNotFoundError: - print(f"Add credentials.json file to: {storage_directory}.") - sys.exit() - google_drive = google_drive_tool.google_drive(credentials_path, folder_name, email) - google_drive.create_folder("test folder") \ No newline at end of file