5
\$\begingroup\$

I feel like there must be a better way to do the following. I would really appreciate any feedback/pointers so that I can improve my code and skills (or lack thereof!).

The goal is to allow users to query our imapala database using a python script, so that they don't have to learn SQL.

I'd like to take the following dictionary of user arguments (as an example):

args = {
'members':'all',
'subject_id':'all',
'genes':'BRCA1,TRT1,SON%,HOX%',
'gt':'hom',
'chrom':'8',
'search_db':'wgs.comgen_variant',
'search_db_cols':'all', 
'annot_db':'ref.ensembl_genes',
'annot_db_cols':'gene_name',
}

And parse them as follows:

  • If a value is 'all', no statement is added to the query

  • If a value ends in "%", create a statement as AND db_name.table_name.column LIKE "value%"

  • If a value doesn't end in "%", create a statement as AND db_name.table_name.column = "value"

The desired output should look like (this is just the beginning of the query):

SELECT wgs.comgen_variant.*, ref.ensembl_genes.gene_name    
FROM wgs.comgen_variant,ref.ensembl_genes 
WHERE (ref.ensembl_genes.gene_name = 'BRCA1' 
        OR ref.ensembl_genes.gene_name = 'TRT1' 
        OR ref.ensembl_genes.gene_name LIKE 'SON%' 
        OR ref.ensembl_genes.gene_name LIKE 'HOX%')
        AND wgs.comgen_variant.chrom = "8"
        AND wgs.comgen_variant.gt = "hom"

I came up with the following code, to separate the wildcards from the non-wildcards and create statements:

# separate wildcards from args
def find_wildcards(x):
    user_args = []
    wildcards = []
    # separate user args by comma
    arg_list = x.replace("'", "").split(',')
    for arg in arg_list:
        if arg.endswith('%'):
            wildcards.append(arg)
        else:
            user_args.append(arg)
    return user_args, wildcards

# function to turn arg_list and wild_list into statements
def process_args(arg_list, wild_list, name_arg, arg_db):
    # if there is only one user arg
    if (len(arg_list) == 1 and len(wild_list)== 0):
        query_args = 'AND {0}.{1} = '.format(arg_db, name_arg) + "'" + ", ".join(str(e) for e in arg_list) + "'"
    # if there are more than one user arguments
    elif (len(arg_list) > 1 and len(wild_list) == 0):
    conditions = ['{0}.{1} = '.format(arg_db, name_arg) + "'" + s + "'" for s in arg_list]
        query_args = "AND (" + " OR ".join(map(str, conditions)) + ")"
    # if there is only one wildcard
    elif (len(wild_list) == 1 and len(arg_list) < 1):
        query_args = 'AND {0}.{1} LIKE '.format(arg_db, name_arg) + "'" + ", ".join(str(e) for e in wild_list) + "'"
    # if there are more than one wildcard
    elif (len(wild_list) > 1 and len(arg_list) == 0):
        conditions = ['{0}.{1} LIKE '.format(arg_db, name_arg) + "'" + s + "'" for s in wild_list]
        query_args = "AND (" + " OR ".join(map(str, conditions)) + ")"
    #if there are more than one wildcard and more than one user arg
    elif (len(arg_list) > 1 and len(wild_list) > 1):
        reg_conditions = ['{0}.{1} = '.format(arg_db, name_arg) + "'" + s + "'" for s in arg_list]
        reg_args = "AND ((" + " OR ".join(map(str, reg_conditions)) + ")"
        wild_conditions = ['{0}.{1} LIKE '.format(arg_db, name_arg) + "'" + s + "'" for s in wild_list]
        wild_args = "AND (" + " OR ".join(map(str, wild_conditions)) + ")"
        query_args = (reg_args + " " + wild_args).replace(' AND ', ' OR ', 2) + ")"
    #if there is one of each
    elif (len(arg_list) == 1 and len(wild_list) == 1):
        reg_conditions = ['{0}.{1} = '.format(arg_db, name_arg) + "'" + s + "'" for s in arg_list]
        reg_args = "AND (" + " OR ".join(map(str, reg_conditions)) + ")"
        wild_conditions = ['{0}.{1} LIKE '.format(arg_db, name_arg) + "'" + s + "'" for s in wild_list]
        wild_args = "AND (" + " OR ".join(map(str, wild_conditions)) + ")"
        query_args = (reg_args + " " + wild_args).replace(') AND (', ' OR ', 1)
    else:
        query_args = ''
    return query_args

# run above function on user args
def create_query(argument, arg_name, arg_db):
    if argument != 'all':
        arg_list, wild_list = find_wildcards(argument)
        arg_query = process_args(arg_list, wild_list, arg_name, arg_db)
    else:
        arg_query = ''
    return arg_query

Then, I run the following to create a statement for each argument and merge them together:

# create members argument
member_query = create_query(members,"sample_id", search_db)

# create subject_id argument
subject_query = create_query(subject_id,"sample_id", search_db)

# create genes argument
genes_query = create_query(genes,"gene_name", annot_db)

# create genotype argument
gt_query = create_query(gt,"gt", search_db)

# create chromosome argument
chrom_query = create_query(chrom,"chrom", search_db)

