From 7930a196c08e95a1f6417e325d170f96ed49ea49 Mon Sep 17 00:00:00 2001 From: GDjkhp Date: Tue, 17 Dec 2024 18:50:56 +0800 Subject: [PATCH] changed to aiogoogle --- api_gdrive.py | 251 +++++++++++++++++++++++++---------------------- deeznuts.py | 30 +++--- requirements.txt | 2 +- ytdlp_.py | 38 ++++--- 4 files changed, 174 insertions(+), 147 deletions(-) diff --git a/api_gdrive.py b/api_gdrive.py index ed2be0d..e81579e 100644 --- a/api_gdrive.py +++ b/api_gdrive.py @@ -1,25 +1,23 @@ import os import glob +import json import mimetypes -from google.oauth2.credentials import Credentials -from googleapiclient.discovery import build -from googleapiclient.http import MediaFileUpload +from aiogoogle import Aiogoogle -class DriveUploader: +class AsyncDriveUploader: def __init__(self, credentials_path): """ - Initialize Google Drive service + Initialize Aiogoogle with user credentials Args: credentials_path (str): Path to OAuth 2.0 credentials file """ - self.creds = Credentials.from_authorized_user_file( - credentials_path, - ['https://www.googleapis.com/auth/drive.file'] - ) - self.drive_service = build('drive', 'v3', credentials=self.creds) + with open(credentials_path, "r", encoding="utf-8") as json_file: + creds_data = json.load(json_file) + self.user = {"access_token": creds_data['token'], "refresh_token": creds_data['refresh_token']} + self.client = {"client_id": creds_data['client_id'], "client_secret": creds_data['client_secret'], "scopes": creds_data['scopes']} - def get_or_create_folder(self, folder_name, parent_folder_id=None): + async def get_or_create_folder(self, folder_name, parent_folder_id=None): """ Find or create a folder in Google Drive. @@ -30,40 +28,48 @@ def get_or_create_folder(self, folder_name, parent_folder_id=None): Returns: str: Folder ID """ - # Build query to find folder - query = [f"name='{folder_name}'", "mimeType='application/vnd.google-apps.folder'"] - if parent_folder_id: - query.append(f"'{parent_folder_id}' in parents") - - results = self.drive_service.files().list( - q=' and '.join(query), - spaces='drive' - ).execute() - - folders = results.get('files', []) - - # If folder exists, return its ID - if folders: - return folders[0]['id'] - - # If folder doesn't exist, create it - folder_metadata = { - 'name': folder_name, - 'mimeType': 'application/vnd.google-apps.folder' - } - - # Add parent if specified - if parent_folder_id: - folder_metadata['parents'] = [parent_folder_id] - - folder = self.drive_service.files().create( - body=folder_metadata, - fields='id' - ).execute() - - return folder['id'] + async with Aiogoogle(user_creds=self.user, client_creds=self.client) as aiogoogle: + drive_v3 = await aiogoogle.discover('drive', 'v3') + + # Build query to find folder + query = [f"name='{folder_name}'", "mimeType='application/vnd.google-apps.folder'"] + if parent_folder_id: + query.append(f"'{parent_folder_id}' in parents") + + # Search for existing folder + results = await aiogoogle.as_user( + drive_v3.files.list( + q=' and '.join(query), + spaces='drive' + ) + ) + + folders = results.get('files', []) + + # If folder exists, return its ID + if folders: + return folders[0]['id'] + + # If folder doesn't exist, create it + folder_metadata = { + 'name': folder_name, + 'mimeType': 'application/vnd.google-apps.folder' + } + + # Add parent if specified + if parent_folder_id: + folder_metadata['parents'] = [parent_folder_id] + + folder = await aiogoogle.as_user( + drive_v3.files.create( + json=folder_metadata, + fields='id' + ) + ) + + return folder['id'] - def make_public_with_link(self, file_or_folder_id): + async def make_public_with_link(self, file_or_folder_id): """ Make a file or folder publicly accessible and get a shareable link. @@ -73,26 +79,33 @@ def make_public_with_link(self, file_or_folder_id): Returns: str: Public sharing link """ - try: - # Create a public permission - self.drive_service.permissions().create( - fileId=file_or_folder_id, - body={'type': 'anyone', 'role': 'reader'} - ).execute() + async with Aiogoogle(user_creds=self.user, client_creds=self.client) as aiogoogle: + drive_v3 = await aiogoogle.discover('drive', 'v3') - # Get the file/folder details to retrieve the web view link - file_metadata = self.drive_service.files().get( - fileId=file_or_folder_id, - fields='webViewLink' - ).execute() + try: + # Create a public permission + await aiogoogle.as_user( + drive_v3.permissions.create( + fileId=file_or_folder_id, + json={'type': 'anyone', 'role': 'reader'} + ) + ) + + # Get the file/folder details to retrieve the web view link + file_metadata = await aiogoogle.as_user( + drive_v3.files.get( + fileId=file_or_folder_id, + fields='webViewLink' + ) + ) + + return file_metadata.get('webViewLink', None) - return file_metadata.get('webViewLink', None) - - except Exception as e: - print(f"Error making file/folder public: {e}") - return None + except Exception as e: + print(f"Error making file/folder public: {e}") + return None - def batch_upload(self, paths, folder_in_drive='Uploaded', make_public=False, recursive=False): + async def batch_upload(self, paths, folder_in_drive='Uploaded', make_public=False, recursive=False): """ Batch upload multiple files and folders. @@ -106,7 +119,7 @@ def batch_upload(self, paths, folder_in_drive='Uploaded', make_public=False, rec list: List of upload results """ # Get or create the root folder in Drive - root_folder_id = self.get_or_create_folder(folder_in_drive) + root_folder_id = await self.get_or_create_folder(folder_in_drive) # Results storage upload_results = [] @@ -126,7 +139,7 @@ def batch_upload(self, paths, folder_in_drive='Uploaded', make_public=False, rec # Determine upload method based on path type if os.path.isfile(matched_path): # Upload single file - result = self._upload_single_file( + result = await self._upload_single_file( matched_path, root_folder_id, make_public @@ -136,14 +149,14 @@ def batch_upload(self, paths, folder_in_drive='Uploaded', make_public=False, rec elif os.path.isdir(matched_path): # Upload folder if recursive: - result = self._upload_folder( + result = await self._upload_folder( matched_path, root_folder_id, make_public ) else: # Non-recursive folder upload (only files in root) - result = self._upload_non_recursive_folder( + result = await self._upload_non_recursive_folder( matched_path, root_folder_id, make_public @@ -155,7 +168,7 @@ def batch_upload(self, paths, folder_in_drive='Uploaded', make_public=False, rec return upload_results - def _upload_single_file(self, file_path, parent_folder_id=None, make_public=False): + async def _upload_single_file(self, file_path, parent_folder_id=None, make_public=False): """ Upload a single file to Google Drive. @@ -167,54 +180,52 @@ def _upload_single_file(self, file_path, parent_folder_id=None, make_public=Fals Returns: dict: Information about uploaded file """ - # Get file size - file_size = os.path.getsize(file_path) + async with Aiogoogle(user_creds=self.user, client_creds=self.client) as aiogoogle: + drive_v3 = await aiogoogle.discover('drive', 'v3') + + # Get file size + file_size = os.path.getsize(file_path) - # Prepare file metadata - file_metadata = { - 'name': os.path.basename(file_path) - } - - # Add parent folder if specified - if parent_folder_id: - file_metadata['parents'] = [parent_folder_id] - - # Detect MIME type - mime_type, _ = mimetypes.guess_type(file_path) - if mime_type is None: - mime_type = 'application/octet-stream' - - # Create media upload - media = MediaFileUpload( - file_path, - mimetype=mime_type, - resumable=True - ) - - # Upload the file - print(file_path) - file = self.drive_service.files().create( - body=file_metadata, - media_body=media, - fields='id,webViewLink' - ).execute() - - # Make file public if requested - public_link = None - if make_public: - public_link = self.make_public_with_link(file['id']) - - return { - 'type': 'file', - 'name': os.path.basename(file_path), - 'full_path': file_path, - 'id': file['id'], - 'size_bytes': file_size, - 'size_human_readable': self._format_file_size(file_size), - 'link': public_link or file.get('webViewLink') - } + # Prepare file metadata + file_metadata = { + 'name': os.path.basename(file_path) + } + + # Add parent folder if specified + if parent_folder_id: + file_metadata['parents'] = [parent_folder_id] + + # Detect MIME type + mime_type, _ = mimetypes.guess_type(file_path) + if mime_type is None: + mime_type = 'application/octet-stream' + + # Upload the file + print(file_path) + file = await aiogoogle.as_user( + drive_v3.files.create( + json=file_metadata, + upload_file=file_path, + fields='id,webViewLink' + ) + ) + + # Make file public if requested + public_link = None + if make_public: + public_link = await self.make_public_with_link(file['id']) + + return { + 'type': 'file', + 'name': os.path.basename(file_path), + 'full_path': file_path, + 'id': file['id'], + 'size_bytes': file_size, + 'size_human_readable': self._format_file_size(file_size), + 'link': public_link or file.get('webViewLink') + } - def _upload_folder(self, local_path, parent_folder_id=None, make_public=False): + async def _upload_folder(self, local_path, parent_folder_id=None, make_public=False): """ Recursively upload a local folder to Google Drive. @@ -230,7 +241,7 @@ def _upload_folder(self, local_path, parent_folder_id=None, make_public=False): folder_name = os.path.basename(local_path) # Create the folder in Google Drive - current_folder_id = self.get_or_create_folder( + current_folder_id = await self.get_or_create_folder( folder_name, parent_folder_id ) @@ -238,7 +249,7 @@ def _upload_folder(self, local_path, parent_folder_id=None, make_public=False): # Make folder public if requested public_link = None if make_public: - public_link = self.make_public_with_link(current_folder_id) + public_link = await self.make_public_with_link(current_folder_id) # Track uploaded files uploaded_files = [] @@ -250,7 +261,7 @@ def _upload_folder(self, local_path, parent_folder_id=None, make_public=False): # Recursively upload files and subdirectories if os.path.isfile(local_item_path): - file_result = self._upload_single_file( + file_result = await self._upload_single_file( local_item_path, current_folder_id, make_public @@ -258,7 +269,7 @@ def _upload_folder(self, local_path, parent_folder_id=None, make_public=False): total_size += file_result.get('size_bytes', 0) uploaded_files.append(file_result) elif os.path.isdir(local_item_path): - subfolder_result = self._upload_folder( + subfolder_result = await self._upload_folder( local_item_path, current_folder_id, make_public @@ -277,7 +288,7 @@ def _upload_folder(self, local_path, parent_folder_id=None, make_public=False): 'files': uploaded_files } - def _upload_non_recursive_folder(self, local_path, parent_folder_id=None, make_public=False): + async def _upload_non_recursive_folder(self, local_path, parent_folder_id=None, make_public=False): """ Upload only files in the root of a folder (non-recursive). @@ -293,7 +304,7 @@ def _upload_non_recursive_folder(self, local_path, parent_folder_id=None, make_p folder_name = os.path.basename(local_path) # Create the folder in Google Drive - current_folder_id = self.get_or_create_folder( + current_folder_id = await self.get_or_create_folder( folder_name, parent_folder_id ) @@ -301,7 +312,7 @@ def _upload_non_recursive_folder(self, local_path, parent_folder_id=None, make_p # Make folder public if requested public_link = None if make_public: - public_link = self.make_public_with_link(current_folder_id) + public_link = await self.make_public_with_link(current_folder_id) # Track uploaded files uploaded_files = [] @@ -313,7 +324,7 @@ def _upload_non_recursive_folder(self, local_path, parent_folder_id=None, make_p # Upload only files, skip subdirectories if os.path.isfile(local_item_path): - file_result = self._upload_single_file( + file_result = await self._upload_single_file( local_item_path, current_folder_id, make_public diff --git a/deeznuts.py b/deeznuts.py index 1a9d6e4..6ab0ded 100644 --- a/deeznuts.py +++ b/deeznuts.py @@ -6,8 +6,7 @@ from discord.ext import commands from discord import app_commands import discord -from concurrent.futures import ThreadPoolExecutor # new hack -from api_gdrive import DriveUploader +from api_gdrive import AsyncDriveUploader from util_discord import description_helper, command_check, check_if_not_owner from util_database import myclient mycol = myclient["utils"]["cant_do_json_shit_dynamically_on_docker"] @@ -65,10 +64,8 @@ async def cook_deez(ctx: commands.Context, links: str): await client.session.close() await info.edit(content="Uploading to Google Drive. This may take a while.") - uploader = DriveUploader('./res/token.json') - with ThreadPoolExecutor() as pool: - bot: commands.Bot = ctx.bot - results = await bot.loop.run_in_executor(pool, uploader.batch_upload, [str(ctx.author.id)], 'NOOBGPT', True, True) + uploader = AsyncDriveUploader('./res/token.json') + results = await uploader.batch_upload([str(ctx.author.id)], 'NOOBGPT', True, True) collect_urls = [] for result in results: if result['type'] == 'folder': @@ -83,13 +80,20 @@ async def cook_deez(ctx: commands.Context, links: str): def folder_printer(result, ctx: commands.Context, collect_urls: list): if result['name'] != "__artwork": # check if not artwork folder - if result['name'] == str(ctx.author.id) and not sub_root_cheeks(result): # check if sub-root folder - data = {"label": "Root", "url": result.get('link', link_null), "emoji": "🫚", - "name": "Root", "size": result['size_human_readable']} - else: # album folder - data = {"label": len(collect_urls), "url": result.get('link', link_null), "emoji": None, - "name": result['name'], "size": result['size_human_readable']} - collect_urls.append(data) + if result['name'] != str(ctx.author.id): # album folder + collect_urls.append( + { + "label": str(len(collect_urls)), "url": result.get('link', link_null), "emoji": None, + "name": result['name'], "size": result['size_human_readable'] + } + ) + elif not sub_root_cheeks(result): # check if sub-root folder + collect_urls.append( + { + "label": "Root", "url": result.get('link', link_null), "emoji": "🫚", + "name": "Root", "size": result['size_human_readable'] + } + ) # print(f"Folder: {result['name']}") # print(f"Public Link: {result.get('link', 'No link')}") diff --git a/requirements.txt b/requirements.txt index 80a7905..0d760a8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,5 +17,5 @@ ffmpeg-python lxml pydub git+https://github.com/Xtr4F/PyCharacterAI -google-api-python-client +aiogoogle streamrip \ No newline at end of file diff --git a/ytdlp_.py b/ytdlp_.py index c8eddfe..b65845f 100644 --- a/ytdlp_.py +++ b/ytdlp_.py @@ -5,9 +5,8 @@ import os import asyncio import time -from concurrent.futures import ThreadPoolExecutor # new hack from util_discord import command_check, description_helper -from api_gdrive import DriveUploader +from api_gdrive import AsyncDriveUploader async def YTDLP(ctx: commands.Context | discord.Interaction, arg1: str, arg2: str): if await command_check(ctx, "ytdlp", "media"): @@ -45,23 +44,17 @@ async def YTDLP(ctx: commands.Context | discord.Interaction, arg1: str, arg2: st await asyncio.to_thread(ydl.download, [arg2]) # old hack if os.path.isfile(filename): try: - uploader = DriveUploader('./res/token.json') - with ThreadPoolExecutor() as pool: - bot: commands.Bot = ctx.bot - results = await bot.loop.run_in_executor(pool, uploader.batch_upload, [filename], 'NOOBGPT', True, True) - link = None - for result in results: - link = f"[{result.get('name')}]({result.get('link')})" - break - if not results: link = "[rickroll placeholder](https://youtube.com/watch?v=dQw4w9WgXcQ)" # should be impossible + uploader = AsyncDriveUploader('./res/token.json') + results = await uploader.batch_upload([filename], 'NOOBGPT', True, True) + links = [{'label': 'Download', 'emoji': '⬇️', 'url': results[0].get('link')}] # file = discord.File(filename) if isinstance(ctx, commands.Context): # await ctx.reply(file=file) - await ctx.reply(link) + await ctx.reply(filename, view=LinksView(links, ctx)) await msg.edit(content=f"`{filename}` has been prepared successfully!\nTook {round(time.time() * 1000)-old}ms") if isinstance(ctx, discord.Interaction): # await ctx.followup.send(file=file) - await ctx.followup.send(link) + await ctx.followup.send(filename, view=LinksView(links, ctx)) await ctx.edit_original_response(content=f"`{filename}` has been prepared successfully!\nTook {round(time.time() * 1000)-old}ms") except: error_message = f"Error: An error occurred while cooking `{filename}`\nFile too large!" @@ -83,6 +76,25 @@ async def YTDLP(ctx: commands.Context | discord.Interaction, arg1: str, arg2: st if isinstance(ctx, discord.Interaction): await ctx.edit_original_response(content=error_message) +class CancelButton(discord.ui.Button): + def __init__(self, ctx: commands.Context): + super().__init__(emoji="❌", style=discord.ButtonStyle.success) + self.ctx = ctx + + async def callback(self, interaction: discord.Interaction): + if interaction.user != self.ctx.author: + return await interaction.response.send_message(f"Only <@{self.ctx.author.id}> can interact with this message.", + ephemeral=True) + await interaction.response.defer() + await interaction.delete_original_response() + +class LinksView(discord.ui.View): + def __init__(self, links: list, ctx: commands.Context): + super().__init__(timeout=None) + for x in links[:24]: + self.add_item(discord.ui.Button(style=discord.ButtonStyle.link, url=x['url'], label=x['label'], emoji=x['emoji'])) + self.add_item(CancelButton(ctx)) + def checkSize(info, *, incomplete): filesize = info.get('filesize') if info.get('filesize') else info.get('filesize_approx') if filesize and filesize > 25000000: # 25mb