SocialEngineerDBuilder/SG.py

# -*- coding: utf-8 -*-
"""
"""
from __future__ import print_function

import sys
import sqlalchemy
from User import user_model
from Common import get_DB_engine,init_User_table,get_DB_session,preConstruct


def gen_parse_class(f_name):
    """
    """
    # Read the first line
    with open(f_name) as f:
        sep,columns = preConstruct(f.readline())

    if sep == 'SPACE':
        sep = ' '
    columns = [c.strip() for c in columns]

    # check column name
    if not check_col_name(columns): raise Exception('Invalid column name!')

    # inner define function
    def parse_line_to_dic(self,line):
        """
        The generated class __init__()
        """
        row_elements = line.split(sep)
        row_elements = [r.strip() for r in row_elements]
        kargs = {}
        for ele_i in range(len(columns)):
            try:
                if columns[ele_i] == 'ignore':
                    pass
                else:
                    kargs[columns[ele_i]] = row_elements[ele_i]
            except Exception,e:
                sys.stdout.write(e)
        # do the init
        super(spec_user_model,self).__init__(**kargs)

    # inner define class
    spec_user_model = type('spec_user_model',(user_model,),{'__init__':parse_line_to_dic})

    # return class like this:
    """
     Class spec_user_model(user_model):
         def __init__(self,):
             row_elements = line.split(sep)
             kargs = {}

             for ele_i in range(len(row_elements)):
                 kargs[ele_i] = row_elements[ele_i]
                 # do the init
                 super().__init__(**kargs)
    """
    # usage():
    # user = spec_user_model(username='user',password='pass',...)
    # session.commit(user)
    return spec_user_model


# do one job
def save_to_db(fname,user_class,single_file=False,lines_num=1000):
    """
    """
    sess = get_DB_session()
    #sess.autocommit = True
    with open(fname) as f:
        if single_file==True:
            f.readline()
        while True:
            line_n = 0
            for line in f:
                try:
                    user = user_class(line)
                    sess.add(user)
                    line_n += 1

                    if line_n % lines_num == 0:
                        sess.commit()
                        print("\r Processed {0} lines\r".format(line_n),file=sys.stdout,end=" ")
                        sys.stdout.flush()
                except Exception,e:
                    print("exception found :\n  line:{0}\n  content:{1}   exception:{2}\n".format(line_n,line,e))
            sess.commit()
            print("\n {0} datalines Finished\r".format(line_n),file=sys.stdout)
            break
    return True

# do all dirty jobs
def save_all_to_db():
    init_User_table()
    if len(sys.argv) == 2:
        user_class = gen_parse_class(sys.argv[1])
        save_to_db(sys.argv[1],user_class,single_file=True)
    elif len(sys.argv) == 3:
        user_class = gen_parse_class(sys.argv[2])
        save_to_db(sys.argv[1],user_class,single_file=False)
    elif len(sys.argv) == 4:
        user_class = gen_parse_class(sys.argv[2])
        save_to_db(sys.argv[1],user_class,single_file=True)
    else:
        raise Exception('Error argument number!')

# build column line
def build_col_line(columns):
    line = ""
    for c in columns:
        # init
        if line == '':
            line = '{0}'.format(c)
            continue
        # ignore case
        if c == 'ignore':
            line = "{0},@dummy".format(line)
        # append col
        else:
            line =  "{0},{1}".format(line,c)
    return line

# build sql string
def build_sql_by_tmpl(source_file,template,single_file=False,model=user_model):
    """
    """
    if single_file == True:
        template = source_file
        ignore_lines = 'IGNORE 1 LINES'
    else:
        ignore_lines = ''

    with open(template) as f:
        sep,columns = preConstuct(f.readline())

    if sep == 'SPACE':
        sep = ' '
    columns = [c.strip() for c in columns]
    if not check_col_name(columns): return False

    col_line = build_col_line(columns)
    table_name = model.__tablename__

    sql_query = """LOAD DATA LOCAL INFILE '{0}'
    INTO TABLE {1}
    FIELDS TERMINATED BY '{2}'
    LINES TERMINATED BY '\\r\\n' 
    {3}
    ({4});
    """.format(
        source_file,
        table_name,
        sep,
        ignore_lines,
        col_line
    )
    return sql_query


def do_load_data_from_file():
    init_User_table()
    if len(sys.argv) == 2:
        do_load_data_from_file(sys.argv[1],sys.argv[1],single_file=True)
    elif len(sys.argv) == 3:
        do_load_data_from_file(sys.argv[1],sys.argv[2],single_file=False)
    elif len(sys.argv) == 4:
        do_load_data_from_file(sys.argv[1],sys.argv[2],single_file=True)
    else:
        raise Exception('Error argument number!')

if __name__ == '__main__':
    save_all_to_db()