Advanced SQL Templates In Python with JinjaSql


In A Simple Approach To Templated SQL Queries In Python, I introduced the basics of SQL templates in Python using JinjaSql. This post further demonstrates the power of Jinja2 within JinjaSql templates using presets, loops, and custom functions. Let’s consider an everyday use case when we have a table with some dimensions and some numerical values, and we want to find some metrics for a given dimension or a combination of dimensions. The example data below is tiny and entirely made up, but it suffices for demonstrating the advanced features. First, I introduce the data set and the questions to be answered with SQL queries. I’ll then build the queries without templates, and finally, show how to use the SQL templates to parameterize and generate these queries. To generate SQL from JinjaSql templates, I’ll use (and modify) the apply_sql_template function introduced in the previous blog and available on GitHub in sql_tempates_base.py.

Example data set

Let’s consider a table transactions that contains records about purchases in some stores. The purchase can be made by cash, with a credit or a debit card, which adds an extra dimension to the data. Here is the code for creating and populating the table.

create table transactions (
    transaction_id int,
    user_id int,
    transaction_date date,
    store_id int,
    payment_method varchar(10),
    amount float
)
;

insert into transactions
(transaction_id, user_id, transaction_date, store_id, payment_method, amount)
values
    (1, 1234, '2019-03-02', 1, 'cash', 5.25),
    (1, 1234, '2019-03-01', 1, 'credit', 10.75),
    (1, 1234, '2019-03-02', 2, 'cash', 25.50),
    (1, 1234, '2019-03-03', 2, 'credit', 17.00),
    (1, 4321, '2019-03-01', 2, 'cash', 20.00),
    (1, 4321, '2019-03-02', 2, 'debit', 30.00),
    (1, 4321, '2019-03-03', 1, 'cash', 3.00)
;

Metrics to compute

When exploring a data set, it is common to look at the main performance metrics across all dimensions. In this example, we want to compute the following metrics:

  • number of transactions
  • average transaction amount
  • the total amount of transactions

We want these metrics for each user, store, and payment method. We also want to look at these metrics by store and payment method together.

Template for a single dimension

The query to obtain the metrics for each store is:

select
    store_id
    , count(*) as num_transactions
    , sum(amount) as total_amount
    , avg(amount) as avg_amount
from
    transactions
group by
    store_id
order by total_amount desc

To get the same metrics for other dimensions, we only need to change the store_id into user_id or payment_method in both the select and group by clauses. So the JinjaSql template may look like

_BASIC_STATS_TEMPLATE = '''
select
    {{ dim | sqlsafe }}
    , count(*) as num_transactions
    , sum(amount) as total_amount
    , avg(amount) as avg_amount
from
    transactions
group by
    {{ dim | sqlsafe }}
order by total_amount desc
'''

with parameters, for example, as

params = {
    'dim': 'payment_method'
}
sql = apply_sql_template(_BASIC_STATS_TEMPLATE, params)

The above template works for a single dimension, but what if we have more than one? To generate a generic query that works with any number of dimensions, let’s create a skeleton of a function that takes a list of dimension column names as an argument and returns the SQL.

def get_basic_stats_sql(dimensions):
  '''
  Returns the sql computing the number of transactions,
  as well as the total and the average transaction amounts
  for the provided list of column names as dimensions.
  '''
  # TODO: construct params
  return apply_sql_template(_BASIC_STATS_TEMPLATE, params)

Essentially, we want to transform a list of column names, such as ['payment_method'] or ['store_id', 'payment_method'] into a single string containing the column names as a comma-separated list. Here we have some options, as it can be done either in Python or in the template.

Passing a string generated outside the template

The first option is to generate the comma-separated string before passing it to the template. We can do it simply by joining the members of the list together:

def get_basic_stats_sql(dimensions):
    '''
    Returns the sql computing the number of transactions,
    as well as the total and the average transaction amounts
    for the provided list of column names as dimensions.
    '''
    params = {
      'dim': '\n    , '.join(dimensions)
    }
    return apply_sql_template(_BASIC_STATS_TEMPLATE, params)

It so happens that the template parameter dim is in the right place, so the resulting query is

>>> print(get_basic_stats_sql(['store_id', 'payment_method']))
select
    store_id
    , payment_method
    , count(*) as num_transactions
    , sum(amount) as total_amount
    , avg(amount) as avg_amount
from
    transactions
group by
    store_id
    , payment_method
order by total_amount desc

Now we can quickly generate SQL queries for all desired dimensions using

dimension_lists = [
    ['user_id'],
    ['store_id'],
    ['payment_method'],
    ['store_id', 'payment_method'],
]

