Advertisement
Guest User

Untitled

a guest
Aug 24th, 2019
109
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 1.45 KB | None | 0 0
  1. import csv
  2. import psycopg2
  3. import os
  4. import re
  5.  
  6. data_dir = "path/to/files"
  7. data_files = [
  8. ('table1', 'table1.csv'),
  9. ('table2', 'table2.csv')
  10. ]
  11.  
  12. with psycopg2.connect(dbname="postgres", user="postgres", password="postgres", host="localhost") as conn:
  13. with conn.cursor() as cur:
  14. for table_name, file_name in data_files:
  15. # Loop through each CSV File
  16. csv_file = os.path.join(data_dir, file_name)
  17.  
  18. with open(csv_file, newline='') as f:
  19. # Read in CSV Data to get headers:
  20. csv_data = csv.DictReader(f)
  21.  
  22. # Format headers: Convert spaces to underscores, lowercase and remove anything not alpha-numeric or underscore:
  23. fieldnames = [re.sub(r"[^a-z0-9_]", "", i.strip().lower().replace(" ", "_")) for i in
  24. csv_data.fieldnames]
  25.  
  26. # Reset file position
  27. f.seek(0, 0)
  28.  
  29. # Drop previous table
  30. cur.execute(f"drop table if exists {table_name};")
  31.  
  32. # Recreate the table with current field names
  33. cur.execute("create table {table_name}({field_names} text)".format(
  34. table_name=table_name,
  35. field_names=" text,".join(fieldnames)))
  36.  
  37. # Add the data
  38. cur.copy_expert(f"COPY {table_name} FROM STDIN WITH CSV HEADER", f)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement