def test_function_calls(self):
self.assert_compiled_select(
'SELECT ABS(-3), POW(2, 3), NOW()',
typed_ast.Select([
typed_ast.SelectField(
typed_ast.FunctionCall(runtime.get_func('abs'), [
typed_ast.FunctionCall(
runtime.get_unary_op('-'),
[typed_ast.Literal(3, tq_types.INT)], tq_types.INT)
], tq_types.INT), 'f0_', None),
typed_ast.SelectField(
typed_ast.FunctionCall(runtime.get_func('pow'), [
typed_ast.Literal(2, tq_types.INT),
typed_ast.Literal(3, tq_types.INT)
], tq_types.INT), 'f1_', None),
typed_ast.SelectField(
typed_ast.FunctionCall(runtime.get_func('now'), [],
tq_types.INT), 'f2_', None)
], typed_ast.NoTable(), typed_ast.Literal(True,
tq_types.BOOL), None,
typed_ast.Literal(True, tq_types.BOOL), None,
None,
self.make_type_context(
[(None, 'f0_', tq_types.INT),
(None, 'f1_', tq_types.INT),
(None, 'f2_', tq_types.INT)],
self.make_type_context([]))))
def test_aggregates(self):
self.assert_compiled_select(
"SELECT MAX(value), MIN(value) FROM table1",
typed_ast.Select(
[
typed_ast.SelectField(
typed_ast.AggregateFunctionCall(
runtime.get_func("max"),
[typed_ast.ColumnRef("table1", "value", tq_types.INT)],
tq_types.INT,
),
"f0_",
),
typed_ast.SelectField(
typed_ast.AggregateFunctionCall(
runtime.get_func("min"),
[typed_ast.ColumnRef("table1", "value", tq_types.INT)],
tq_types.INT,
),
"f1_",
),
],
typed_ast.Table("table1", self.table1_type_ctx),
typed_ast.Literal(True, tq_types.BOOL),
typed_ast.GroupSet(set(), []),
None,
self.make_type_context(
[(None, "f0_", tq_types.INT), (None, "f1_", tq_types.INT)], self.make_type_context([])
),
),
)
def test_select_grouped_and_non_grouped_fields(self):
self.assert_compiled_select(
"SELECT value, SUM(value2) FROM table1 GROUP BY value",
typed_ast.Select(
[
typed_ast.SelectField(typed_ast.ColumnRef("table1", "value", tq_types.INT), "value"),
typed_ast.SelectField(
typed_ast.FunctionCall(
runtime.get_func("sum"),
[typed_ast.ColumnRef("table1", "value2", tq_types.INT)],
tq_types.INT,
),
"f0_",
),
],
typed_ast.Table("table1", self.table1_type_ctx),
typed_ast.Literal(True, tq_types.BOOL),
typed_ast.GroupSet(alias_groups={"value"}, field_groups=[]),
None,
self.make_type_context(
[(None, "value", tq_types.INT), (None, "f0_", tq_types.INT)],
self.make_type_context([("table1", "value", tq_types.INT)]),
),
),
)
def compile_FunctionCall(self, expr, type_ctx):
# Innermost aggregates are special, since the context to use changes
# inside them. We also need to generate an AggregateFunctionCall AST so
# that the evaluator knows to change the context.
if self.is_innermost_aggregate(expr):
if type_ctx.aggregate_context is None:
raise exceptions.CompileError('Unexpected aggregate function.')
sub_expr_ctx = type_ctx.aggregate_context
ast_type = typed_ast.AggregateFunctionCall
else:
sub_expr_ctx = type_ctx
ast_type = typed_ast.FunctionCall
func = runtime.get_func(expr.name)
compiled_args = [
self.compile_expr(sub_expr, sub_expr_ctx) for sub_expr in expr.args
]
try:
result_type = func.check_types(*(arg.type
for arg in compiled_args))
except TypeError:
raise exceptions.CompileError(
'Invalid types for function {}: {}'.format(
expr.name, [arg.type for arg in compiled_args]))
return ast_type(func, compiled_args, result_type)
def test_within_clause(self):
self.assert_compiled_select(
'SELECT r1.s, COUNT(r1.s) WITHIN r1 AS num_s_in_r1 '
'FROM record_table',
typed_ast.Select(select_fields=[
typed_ast.SelectField(
typed_ast.ColumnRef('record_table', 'r1.s',
tq_types.STRING), 'r1.s', None),
typed_ast.SelectField(
typed_ast.FunctionCall(runtime.get_func('count'), [
typed_ast.ColumnRef('record_table', 'r1.s',
tq_types.STRING)
], tq_types.INT), 'num_s_in_r1', 'r1')
],
table=typed_ast.Table('record_table',
self.record_table_type_ctx),
where_expr=typed_ast.Literal(True, tq_types.BOOL),
group_set=typed_ast.GroupSet(set(), []),
having_expr=typed_ast.Literal(
True, tq_types.BOOL),
orderings=None,
limit=None,
type_ctx=self.make_type_context(
[(None, 'r1.s', tq_types.STRING),
(None, 'num_s_in_r1', tq_types.INT)],
self.make_type_context([]))))
def test_aggregates(self):
self.assert_compiled_select(
'SELECT MAX(value), MIN(value) FROM table1',
typed_ast.Select([
typed_ast.SelectField(
typed_ast.AggregateFunctionCall(
runtime.get_func('max'),
[typed_ast.ColumnRef('table1', 'value', tq_types.INT)],
tq_types.INT), 'f0_', None),
typed_ast.SelectField(
typed_ast.AggregateFunctionCall(
runtime.get_func('min'),
[typed_ast.ColumnRef('table1', 'value', tq_types.INT)],
tq_types.INT), 'f1_', None)
], typed_ast.Table('table1', self.table1_type_ctx),
typed_ast.Literal(True, tq_types.BOOL),
typed_ast.GroupSet(set(), []),
typed_ast.Literal(True, tq_types.BOOL), None,
None,
self.make_type_context(
[(None, 'f0_', tq_types.INT),
(None, 'f1_', tq_types.INT)],
self.make_type_context([]))))
def test_function_calls(self):
self.assert_compiled_select(
"SELECT ABS(-3), POW(2, 3), NOW()",
typed_ast.Select(
[
typed_ast.SelectField(
typed_ast.FunctionCall(
runtime.get_func("abs"),
[
typed_ast.FunctionCall(
runtime.get_unary_op("-"), [typed_ast.Literal(3, tq_types.INT)], tq_types.INT
)
],
tq_types.INT,
),
"f0_",
),
typed_ast.SelectField(
typed_ast.FunctionCall(
runtime.get_func("pow"),
[typed_ast.Literal(2, tq_types.INT), typed_ast.Literal(3, tq_types.INT)],
tq_types.INT,
),
"f1_",
),
typed_ast.SelectField(typed_ast.FunctionCall(runtime.get_func("now"), [], tq_types.INT), "f2_"),
],
typed_ast.NoTable(),
typed_ast.Literal(True, tq_types.BOOL),
None,
None,
self.make_type_context(
[(None, "f0_", tq_types.INT), (None, "f1_", tq_types.INT), (None, "f2_", tq_types.INT)],
self.make_type_context([]),
),
),
)
def test_case(self):
self.assert_compiled_select(
'SELECT CASE WHEN TRUE THEN 1 WHEN FALSE THEN 2 END',
typed_ast.Select(
select_fields=[
typed_ast.SelectField(
typed_ast.FunctionCall(runtime.get_func('if'), [
typed_ast.Literal(True, tq_types.BOOL),
typed_ast.Literal(1, tq_types.INT),
typed_ast.FunctionCall(runtime.get_func('if'), [
typed_ast.Literal(False, tq_types.BOOL),
typed_ast.Literal(2, tq_types.INT),
typed_ast.Literal(None, tq_types.NONETYPE),
], tq_types.INT)
], tq_types.INT), 'f0_', None)
],
table=typed_ast.NoTable(),
where_expr=typed_ast.Literal(True, tq_types.BOOL),
group_set=None,
having_expr=typed_ast.Literal(True, tq_types.BOOL),
orderings=None,
limit=None,
type_ctx=self.make_type_context([(None, 'f0_', tq_types.INT)],
self.make_type_context([]))))
def test_select_grouped_and_non_grouped_fields(self):
self.assert_compiled_select(
'SELECT value, SUM(value2) FROM table1 GROUP BY value',
typed_ast.Select([
typed_ast.SelectField(
typed_ast.ColumnRef('table1', 'value', tq_types.INT),
'value', None),
typed_ast.SelectField(
typed_ast.FunctionCall(runtime.get_func('sum'), [
typed_ast.ColumnRef('table1', 'value2', tq_types.INT)
], tq_types.INT), 'f0_', None)
], typed_ast.Table('table1', self.table1_type_ctx),
typed_ast.Literal(True, tq_types.BOOL),
typed_ast.GroupSet(alias_groups={'value'},
field_groups=[]),
typed_ast.Literal(True, tq_types.BOOL), None,
None,
self.make_type_context(
[(None, 'value', tq_types.INT),
(None, 'f0_', tq_types.INT)],
self.make_type_context([('table1', 'value',
tq_types.INT)]))))
