Update app.py

app.py

@@ -5,6 +5,7 @@ import json
 from simple_salesforce import Salesforce
 from google.oauth2 import service_account
 from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload
 
 # Google Drive Authentication Setup
 SCOPES = ['https://www.googleapis.com/auth/drive']
@@ -15,21 +16,35 @@ drive_service = build('drive', 'v3', credentials=credentials)
 
 # Salesforce Authentication
 sf = Salesforce(
-    username='YOUR_SALESFORCE_USERNAME',
-    password='YOUR_SALESFORCE_PASSWORD',
-    security_token='YOUR_SALESFORCE_SECURITY_TOKEN',
+    username=os.getenv('SF_USERNAME'),
+    password=os.getenv('SF_PASSWORD'),
+    security_token=os.getenv('SF_SECURITY_TOKEN'),
     domain='login'
 )
 
 # Local folder to save retrieved files
-LOCAL_FOLDER = 'retrieved_files'
+LOCAL_FOLDER = r"C:\Users\bharg\Documents\Work\Sathkrutha\Retrieval"
 if not os.path.exists(LOCAL_FOLDER):
     os.makedirs(LOCAL_FOLDER)
 
-def get_salesforce_files(custom_object_id):
-    query = f"SELECT Id, Title, ContentDocumentId FROM ContentDocumentLink WHERE LinkedEntityId = '{custom_object_id}'"
-    results = sf.query(query)
-    return results['records']
+def get_all_salesforce_files(custom_object_api_name):
+    custom_object_api_name = 'MyCustomer__c'
+    # Step 1: Query all record IDs from the custom object
+    query_all_records = f"SELECT Id FROM {custom_object_api_name}"
+    all_records = sf.query(query_all_records)
+    record_ids = [record['Id'] for record in all_records['records']]
+    
+    all_files = []
+
+    # Step 2: Loop through each record ID and get linked files
+    for record_id in record_ids:
+        query_files = f"SELECT Id, Title, ContentDocumentId FROM ContentDocumentLink WHERE LinkedEntityId = '{record_id}'"
+        results = sf.query(query_files)
+        
+        if results['records']:
+            all_files.extend(results['records'])
+    
+    return all_files
 
 def download_google_drive_file(file_id, file_name):
     request = drive_service.files().get_media(fileId=file_id)
@@ -37,13 +52,14 @@ def download_google_drive_file(file_id, file_name):
     with open(file_path, 'wb') as f:
         downloader = MediaIoBaseDownload(f, request)
         done = False
-        while done is False:
+        while not done:
             status, done = downloader.next_chunk()
-            print(f"Download progress: {int(status.progress() * 100)}%")
+            if status:
+                print(f"Download progress: {int(status.progress() * 100)}%")
     return file_path
 
-def save_salesforce_files(custom_object_id):
-    files = get_salesforce_files(custom_object_id)
+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']
         content_doc = sf.ContentDocument.get(content_doc_id)
@@ -57,6 +73,6 @@ def save_salesforce_files(custom_object_id):
             print(f"Failed to save file {google_drive_file_name}")
 
 if __name__ == "__main__":
-    # Replace 'YOUR_CUSTOM_OBJECT_ID' with the Salesforce custom object Id
-    custom_object_id = 'YOUR_CUSTOM_OBJECT_ID'
-    save_salesforce_files(custom_object_id)
+    # Replace 'YOUR_CUSTOM_OBJECT_API_NAME' with the Salesforce custom object API name
+    custom_object_api_name = 'MyCustomer__c'
+    save_salesforce_files(custom_object_api_name)