Update app.py

app.py

@@ -28,63 +28,55 @@ temp_dir = tempfile.mkdtemp()
 
 # Function to retrieve all files from the Salesforce custom object
 def get_all_salesforce_files(custom_object_api_name):
-    # Step 1: Query all record IDs from the custom object
-    query_all_records = f"SELECT Id FROM {custom_object_api_name}"
+    # Step 1: Query all record IDs from the custom object, including the Google Drive File ID field
+    query_all_records = f"SELECT Id, Google_Fille_ID__c FROM {custom_object_api_name} WHERE Google_Fille_ID__c != NULL"
     all_records = sf.query(query_all_records)
-    record_ids = [record['Id'] for record in all_records['records']]
+    records = all_records['records']
     
-    all_files = []
-
-    # Step 2: Loop through each record ID and get linked files from the Files object
-    for record_id in record_ids:
-        query_files = f"SELECT Id, ContentDocumentId FROM ContentDocumentLink WHERE LinkedEntityId = '{record_id}'"
-        results = sf.query(query_files)
-        
-        if results['records']:
-            all_files.extend(results['records'])
-    
-    return all_files
+    return records
 
 # Function to download a file from Google Drive and save it to a temporary location
 def download_google_drive_file(file_id, file_name):
-    request = drive_service.files().get_media(fileId=file_id)
-    file_path = os.path.join(temp_dir, file_name)
-    with open(file_path, 'wb') as f:
-        downloader = MediaIoBaseDownload(f, request)
-        done = False
-        while not done:
-            status, done = downloader.next_chunk()
-            if status:
-                print(f"Download progress: {int(status.progress() * 100)}%")
-    return file_path
+    try:
+        request = drive_service.files().get_media(fileId=file_id)
+        file_path = os.path.join(temp_dir, file_name)
+        with open(file_path, 'wb') as f:
+            downloader = MediaIoBaseDownload(f, request)
+            done = False
+            while not done:
+                status, done = downloader.next_chunk()
+                if status:
+                    print(f"Download progress: {int(status.progress() * 100)}%")
+        return file_path
+    except Exception as e:
+        print(f"Error downloading file {file_name}: {e}")
+        return None
 
 # Function to save all files linked to records from the Salesforce custom object
 def save_salesforce_files(custom_object_api_name):
-    files = get_all_salesforce_files(custom_object_api_name)
-    for file in files:
-        content_doc_id = file['ContentDocumentId']
-        
-        # Query ContentDocument to get the Title and FileExtension
-        content_doc_query = f"SELECT Title, FileExtension FROM ContentDocument WHERE Id = '{content_doc_id}'"
-        content_doc_result = sf.query(content_doc_query)
-
-        if not content_doc_result['records']:
-            print(f"No details found for ContentDocumentId: {content_doc_id}")
-            continue
-        
-        content_doc = content_doc_result['records'][0]
-        file_name = content_doc['Title'] + '.' + content_doc['FileExtension']
+    records = get_all_salesforce_files(custom_object_api_name)
+    for record in records:
+        google_drive_file_id = record['Google_Drive_File_ID__c']
+        record_id = record['Id']
         
-        # Download the file from Google Drive using the ContentDocument ID (assuming the file is stored in Google Drive)
-        try:
-            print(f"Downloading file {file_name} from Google Drive...")
-            file_path = download_google_drive_file(content_doc_id, file_name)
-            if os.path.exists(file_path):
-                print(f"File {file_name} saved successfully to {file_path}")
-            else:
-                print(f"Failed to save file {file_name}")
-        except Exception as e:
-            print(f"Error downloading file {file_name}: {e}")
+        # Check if the Google Drive File ID exists
+        if google_drive_file_id:
+            # Get the file metadata from Google Drive to obtain the filename
+            try:
+                file_metadata = drive_service.files().get(fileId=google_drive_file_id, fields='name').execute()
+                file_name = file_metadata['name']
+                
+                print(f"Downloading file {file_name} from Google Drive...")
+                file_path = download_google_drive_file(google_drive_file_id, file_name)
+                
+                if file_path and os.path.exists(file_path):
+                    print(f"File {file_name} saved successfully to {file_path}")
+                else:
+                    print(f"Failed to save file {file_name}")
+            except Exception as e:
+                print(f"Error retrieving file metadata for Google Drive ID {google_drive_file_id}: {e}")
+        else:
+            print(f"No Google Drive File ID found for record: {record_id}")
 
 if __name__ == "__main__":
     # Replace 'YOUR_CUSTOM_OBJECT_API_NAME' with the Salesforce custom object API name