diff --git a/etl2/etl.py b/etl2/etl.py index cb62b40..66fb854 100644 --- a/etl2/etl.py +++ b/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() diff --git a/etl2/etl_drop_tables/etl_drop_tables.py b/etl2/etl_drop_tables/etl_drop_tables.py new file mode 100644 index 0000000..fbb9486 --- /dev/null +++ b/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)