insert file contents refactored to take in variable number of columns

db/util/data/products.csv

@@ -1,4 +1,4 @@
-name,regionId,categoryId
+name,regionid,categoryid
 Allspice - Whole,8,1
 Cocoa Powder,1,1
 Cardamom - Green,3,1

db/util/insert_file_contents.py

@@ -1,33 +1,30 @@
 import csv
+from re import M
 from psycopg2 import sql
 
 # function to read from a given csv file into postgres
-def insert_file_contents(conn, cur, file_path, table_name, column_count):
+def insert_file_contents(conn, cur, file_path, table_name):
     with open(file_path, 'r', encoding='utf-8-sig') as f:
         reader = csv.reader(f)
-        header_names = ""
         first_row_accessed = False
+        header_names = ""
+        num_columns = 0
 
         for row in reader:
             # get row values from first row of reader
             if not first_row_accessed:
                 header_names = [item for item in row]
+                num_columns = len(header_names)
                 first_row_accessed = True
                 continue
 
-            # execute table insertion for each row of csv based on number of columns
-            if column_count == 3:
-                cur.execute(sql.SQL("INSERT INTO {table} ({h1}, {h2}, {h3}) VALUES (%s, %s, %s)".format(
-                        table=table_name,
-                        h1=header_names[0],
-                        h2=header_names[1],
-                        h3=header_names[2]
-                    )),
-                    row
-                )
-            elif column_count == 1:
-                cur.execute(sql.SQL("INSERT INTO {table} ({field}) VALUES (%s)".format(table=table_name, field=header_names[0])), row)
-            else:
-                raise
+            mapped_columns = [header_names[i] for i in range(num_columns)]
+            prepared_q = sql.SQL("INSERT INTO {TABLE} ({COLS}) VALUES ({VALS})").format(
+                TABLE=sql.Identifier(table_name),
+                COLS=sql.SQL(', ').join(map(sql.Identifier, mapped_columns)),
+                VALS=sql.SQL(', ').join(sql.Placeholder() * len(mapped_columns))
+            )
+            
+            cur.execute(prepared_q, [item for item in row])
 
     conn.commit()

db/util/main.py

@@ -1,8 +1,7 @@
-import psycopg2
-import csv
-import os
 from insert_file_contents import insert_file_contents
+import psycopg2
 from psycopg2 import sql
+import os
 
 # read data from environment if present
 env_path = "../../.env"
@@ -28,8 +27,8 @@ conn = psycopg2.connect("dbname=e-commerce-092122 user=mikayladobson")
 cur = conn.cursor()
 
 # read contents of each file into postgres
-insert_file_contents(conn, cur, "./data/categories.csv", 'category', 1)
-insert_file_contents(conn, cur, "./data/regions.csv", 'region', 1)
-insert_file_contents(conn, cur, "./data/products.csv", 'product', 3)
+insert_file_contents(conn, cur, "./data/categories.csv", 'category')
+insert_file_contents(conn, cur, "./data/regions.csv", 'region')
+insert_file_contents(conn, cur, "./data/products.csv", 'product')
 
 print("Database insertions executed successfully.")