# list to store final query object
final_query = []

# merge query arguments
def merge_query(query_arg):
    if len(query_arg) > 0:
        final_query.append(query_arg)

# create list to apply function on   
query_list = [member_query, subject_query, genes_query, gt_query,chrom_query]

# apply function to list of query statements
for query in query_list:
     merge_query(query)

# create final impala query, replacing first item with WHERE instead of AND
impala_query= ' '.join(final_query).replace('AND', 'WHERE', 1)

Then I run the following to determine which columns to return, and write the beginning of the query:

# function to create list of columns to return
def get_cols(db, col_arg):
    if col_arg == 'all':
        col_query = '{0}.*'.format(db)
    else:
        col_list = col_arg.replace("'", "").split(',')
        if len(col_list) > 1:
            for col in col_list:
                cols = ["{0}.".format(db) + s for s in col_list]
                col_query = ",".join(map(str, cols))
        elif len(col_list) == 1:
            col_query = "{0}.".format(db) + ",".join(map(str, col_list))
        else:
            print "Check your search_db_cols and annot_db_cols arguments and try again please."
            col_query = ''
    return col_query

 # merge all query statements 
col_statement = get_cols(search_db, search_db_cols) + ", " + get_cols(annot_db, annot_db_cols)

# write beginning of query
query_start = "SELECT " + col_statement + " FROM {0},{1} {2}".format(search_db, annot_db, impala_query)
\$\endgroup\$
2
\$\begingroup\$

First of all, I apologize if this is incomplete - I'm working from my laptop and I tend to code worse with fewer and smaller screens.

That being said, there is one huge issue I see with your code

Don't trust users. Ever

I don't see anything in your code that would prevent a SQL injection attack - for example, if I were to supply you with arguments like

args = {
    'members': 'all',
    'subject_id': 'all',
    'genes': 'BRCA1,TRT1,SON%,HOX%',
    'gt': 'hom',
    'chrom': '8',
    'search_db': 'wgs.comgen_variant; BEGIN TRANSACTION badness; DROP TABLE IMPORTANT_TABLE; COMMIT TRANSACTION badness;',
    'search_db_cols': 'all', 
    'annot_db': 'ref.ensembl_genes',
    'annot_db_cols': 'gene_name'
}

you'd probably have a bad time. I'm sure that there are much more subtle ways to achieve similarly nasty results, but this is something to keep in mind. I don't have a good suggestion on how to prevent SQL injection - I strongly prefer using libraries where they handle that and I just have to call their methods properly. If possible I'd work on leveraging existing code and techniques to avoid this problem.

Other stuff

Naming

In your find_wildcards function you name your input parameter x which tells me absolutely nothing about what it is or what it should be. Consider a more useful name.

Documentation

It will greatly help you, and any future contributors, if you include documentation on what each function is supposed to do - I recommend adding a docstring to every function, as well as a module level one.

String formatting

You use a weird mixture of string concatenation and string.format() to build your strings - pick one (and it better be formatting) and stick with it. For example, make

query_args = 'AND {0}.{1} = '.format(arg_db, name_arg) + "'" + ", ".join(str(e) for e in arg_list) + "'"

into

query_args = "AND {}.{} = '{}'".format(arg_db, name_arg, ', '.join(str(e) for e in arg_list))

Note that you also don't need to specify {index} - I generally find it just adds noise unless you want the same value inserted in multiple places.

Don't make unnecessary intermediate lists

Sometimes you use a list comprehension (which is good) and then map a function to that list. This is okay, but it could be better. Consider

conditions = ['{0}.{1} = '.format(arg_db, name_arg) + "'" + s + "'" for s in arg_list]
query_args = "AND (" + " OR ".join(map(str, conditions)) + ")"

This could be easily rewritten as

query_args = "AND ({})".format(" OR ".join("{}.{} = '{}'".format(arg_db, name_arg, s) for s in arg_list))

Some people might say that's hard to read - I'm not one of them, but I get the point. Thank god for indentation.

query_args = "AND ({})".format(" OR ".join(
    "{}.{} = '{}'".format(arg_db, name_arg, s)
    for s in arg_list))

Something else you could do - the form "{}.{}".format(arg_db, name_arg) is everywhere. Instead of repeating that formatting everywhere, save that as a local variable and insert it as needed.

There looks like a lot of repetition here. You could probably reduce that by being clever with function parameters and string formatting.

Finally, this section is pretty hard to understand. You wouldn't go wrong with some comments explaining why you're doing the things you do.

Simply booleans

Instead of

if len(query_arg) > 0:

just do

if query_arg:

Simplify and explain get_cols

This function is also hard to read and follow. Adding comments and documentation would go a long way to help with this

Make a main function

At the end you have quite a few steps required to actually get the query you want - it would be much easier to write a wrapper function that puts all of that work together and outputs the final query.

\$\endgroup\$
  • \$\begingroup\$ Thank you so much, this was extremely helpful! I really appreciate you taking the time to give me such thorough feedback. I'll try implementing your suggestions and get back to you if I come up with any questions/solutions that might be useful to the discussion! \$\endgroup\$ – SummerEla Jul 26 '15 at 17:14

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.