diff --git a/delete_database.py b/delete_database.py new file mode 100755 index 0000000..c32834f --- /dev/null +++ b/delete_database.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python +import sys +import os +import re +# import time +import logging +import argparse as arg +import psycopg2 as pg +from psycopg2.extensions import AsIs + +from classes.CustomFormater import CustomFormatter +from classes.Timer import Timer + +# Schema name is NAMEDATALEN-1 (PGSQL source code) +# -> src/include/pg_config_manual.h +def check_schema_name(arg_value, pat=re.compile(r"^[a-z0-9A-Z]{1,63}$")): + if not pat.match(arg_value): + raise ValueError + return arg_value + + +def parse_args(): + """ + Parse arguments + """ + parser = arg.ArgumentParser('Process csv files from INSEE') + + parser.add_argument('--connection-file', '-f', + help='Postgresql connexion file', + default='.pgconn' + ) + + # As whe use AsIs function to not include '' in our sql queries for schema + # name, me mus ensure that it is not contains SQL special characters + # like comments --, final coma etc and avoid SQL injection + parser.add_argument('--schema-name', + help='Database schema name', + type=check_schema_name, + default='insee' + ) + + debug_group = parser.add_mutually_exclusive_group() + debug_group.add_argument('--verbose', '-V', + help='Verbose output', + action='store_true') + debug_group.add_argument('--debug', '-d', + help='Activate debug mode', + action='store_true') + return parser.parse_args() + + +if __name__ == '__main__': + args = parse_args() + #logging.basicConfig(level=logging.DEBUG) + logger = logging.getLogger() + tty_handler = logging.StreamHandler() + + # create console handler with a higher log level + tty_handler.setFormatter(CustomFormatter()) + logger.addHandler(tty_handler) + + if args.verbose is True: + logger.setLevel(logging.INFO) + logger.info('VERBOSE mode activated') + + if args.debug is True: + logger.setLevel(logging.DEBUG) + logger.debug('DEBUG mode activated') + + t = Timer(logger=logger.info) + logging.debug('Import pgsql connection file {}'.format(args.connection_file)) + with open(args.connection_file) as cf: + pg_conn = cf.read() + + t.start('Delete schema') + logger.debug('Delete schema {}'.format(args.schema_name)) + conn = pg.connect(pg_conn) + with conn.cursor() as curs: + try: + curs.execute('DROP SCHEMA IF EXISTS %s CASCADE', + (AsIs(args.schema_name),)) + conn.commit() + except Exception as e: + logger.error('Can\'t drop schema {}: {}'.format( + args.schema_name, + e + ) + ) + + t.stop()