provenance.py

import psycopg2
import sqlparse as sp
import argparse
import sys
import copy
import time
import operator
import provenanceUtils as ut
import traceback

#When craeting the auxillary provenance tables, need to add a column called provtuple, which is a sorted array
#[t1, t2, t3...]

tablesWithProvenance = {}

def tablesInQuery(query):
    parsed = sp.parse(query)[0]
    output = []
    state = ""
    for token in parsed.tokens:
        cur = token.value
        if cur == " ":
            continue
        if cur.lower() == "from":
            state = "first table"
            continue

        if state == "first table":
            if "(" not in cur and "join" in cur:
                state = "second table"
                continue
            if "(" in cur and ")" in cur:
                #subquery
                output.extend(tablesInQuery(cur[cur.index("(")+1 : cur.index(")")]))
            else:
                output.append(cur.split("as")[0])
            continue
        
        if state == "second table":
            if "(" in cur and ")" in cur:
                #subquery
                output.extend(tablesInQuery(cur[cur.index("(")+1 : cur.index(")")]))
            else:
                output.append(cur.split("as")[0])
            break

    return output


#Preprocess query
def preprocess_query(query, cur, conn):
    tables = tablesInQuery(query)
    output = []
    for table in tables:
        if table not in tablesWithProvenance:
            command1 = "select add_provenance('{}');".format(table)
            command2 = "select create_provenance_mapping('{}_provtuple_mapping','{}','provtuple');".format(table, table)
            ut.commitAndExecute(conn, cur, command1)
            ut.commitAndExecute(conn, cur, command2)
            tablesWithProvenance[table] = 1
        formula = " formula(provenance(), '{}_provtuple_mapping'), ".format(table)
        index = query.find('select') + 6
        tmp = query[:index] + formula + query[index:]
        output.append(tmp)
        
    return (tables,output)

#formulas:[[formula1, formula2, ...],
#           [formula1, ...]
#           ...
#         ]
#fs: [formula1, formula2, ...]

def process_formula(fs, tables):
    elements = []
    for i, formula in enumerate(fs):
        nf, idx = ut.remove1sAndAddTablename(formula, tables[i])
        elements.append(((nf, idx)))
    elements.sort(key=operator.itemgetter(1))
    formula_tmp = [x[0] for x in elements]
    output = " ".join(formula_tmp)
    return output

def get_header(query):
    parsed = sp.parse(query)[0]
    headers = (parsed.tokens)[2].value.split(",")
    output = ""
    for h in headers:
        h = h.strip()
        header = ""
        if "." in h:
            header = h.split(".")[1]
        else:
            header = h
        output += header + "   |   "
    return output


def exe_query(query, out_file, cur, conn):
    out_file.write(query + "\n")
    print("executing query: " + query)
    tables, queries = preprocess_query(query, cur, conn)
    result_t = None
    formulas = []
    for processed_query in queries:
        ut.commitAndExecute(conn, cur, processed_query)
        output = cur.fetchall()
        if result_t == None:
            result_t = output
        for i, tu in enumerate(output):
            if len(formulas) == i:
                formulas.append([tu[0]])
            else:
                formulas[i].append(tu[0])
    
    print("formulas", formulas)
    processed_formulas = []
    for fs in formulas:
        rowformula = process_formula(fs, tables)
        processed_formulas.append(rowformula)
    print("processed formula", processed_formulas)
    result__ = []
    for i, row in enumerate(result_t):
        rowl = list(row)
        rowl[0] = processed_formulas[i]
        result__.append(rowl)
    print("result", result__)

    # Combine rows with same output
    result = ut.combineSameRow(result__)
    
    # Now print the processed query
    header = get_header(query)
    out_file.write("provenance formula   |   " + header+ "provsql-token"+"\n\n")
    for row in result:
        line = " | ".join(row)
        out_file.write(line+"\n")
    out_file.write("\n---------------------------\n\n")
    

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("env", help="Environment configuration file")
    parser.add_argument("-q", "--query", help="Relative path for query file.")
    parser.add_argument("-o", "--output", help="Relative path for output file. Default as stdout")
    args = parser.parse_args() 

    env_config = ''
    queries = []
    out_file = sys.stdout
    with open(args.env) as fp:
        line = fp.readline()
        while line:
            tmp = line.strip()
            env_config  = env_config + tmp + ' '
            line = fp.readline()
    print("database configuration " + env_config)

    if args.query:
        with open(args.query) as fp:
            line = fp.readline()
            while line:
                queries.append(line.strip())
                line = fp.readline()
    else:
        print("Please specify at least one query file.")
        exit(1)

    if args.output:
        out_file = open(args.output, 'w')
    
    conn = None
    try:
        print('connecting to databse...')
        conn = psycopg2.connect(env_config)
        cur = conn.cursor()
        ut.commitAndExecute(conn, cur, 'SET search_path TO public, provsql;')

        for query in queries:
            exe_query(query, out_file, cur, conn)
            print("FINISHED: "+query)

        print("Done!")
    except (Exception, psycopg2.DatabaseError) as error:
        print(traceback.format_exc())
        exit(1)
    finally:
        out_file.close()
        if conn is not None:
            conn.close()
            print('Database connection closed.')

if __name__ == '__main__':
    main()