savestate

prepared_steps/1_extract_data/extract_via_simple_salesforce.py

@@ -10,8 +10,9 @@ import time
 import sys
 import json
 import pandas as pd
-from dotenv import load_dotenv, find_dotenv
-from simple_salesforce import Salesforce
+from sys import path
+path.append('../..')
+from sf_auth import get_sf_connection
 
 done = False
 
@@ -25,44 +26,7 @@ def animate():
         time.sleep(0.1)
     sys.stdout.write('\rDone!     ')
 
-def get_credentials(context):
-    """
-    Get credentials for a given context from the .env file
-    
-    Args:
-        context (str): Context name (e.g., 'qa2', 'prod')
-    
-    Returns:
-        dict: Credentials dictionary with username, password, and security_token
-    """
-    context = context.upper()
-            
-    # Initialize credentials dictionary
-    credentials = {
-        'USERNAME': None,
-        'PASSWORD': None,
-        'SECURITY_TOKEN': None
-    }
-    
-    if context != 'PROD':
-        credentials['DOMAIN'] = 'test'
-
-    # Load the .env file explicitly from one directory above
-    env_file = find_dotenv("../.env")
-    load_dotenv(env_file)
-    
-    # Load all environment variables
-    env_vars = os.environ
-
-
-    
-    for key, value in env_vars.items():
-        #print(f'{context}_SF_', key, value)
-        if f'{context}_SF_' in key:
-            credential_key = key.split(f'{context}_SF_')[-1].upper()
-            credentials[credential_key] = value
-            
-    return credentials
+# ...existing code for animate() function...
 
 def extract_data(object_id, query, output_path='output', context='qa2'):
     """
@@ -74,45 +38,15 @@ def extract_data(object_id, query, output_path='output', context='qa2'):
         context (str): Context name for credentials (e.g., 'qa2', 'prod')
     """
     try:
-
         global done
         done = False
-        # Get credentials based on context
-        credentials = get_credentials(context)
         
-        if not all(credentials.values()):
-            raise ValueError(f"Missing credentials for context: {context}")
+        # Get Salesforce connection using the new module
+        sf = get_sf_connection(context)
         
-        if(context == 'prod'):
-            # Initialize Salesforce bulk connector
-            sf = Salesforce(
-                username=credentials['USERNAME'],
-                password=credentials['PASSWORD'],
-                security_token=credentials['SECURITY_TOKEN'],
-                version='62.0'
-            )
-        else:
-            # Initialize Salesforce bulk connector
-            sf = Salesforce(
-                username=credentials['USERNAME'],
-                password=credentials['PASSWORD'],
-                security_token=credentials['SECURITY_TOKEN'],
-                domain=credentials['DOMAIN'],
-                version='62.0'
-            )
-        
-        # Create a simple query for the desired object
-        #soql_query = f"""
-        #    SELECT Id, City__c, Country__c, GeoY__c, GeoX__c, PostalCode__c, Street__c, Extension__c, HouseNo__c, FlatNo__c, Floor__c FROM SCInstalledBaseLocation__c WHERE Country__c = 'NL'
-        #"""
-
         t = threading.Thread(target=animate)
         t.start()
 
-        #sf.bulk2.__getattr__("SCInstalledBaseLocation__c").download(
-        #    soql_query, path="./", max_records=2000000
-        #)
-
         results = sf.bulk2.__getattr__(object_id).query(
             query, max_records=2000000
         )
@@ -124,37 +58,7 @@ def extract_data(object_id, query, output_path='output', context='qa2'):
         time.sleep(10)
         done = True
         t.do_run = False
-        
-        """
-        # Execute the Bulk query job
-        job = sf.bulk2.__getattr__("SCInstalledBase__c").query(soql_query)
-        
-        # Polling for job completion (might take a moment)
-        job_id = job['id']
-        while True:
-            status = sf.bulk.job(job_id).get()['status']
-            if status == 'Complete' or status == 'Closed' :
-                break
-            if status == 'Aborted':
-                exit(1)
-            if status == 'Failed':
-                raise ValueError(f'Job failed: {job_id}')
-            
-        
-        # Get the results
-        result = sf.bulk.result(job_id)
-        df = pd.DataFrame(result.records)
-        
-        # Create output directory if it doesn't exist
-        os.makedirs(output_path, exist_ok=True)
-        
-        # Save to CSV file
-        csv_file = os.path.join(output_path, f'{object_id}_data.csv')
-        df.to_csv(csv_file, index=False)
-        
-        print(f'Successfully extracted {len(df)} records from {object_id}')
-        return csv_file
-        """
+
     except Exception as e:
         done = True
         t.do_run = False

prepared_steps/1_extract_data/queries.json

@@ -35,7 +35,10 @@
             "query": "SELECT Id, ValidTo__c, ValidFrom__c, Country__c, City__c, PostalCode__c, District__c, Street__c, HouseNo__c, Extension__c, FlatNo__c, Floor__c, GeoY__c, GeoX__c, Resource__c, Resource__r.Employee__r.Name, Stock__c, Stock__r.ID2__c FROM SCResourceAssignment__c WHERE Country__c = '{country}' AND ValidTo__c >= TODAY"
         },{
             "sobject": "ContractTemplates",
-            "query": "SELECT id, name, TemplateName__c, status__c, Brand__r.Name, Country__c, Runtime__c FROM sccontract__c WHERE Template__c = null AND Country__c = '{country}'"
+            "query": "SELECT id, name, TemplateName__c, status__c, Brand__r.Name, Country__c, Runtime__c FROM SCContract__c WHERE Template__c = null AND Country__c = '{country}'"
+        },{
+            "sobject": "SCContract__c",
+            "query": "SELECT id, name, Template__c, status__c, Brand__r.Name, Country__c, Runtime__c, EndDate__c, StartDate__c, Account__c, AccountOwner__c, IoT_Registration_Status__c FROM SCContract__c WHERE Template__c != null AND EndDate__c >= TODAY AND Country__c = '{country}' limit 3"
         }
     ]
 }