Queries

Compose the final SQL query.

---

source

Table

 Table (name)

A class with star and as_ methods to be used as a Table/AliasesQuery in a query. Any other attribute will be treated as a field.

vw = Table('vw')
test_eq(vw.as_('a').get_sql(), 'vw as a')
test_eq((vw.column + 2 > 1).get_sql(), 'a.column + 2 > 1')
test_eq(vw.star().get_sql(), 'a.*')

tbl = Table('PROD_BI.PRES.tbl')
test_eq(tbl.get_sql(), 'PROD_BI.PRES.tbl')
test_eq(tbl.as_('t').get_sql(), 'PROD_BI.PRES.tbl as t')
test_eq(tbl.as_('t').get_sql(quote_char='"'), '"PROD_BI"."PRES"."tbl" as "t"')
test_eq((tbl.column + 2 > 1).get_sql(), 't.column + 2 > 1')
test_eq((tbl.column + 2 > 1).get_sql(quote_char='"'), '"t"."column" + 2 > 1')

---

source

Query

 Query ()

The main class to construct a query. It returns a SelectQuery object when called with the classmethod from_ or with_. One can extend this class to add more classmethods to construct different types of queries.

---

source

UnionQuery

 UnionQuery (q1, q2, union_type='UNION')

The class to construct a union query.

---

source

SelectQuery

 SelectQuery ()

The class to construct a select query. It returns a Joiner object when called with the method join. It returns a Selector object when called with the method select.

---

source

Selector

 Selector (select_query, *args)

Select clause could be followed by a distinct clause

---

source

Joiner

 Joiner (select_query, query, how=None)

Join clause has to be followed by an on clause

---

source

QueryBase

 QueryBase ()

An empty query class to be inherited by all query classes. A convenient tool to make all query objects belong to this same class.

q0 = (Query
      .from_(Table('tbl').as_('a'))
      .select('col1')
      .where((Field('col2')-100>2) & (Field('col3')/9<=1))
      .limit(100)
)
test_eq(q0.get_sql(), 
"""select col1
from tbl as a
where col2 - 100 > 2 and col3 / 9 <= 1
limit 100""")

qj = (Query
     .with_(Query.from_(Table('tbl').as_('a')).select('col1'), 's')     
     .with_(q0, 'm')
     .from_('s')
     .join('m').on('s.col1=m.col1')
     .where(Field('col1')>=10)
     .select('s.col1', 'm.col2', 'm.col3')
)
test_eq(qj.get_sql(), 
"""with s as (
select col1
from tbl as a)

, m as (
select col1
from tbl as a
where col2 - 100 > 2 and col3 / 9 <= 1
limit 100)

select s.col1, m.col2, m.col3
from s
join m on s.col1=m.col1
where col1 >= 10""")

qd = (Query
     .from_('s')
     .where(Field('col1')>=10)
     .select('col1', 'col2', 'col3').distinct()
     .orderby('col1', 'col2')
)
test_eq(qd.get_sql(),
"""select distinct col1, col2, col3
from s
where col1 >= 10
order by col1, col2""")

---

source

Exists

 Exists (query:__main__.Query)

Exists statement

tbl = Table('tbl')
test_eq(Exists(Query.from_(tbl)
               .select(tbl.col1)
               .where(tbl.col2>10)).get_sql(quote_char='"'), 
        'EXISTS (select "tbl"."col1"\nfrom "tbl"\nwhere "tbl"."col2" > 10)')
test_eq(Query.from_(tbl)
        .select(tbl.col1)
        .where(
            Exists(Query.from_(tbl).
                   select(tbl.col1)
                   .where(tbl.col2>10))
        ).get_sql(quote_char='"'),
        'select "tbl"."col1"\nfrom "tbl"\nwhere EXISTS (select "tbl"."col1"\nfrom "tbl"\nwhere "tbl"."col2" > 10)')

Extending Query

We can extend SelectQuery to support more complex queries. There are 3 things we need to implement: 1. A method to record all necessary information for the SQL clause. Let’s call it custom for example. 2. A method to generate the SQL clause str. This method’s name has to be parse_ + the previous method’s name (parse_custom). 3. The class variable sql_keys has to be overwritten to include the new method’s name in the appropriate position. The order of the keys in this list determines the order of the SQL clauses in the final SQL str.

For example, for Snowflake SQL, to generate the PIVOT clause, we need to implement both a pivot method and a parse_pivot method, and also add pivot into sql_keys at the right place.

class SFSelectQuery(SelectQuery):
    # the order to put together the final sql query
    sql_keys = ['with', 'select', 'from', 'join', 'pivot', 'where', 'groupby', 'having', 'orderby', 'limit']

    def __init__(self) -> None:
        super().__init__()

    def pivot(self, agg_func, pivot_col, value_col, pivot_values, alias):
        self.dic['pivot'] = {
            'agg_func': agg_func,
            'pivot_col': pivot_col,
            'value_col': value_col,
            'pivot_values': pivot_values,
            'alias': alias
        }
        return self

    def parse_pivot(self, **kwargs):
        dialect = kwargs.get('dialect', None)
        if dialect == 'snowflake':
            dic = self.dic['pivot']
            agg_func = dic['agg_func']
            pivot_col = dic['pivot_col']
            value_col = dic['value_col']
            pivot_values = dic['pivot_values']
            alias = dic['alias']
            pivot_values = ', '.join([f"'{v}'" for v in pivot_values])
            return f"pivot({execute(agg_func(value_col), **kwargs)} for {execute(pivot_col, **kwargs)} in ({pivot_values})) as {alias}"
        else:
            raise NotImplementedError(f"dialect {dialect} not implemented")

To use this new SFSelectQuery in Query, we simply assign it to the class variable q. And we can now use .pivot in our query.

Query.q = SFSelectQuery

vw = Table('vw')
print(
    Query
    .from_(vw)
    .pivot(fn.Sum, vw.amount, vw.month, ['JAN', 'FEB', 'MAR', 'APR'], 'p')
    .where(vw.column > 1).select(vw.star()).get_sql(dialect='snowflake')
)

select vw.*
from vw
pivot(SUM(vw.month) for vw.amount in ('JAN', 'FEB', 'MAR', 'APR')) as p
where vw.column > 1