Create workflows/storage.py

workflows/storage.py

@@ -0,0 +1,55 @@
+from prefect import task
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaFileUpload, MediaIoBaseUpload
+from config.settings import settings
+from io import BytesIO
+import base64
+
+@task(retries=2)
+def upload_to_google_drive(image_data: str, filename: str, folder_id: str = None) -> str:
+    """Upload image to Google Drive"""
+    
+    # Parse credentials
+    credentials = service_account.Credentials.from_service_account_info(
+        settings.google_credentials_dict,
+        scopes=['https://www.googleapis.com/auth/drive.file']
+    )
+    
+    # Build Drive service
+    service = build('drive', 'v3', credentials=credentials)
+    
+    # Prepare file
+    if image_data.startswith('data:image'):
+        # Base64 data URL
+        header, encoded = image_data.split(',', 1)
+        image_bytes = base64.b64decode(encoded)
+        media = MediaIoBaseUpload(
+            BytesIO(image_bytes),
+            mimetype='image/png',
+            resumable=True
+        )
+    else:
+        # URL or file path - download first
+        import requests
+        response = requests.get(image_data)
+        media = MediaIoBaseUpload(
+            BytesIO(response.content),
+            mimetype='image/png',
+            resumable=True
+        )
+    
+    # File metadata
+    file_metadata = {
+        'name': filename,
+        'parents': [folder_id] if folder_id else []
+    }
+    
+    # Upload
+    file = service.files().create(
+        body=file_metadata,
+        media_body=media,
+        fields='id, webViewLink'
+    ).execute()
+    
+    return file.get('webViewLink')