[Python] Querying data iteratively, is there a better way of doing this?

I'm trying to gather data using a set of IDs from one dataset and querying them against a PostgreSQL table using psycopg2 in Python. Here is the code:

for i in range(0, len(ids_df), 2000):
print(f"Iteration #: {i}")
ids_subset = ids_df.iloc[i:i+2000]['customerid'].tolist()
cust_query = f"select cust_id, status, update_dtime from customers WHERE cust_id IN ({', '.join(map(str, ids_subset))})"
try:
with psycopg2.connect(**db_params) as conn:
with conn.cursor() as cursor:
cursor.execute(cust_query)
result = cursor.fetchall()
column_names = [desc[0] for desc in cursor.description]
mode = "w" if i == 0 else "a"
with open("file.csv", mode, newline="", encoding="utf-8") as f:
writer = csv.writer(f)
if i == 0:
writer.writerow(column_names)
if not result:
break
writer.writerows(result)
except Exception as e:
print(f"Error processing {db_params['dbname']}: {e}")


I then take the CSV and run my analysis. It works, but I'm dealing with 6 million records and the process can take 35-40 minutes to run. I also have concerns about the load I might be putting on the server side.

I picked 2000 as the step value because it was what worked without returning no rows. Anything greater than some amount in between 2000-3000 wasn't working.

I considered using a server side cursor with psycopg2, but I don't understand how that would work when I'm searching for specific IDs.

Continue reading...