dimension_queries = [get_basic_stats_sql(dims) for dims in dimension_lists]

Preset variables inside the template

An alternative to passing a pre-built string as a template parameter is to move the column list SQL generation into the template itself by setting a new variable at the top:

_PRESET_VAR_STATS_TEMPLATE = '''
{% set dims = '\n    , '.join(dimensions) %}
select
    {{ dims | sqlsafe }}
    , count(*) as num_transactions
    , sum(amount) as total_amount
    , avg(amount) as avg_amount
from
    transactions
group by
    {{ dims | sqlsafe }}
order by total_amount desc
'''

This template is more readable than the previous version since all transformations happen in one place in the template, and at the same time, there’s no clutter. The function should change to

def get_stats_sql(dimensions):
    '''
    Returns the sql computing the number of transactions,
    as well as the total and the average transaction amounts
    for the provided list of column names as dimensions.
    '''
    params = {
      'dimensions': dimensions
    }
    return apply_sql_template(_PRESET_VAR_STATS_TEMPLATE, params)

Loops inside the template

We can also use loops inside the template to generate the columns.

_LOOPS_STATS_TEMPLATE = '''
select
    {{ dimensions[0] | sqlsafe }}\
    {% for dim in dimensions[1:] %}
    , {{ dim | sqlsafe }}{% endfor %}
    , count(*) as num_transactions
    , sum(amount) as total_amount
    , avg(amount) as avg_amount
from
    transactions
group by
    {{ dimensions[0] | sqlsafe }}\
    {% for dim in dimensions[1:] %}
    , {{ dim | sqlsafe }}{% endfor %}
order by total_amount desc
'''

This example may not be the best use of the loops because a preset variable does the job just fine without the extra complexity. However, loops are a powerful construct, especially when there is additional logic inside the loop, such as conditions ({% if ... %} - {% endif %}) or nested loops.

So what is happening in the template above? The first element of the list dimensions[0] stands alone because it doesn’t need a comma in front of the column name. We wouldn’t need that if there were a defined first column in the query, and the for-loop would look simply as

    {% for dim in dimensions %}
    , {{ dim | sqlsafe }}
    {% endfor %}

Then, the for-loop construct goes over the remaining elements dimensions[1:]. The same code appears in the group by clause, which is also not ideal and only serves the purpose of showing the loop functionality.

One may wonder why the formatting of the loop is so strange. The reason is that the flow elements of the SQL template, such as {% endfor %}, generate a blank line if they appear on a separate line. To avoid that, in the template above, both {% for ... %} and {% endfor %} are technically on the same line as the previous code (hence the backslash \ after the first column name). SQL, of course, doesn’t care about whitespace, but humans who read SQL may (and should) care. An alternative to fighting with formatting within the template that would make it more readable is to strip the blank lines from the generated query before printing or logging it. A useful function for that purpose is

import os
def strip_blank_lines(text):
    '''
    Removes blank lines from the text, including those containing only spaces.
    https://stackoverflow.com/questions/1140958/whats-a-quick-one-liner-to-remove-empty-lines-from-a-python-string
    '''
    return os.linesep.join([s for s in text.splitlines() if s.strip()])

A better-formatted template then becomes

_LOOPS_STATS_TEMPLATE = '''
select
    {{ dimensions[0] | sqlsafe }}
    {% for dim in dimensions[1:] %}
    , {{ dim | sqlsafe }}
    {% endfor %}
    , count(*) as num_transactions
    , sum(amount) as total_amount
    , avg(amount) as avg_amount
from
    transactions
group by
    {{ dimensions[0] | sqlsafe }}
    {% for dim in dimensions[1:] %}
    , {{ dim | sqlsafe }}
    {% endfor %}
order by total_amount desc
'''

And the call to print the query is

print(strip_blank_lines(get_loops_stats_sql(['store_id', 'payment_method'])))

All the SQL templates so far used a list of dimensions to produce precisely the same query.

Custom dimensions with looping over a dictionary

In the loop example above, we see how to iterate over a list. It is also possible to iterate over a dictionary. This comes in handy, for example, when we want to alias or transform some or all of the columns that form dimensions. Suppose we wanted to combine the debit and credit cards as a single value and compare it to cash transactions. We can accomplish that by first creating a dictionary defining a transformation for the payment_method and keeping the store_id unchanged.

custom_dimensions = {
    'store_id': 'store_id',
    'card_or_cash': "case when payment_method = 'cash' then 'cash' else 'card' end",
}

Here, both credit and debit values are replaced with card. Then, the template may look like the following:

