clean postgre single table

clean postgre single table

DBCleanByTableName.py

from Query import Query
import time
class DBCleanByTableName:

    vhm_id_range = 4294967295l
    _query = Query()
    _delete_tables = []

    def clean_single_table(self, table_name, owner_id):
        """
        Clean up single table by table_name and owner_id
        """
        try:
            self._delete_tables = []
            info('[{0}]Start delete data from table[{1}]...', self.__get_time(), table_name)
            self.__purge_singe_table_data(table_name, owner_id)
            info('[{0}]Delete table[{1}] sucessful! Constraint_tables : {2}', self.__get_time(), table_name, str(self._delete_tables))
        except Exception,e:
            error('Delete data from table[{0}] fail!', table_name)
            raise e
        finally:
            self._query.close_pgdb_conn()

    def __purge_singe_table_data(self, table_name, owner_id):
        min_id = owner_id << 32
        max_id = min_id + self.vhm_id_range
        constraint_tables = self._query.get_constraint_table(table_name)
        if not constraint_tables:
            self.__do_purge_data(table_name, min_id, max_id, owner_id)
        else:
            for item in constraint_tables:
                self._delete_tables.append(item[0])
                self.__purge_singe_table_data(item[0], owner_id)
            self.__do_purge_data(table_name, min_id, max_id, owner_id)

    def __do_purge_data(self, table_name, min_id, max_id, owner_id):
        table_columns = self._query.get_column_by_table(table_name)
        if table_columns:
            if ['id'] in table_columns and ['owner_id'] in table_columns:
                column = 'id'
            else:
                column = table_columns[0].get(0)
        else:
            error("Don't get columns of {0}", table_name)
            raise
        delete_sql = 'DELETE FROM {0} where {1} between {2} and {3}'.format(table_name, column, min_id, max_id)
        self._query.delete_table_by_sql(delete_sql)

    def __get_time(self):
        return time.strftime('%Y-%m-%d %X', time.localtime())


if __name__ == '__main__':
    result = DBCleanByTableName().clean_single_table('hm_nwk_plcy',102)

Query.py

import psycopg2
import psycopg2.extras
class Query:
    """
    Connect pgdb and execute sql
    """

    _env = Env()
    _db_host = _env.get('db.host')
    _db_port = _env.get('db.port')
    _db_user = _env.get('db.user.name')
    _db_password = _env.get('db.user.password')
    _db_name = _env.get('db.name')

    try:
        _pgdb_conn = psycopg2.connect(host=_db_host, dbname=_db_name, port=_db_port, user=_db_user, password=_db_password)
        _cursor = _pgdb_conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        info('[Get pgdb connect sucess!] Host : {0}, DB name : {1}', _db_host, _db_name)
    except Exception, e:
        error("conntect postgre database failed, ret = {0}", e)
        _cursor.close()
        _pgdb_conn.close()
        raise e

    get_constraint_table_sql = 'SELECT tc.table_name FROM information_schema.table_constraints AS tc JOIN information_schema.key_column_usage AS kcu ' \
                               'ON tc.constraint_name = kcu.constraint_name JOIN information_schema.constraint_column_usage AS ccu ' \
                               'ON ccu.constraint_name= tc.constraint_name WHERE constraint_type = \'FOREIGN KEY\' AND ccu.table_name=\'{0}\''

    get_column_by_table_sql = 'SELECT a.attname FROM pg_class as c,pg_attribute as a  where c.relname =\'{0}\' and a.attrelid = c.oid and a.attnum>0 '

    def get_constraint_table(self, table_name):
        """
        Get the constraint tables of table_name
        """
        sql = self.get_constraint_table_sql.format(table_name)
        self._cursor.execute(sql)
        constraint_tables = self._cursor.fetchall()
        self._pgdb_conn.commit()
        return constraint_tables

    def get_column_by_table(self, table_name):
        """
        Get all columns of table_name
        """
        sql = self.get_column_by_table_sql.format(table_name)
        self._cursor.execute(sql)
        table_columns = self._cursor.fetchall()
        self._pgdb_conn.commit()
        return table_columns

    def delete_table_by_sql(self, delete_sql):
        self._cursor.execute(delete_sql)
        self._pgdb_conn.commit()

    def close_pgdb_conn(self):
        self._cursor.close()
        self._pgdb_conn.close()
        info('[Close pgdb connect sucess!] Host : {0}, DB name : {1}', self._db_host, self._db_name)