add etl for dropping tables and update etl script

etl2/etl.py

@@ -5,6 +5,7 @@
 from termcolor import colored
 from etl_routemodel.etl_routemodel import main as run_routemodel_etl
 from etl_weather.etl_weather import main as run_weather_etl
+from etl_drop_tables.etl_drop_tables import main as run_drop_tables_etl
 
 
 def validate_db_creds(db_user, db_password, db_host, db_name):
@@ -117,6 +118,54 @@ def cmd_weather():
         print(colored("Incorrect database credentials", "red"))
 
 
+def cmd_drop_tables():
+    answers = questionary.form(
+        table_names=questionary.checkbox(
+            "Select tables to drop",
+            choices=[
+                "routemodel",
+                "weather",
+            ],
+            instruction="Use 'space' to select; 'enter' to continue",
+        ),
+        db_host=questionary.text("Database host:port", default=""),
+        db_name=questionary.text("Database name", default=""),
+        db_user=questionary.text("Database user", default=""),
+        db_password=questionary.password("Database user password", default=""),
+        confirm=questionary.confirm(
+            "Confirm routemodel ETL operation",
+            default=False,
+            auto_enter=False,
+        ),
+    ).ask()
+
+    (
+        table_names,
+        db_host,
+        db_name,
+        db_user,
+        db_password,
+        confirm,
+    ) = map(
+        answers.get,
+        (
+            "table_names",
+            "db_host",
+            "db_name",
+            "db_user",
+            "db_password",
+            "confirm",
+        ),
+    )
+    if not confirm:
+        print(colored("drop tables ETL cancelled", "red"))
+    elif confirm and validate_db_creds(db_user, db_password, db_host, db_name):
+        run_drop_tables_etl(db_user, db_password, db_host, db_name, table_names)
+        print(colored("drop tables ETL success", "green"))
+    else:
+        print(colored("Incorrect database credentials", "red"))
+
+
 if __name__ == "__main__":
     print(
         colored(
@@ -125,10 +174,12 @@ def cmd_weather():
         )
     )
     etl_name = questionary.select(
-        "Select ETL operation", choices=["routemodel", "weather"]
+        "Select ETL operation", choices=["routemodel", "weather", "drop_tables"]
     ).ask()
 
     if etl_name == "routemodel":
         cmd_routemodel()
     elif etl_name == "weather":
         cmd_weather()
+    elif etl_name == "drop_tables":
+        cmd_drop_tables()

etl2/etl_drop_tables/etl_drop_tables.py

@@ -0,0 +1,27 @@
+import psycopg2
+
+
+def drop_tables(db_user, db_password, db_host, db_name, table_names):
+    host, port = db_host.split(":")
+    conn = psycopg2.connect(
+        host=host, port=port, user=db_user, password=db_password, dbname=db_name
+    )
+    cursor = conn.cursor()
+    for table_name in table_names:
+        cursor.execute(f"DROP TABLE IF EXISTS public.{table_name} CASCADE;")
+    conn.commit()
+    conn.close()
+    print(f"Successfully dropped tables: {', '.join(table_names)}")
+
+
+def main(db_user, db_password, db_host, db_name, table_names):
+    drop_tables(db_user, db_password, db_host, db_name, table_names)
+
+
+if __name__ == "__main__":
+    db_user = ""
+    db_password = ""
+    db_host = ""
+    db_name = ""
+    table_names = [""]
+    main(db_user, db_password, db_host, db_name, table_names)