Skip to content

Commit

Permalink
Support parameters in table function arguments
Browse files Browse the repository at this point in the history
  • Loading branch information
kasiafi committed Jun 21, 2022
1 parent 058904c commit 0d2dc80
Show file tree
Hide file tree
Showing 7 changed files with 84 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,8 @@
import io.trino.sql.tree.Explain;
import io.trino.sql.tree.ExplainAnalyze;
import io.trino.sql.tree.Expression;
import io.trino.sql.tree.ExpressionRewriter;
import io.trino.sql.tree.ExpressionTreeRewriter;
import io.trino.sql.tree.FetchFirst;
import io.trino.sql.tree.FieldReference;
import io.trino.sql.tree.FrameBound;
Expand Down Expand Up @@ -1677,9 +1679,25 @@ else if (argument.getValue() instanceof Expression) {
if (argument.getValue() instanceof FunctionCall && ((FunctionCall) argument.getValue()).getName().hasSuffix(QualifiedName.of("descriptor"))) { // function name is always compared case-insensitive
throw semanticException(INVALID_FUNCTION_ARGUMENT, argument, "'descriptor' function is not allowed as a table function argument");
}
// inline parameters
Expression inlined = ExpressionTreeRewriter.rewriteWith(new ExpressionRewriter<>()
{
@Override
public Expression rewriteParameter(Parameter node, Void context, ExpressionTreeRewriter<Void> treeRewriter)
{
if (analysis.isDescribe()) {
// We cannot handle DESCRIBE when a table function argument involves a parameter.
// In DESCRIBE, the parameter values are not known. We cannot pass a dummy value for a parameter.
// The value of a table function argument can affect the returned relation type. The returned
// relation type can affect the assumed types for other parameters in the query.
throw semanticException(NOT_SUPPORTED, node, "DESCRIBE is not supported if a table function uses parameters");
}
return analysis.getParameters().get(NodeRef.of(node));
}
}, expression);
Type expectedArgumentType = ((ScalarArgumentSpecification) argumentSpecification).getType();
// currently, only constant arguments are supported
Object constantValue = ExpressionInterpreter.evaluateConstantExpression(expression, expectedArgumentType, plannerContext, session, accessControl, analysis.getParameters());
Object constantValue = ExpressionInterpreter.evaluateConstantExpression(inlined, expectedArgumentType, plannerContext, session, accessControl, analysis.getParameters());
return ScalarArgument.builder()
.type(expectedArgumentType)
.value(constantValue)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -977,4 +977,14 @@ protected Void visitJsonArray(JsonArray node, C context)

return null;
}

@Override
protected Void visitTableFunctionInvocation(TableFunctionInvocation node, C context)
{
for (TableFunctionArgument argument : node.getArguments()) {
process(argument.getValue(), context);
}

return null;
}
}
8 changes: 7 additions & 1 deletion docs/src/main/sphinx/functions/table.rst
Original file line number Diff line number Diff line change
Expand Up @@ -68,4 +68,10 @@ skipped arguments are declared with default values.
You cannot mix the argument conventions in one invocation.

All arguments must be constant expressions, and they can be of any SQL type,
which is compatible with the declared argument type.
which is compatible with the declared argument type. You can also use
parameters in arguments::

PREPARE stmt FROM
SELECT * FROM TABLE(my_function("row_count" => ? + 1, "column_count" => ?));

EXECUTE stmt USING 100, 1;
Original file line number Diff line number Diff line change
Expand Up @@ -1591,6 +1591,17 @@ public void testNativeQuerySimple()
assertQuery("SELECT * FROM TABLE(system.query(query => 'SELECT 1'))", "VALUES 1");
}

@Test
public void testNativeQueryParameters()
{
Session session = Session.builder(getSession())
.addPreparedStatement("my_query_simple", "SELECT * FROM TABLE(system.query(query => ?))")
.addPreparedStatement("my_query", "SELECT * FROM TABLE(system.query(query => format('SELECT %s FROM %s', ?, ?)))")
.build();
assertQuery(session, "EXECUTE my_query_simple USING 'SELECT 1 a'", "VALUES 1");
assertQuery(session, "EXECUTE my_query USING 'a', '(SELECT 2 a) t'", "VALUES 2");
}

@Test
public void testNativeQuerySelectFromNation()
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import io.trino.Session;
import io.trino.plugin.jdbc.BaseJdbcConnectorTest;
import io.trino.sql.planner.plan.AggregationNode;
import io.trino.testing.MaterializedResult;
Expand Down Expand Up @@ -582,6 +583,18 @@ public void testNativeQuerySimple()
assertQueryFails("SELECT * FROM TABLE(system.query(query => 'SELECT 1'))", "line 1:21: Table function system.query not registered");
}

@Override
public void testNativeQueryParameters()
{
// table function disabled for ClickHouse, because it doesn't provide ResultSetMetaData, so the result relation type cannot be determined
Session session = Session.builder(getSession())
.addPreparedStatement("my_query_simple", "SELECT * FROM TABLE(system.query(query => ?))")
.addPreparedStatement("my_query", "SELECT * FROM TABLE(system.query(query => format('SELECT %s FROM %s', ?, ?)))")
.build();
assertQueryFails(session, "EXECUTE my_query_simple USING 'SELECT 1 a'", "line 1:21: Table function system.query not registered");
assertQueryFails(session, "EXECUTE my_query USING 'a', '(SELECT 2 a) t'", "line 1:21: Table function system.query not registered");
}

@Override
public void testNativeQuerySelectFromNation()
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -426,6 +426,18 @@ public void testNativeQuerySimple()
assertQuery("SELECT * FROM TABLE(system.query(query => 'SELECT CAST(1 AS number(2, 1)) FROM DUAL'))", ("VALUES 1"));
}

@Override
public void testNativeQueryParameters()
{
// override because Oracle requires the FROM clause, and it needs explicit type
Session session = Session.builder(getSession())
.addPreparedStatement("my_query_simple", "SELECT * FROM TABLE(system.query(query => ?))")
.addPreparedStatement("my_query", "SELECT * FROM TABLE(system.query(query => format('SELECT %s FROM %s', ?, ?)))")
.build();
assertQuery(session, "EXECUTE my_query_simple USING 'SELECT CAST(1 AS number(2, 1)) a FROM DUAL'", "VALUES 1");
assertQuery(session, "EXECUTE my_query USING 'a', '(SELECT CAST(2 AS number(2, 1)) a FROM DUAL) t'", "VALUES 2");
}

@Override
public void testNativeQueryInsertStatementTableDoesNotExist()
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -572,6 +572,18 @@ public void testNativeQuerySimple()
assertQueryFails("SELECT * FROM TABLE(system.query(query => 'SELECT 1'))", "line 1:21: Table function system.query not registered");
}

@Override
public void testNativeQueryParameters()
{
// not implemented
Session session = Session.builder(getSession())
.addPreparedStatement("my_query_simple", "SELECT * FROM TABLE(system.query(query => ?))")
.addPreparedStatement("my_query", "SELECT * FROM TABLE(system.query(query => format('SELECT %s FROM %s', ?, ?)))")
.build();
assertQueryFails(session, "EXECUTE my_query_simple USING 'SELECT 1 a'", "line 1:21: Table function system.query not registered");
assertQueryFails(session, "EXECUTE my_query USING 'a', '(SELECT 2 a) t'", "line 1:21: Table function system.query not registered");
}

@Override
public void testNativeQuerySelectFromNation()
{
Expand Down

0 comments on commit 0d2dc80

Please sign in to comment.