_CUSTOM_STATS_TEMPLATE = '''
{% set dims = '\n    , '.join(dimensions.keys()) %}
select
    sum(amount) as total_amount
    {% for dim, def in dimensions.items() %}
    , {{ def | sqlsafe }} as {{ dim | sqlsafe }}
    {% endfor %}
    , count(*) as num_transactions
    , avg(amount) as avg_amount
from
    transactions
group by
    {{ dims | sqlsafe }}
order by total_amount desc
'''

Note that I moved the total_amount as the first column just to simplify this example and avoid having to deal with the first element in the loop separately. Also, note that the group by clause is using a preset variable and is different from the code in the select query because it only lists the names of the generated columns. The resulting SQL query is

>>> print(strip_blank_lines(
...     apply_sql_template(template=_CUSTOM_STATS_TEMPLATE,
...                        parameters={'dimensions': custom_dimensions})))
select
    sum(amount) as total_amount
    , store_id as store_id
    , case when payment_method = 'cash' then 'cash' else 'card' end as card_or_cash
    , count(*) as num_transactions
    , avg(amount) as avg_amount
from
    transactions
group by
    store_id
    , card_or_cash
order by total_amount desc

Calling custom Python functions from within JinjaSql templates

What if we want to use a Python function to generate a portion of the code? Jinja2 allows one to register custom functions and functions from other packages for use within the SQL templates. Let’s start with defining a function that generates the string that we insert into the SQL for transforming custom dimensions.

def transform_dimensions(dimensions: dict) -> str:
    '''
    Generate SQL for aliasing or transforming the dimension columns.
    '''
    return '\n    , '.join([
        '{val} as {key}'.format(val=val, key=key)
        for key, val in dimensions.items()
    ])

The output of this function is what we expect to appear in the select clause:

>>> print(transform_dimensions(custom_dimensions))
store_id as store_id
    , case when payment_method = 'cash' then 'cash' else 'card' end as card_or_cash

Now we need to register this function with Jinja2. To do that, we’ll modify the apply_sql_template function from the previous blog as follows.

from jinjasql import JinjaSql

def apply_sql_template(template, parameters, func_list=None):
    '''
    Apply a JinjaSql template (string) substituting parameters (dict) and return
    the final SQL. Use the func_list to pass any functions called from the template.
    '''
    j = JinjaSql(param_style='pyformat')
    if func_list:
        for func in func_list:
            j.env.globals[func.__name__] = func
    query, bind_params = j.prepare_query(template, parameters)
    return get_sql_from_template(query, bind_params)

This version has an additional optional argument func_list that needs to be a list of functions.

Let’s change the template to take advantage of the transform_dimensions function.

_FUNCTION_STATS_TEMPLATE = '''
{% set dims = '\n    , '.join(dimensions.keys()) %}
select
    {{ transform_dimensions(dimensions) | sqlsafe }}
    , sum(amount) as total_amount
    , count(*) as num_transactions
    , avg(amount) as avg_amount
from
    transactions
group by
    {{ dims | sqlsafe }}
order by total_amount desc
'''

Now we also don’t need to worry about the first column not having a comma. The following call produces a SQL query similar to that in the previous section.

>>> print(strip_blank_lines(
...     apply_sql_template(template=_FUNCTION_STATS_TEMPLATE,
...                        parameters={'dimensions': custom_dimensions},
...                        func_list=[transform_dimensions])))

select
    store_id as store_id
    , case when payment_method = 'cash' then 'cash' else 'card' end as card_or_cash
    , sum(amount) as total_amount
    , count(*) as num_transactions
    , avg(amount) as avg_amount
from
    transactions
group by
    store_id
    , card_or_cash
order by total_amount desc

Note how we pass transform_dimensions to apply_sql_template as a list [transform_dimensions]. Multiple functions can be passed into SQL templates as a list of functions, for example, [func1, func2].

Conclusion

This tutorial is a follow up to my first post on the basic use of JinjaSql. It demonstrates the use of preset variables, loops over lists and dictionaries, and custom Python functions within JinjaSql templates for advanced SQL code generation. In particular, the addition of custom functions registration to the apply_sql_template function makes templating much more powerful and versatile. Parameterized SQL queries continue to be indispensable for automated report generation and for reducing the amount of SQL code that needs to be maintained. An added benefit is that with reusable SQL code snippets, it becomes easier to use standard Python unit testing techniques to verify that the generated SQL is correct.

The code in this post is licensed under the MIT License.

Photo by Sergei Izrailev

Copyright (c) 2018-2020 Sergei Izrailev. All opinions are my own.