Backport #71426 to 24.10: Fix crash in mongodb table function

This commit is contained in:
robot-clickhouse 2024-11-05 19:08:31 +00:00
parent 3ffd2ba35f
commit 1c737c48cb
6 changed files with 60 additions and 28 deletions

View File

@ -724,7 +724,10 @@ void ASTFunction::formatImplWithoutAlias(const FormatSettings & settings, Format
{
if (secret_arguments.are_named)
{
assert_cast<const ASTFunction *>(argument.get())->arguments->children[0]->formatImpl(settings, state, nested_dont_need_parens);
if (const auto * func_ast = typeid_cast<const ASTFunction *>(argument.get()))
func_ast->arguments->children[0]->formatImpl(settings, state, nested_dont_need_parens);
else
argument->formatImpl(settings, state, nested_dont_need_parens);
settings.ostr << (settings.hilite ? hilite_operator : "") << " = " << (settings.hilite ? hilite_none : "");
}
if (!secret_arguments.replacement.empty())

View File

@ -15,7 +15,7 @@
#include <TableFunctions/registerTableFunctions.h>
#include <Storages/checkAndGetLiteralArgument.h>
#include <Storages/ColumnsDescription.h>
#include <TableFunctions/TableFunctionMongoDB.h>
namespace DB
{
@ -85,17 +85,11 @@ void TableFunctionMongoDB::parseArguments(const ASTPtr & ast_function, ContextPt
{
if (const auto * ast_func = typeid_cast<const ASTFunction *>(args[i].get()))
{
const auto * args_expr = assert_cast<const ASTExpressionList *>(ast_func->arguments.get());
auto function_args = args_expr->children;
if (function_args.size() != 2)
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Expected key-value defined argument");
auto arg_name = function_args[0]->as<ASTIdentifier>()->name();
const auto & [arg_name, arg_value] = getKeyValueMongoDBArgument(ast_func);
if (arg_name == "structure")
structure = checkAndGetLiteralArgument<String>(function_args[1], "structure");
structure = checkAndGetLiteralArgument<String>(arg_value, arg_name);
else if (arg_name == "options")
main_arguments.push_back(function_args[1]);
main_arguments.push_back(arg_value);
}
else if (i == 5)
{
@ -117,15 +111,11 @@ void TableFunctionMongoDB::parseArguments(const ASTPtr & ast_function, ContextPt
{
if (const auto * ast_func = typeid_cast<const ASTFunction *>(args[i].get()))
{
const auto * args_expr = assert_cast<const ASTExpressionList *>(ast_func->arguments.get());
auto function_args = args_expr->children;
if (function_args.size() != 2)
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Expected key-value defined argument");
auto arg_name = function_args[0]->as<ASTIdentifier>()->name();
const auto & [arg_name, arg_value] = getKeyValueMongoDBArgument(ast_func);
if (arg_name == "structure")
structure = checkAndGetLiteralArgument<String>(function_args[1], "structure");
structure = checkAndGetLiteralArgument<String>(arg_value, arg_name);
else if (arg_name == "options")
main_arguments.push_back(arg_value);
}
else if (i == 2)
{
@ -145,6 +135,20 @@ void TableFunctionMongoDB::parseArguments(const ASTPtr & ast_function, ContextPt
}
std::pair<String, ASTPtr> getKeyValueMongoDBArgument(const ASTFunction * ast_func)
{
const auto * args_expr = assert_cast<const ASTExpressionList *>(ast_func->arguments.get());
const auto & function_args = args_expr->children;
if (function_args.size() != 2 || ast_func->name != "equals" || !function_args[0]->as<ASTIdentifier>())
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Expected key-value defined argument, got {}", ast_func->formatForErrorMessage());
const auto & arg_name = function_args[0]->as<ASTIdentifier>()->name();
if (arg_name == "structure" || arg_name == "options")
return std::make_pair(arg_name, function_args[1]);
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Expected key-value defined argument, got {}", ast_func->formatForErrorMessage());
}
void registerTableFunctionMongoDB(TableFunctionFactory & factory)
{
factory.registerFunction<TableFunctionMongoDB>(

View File

@ -0,0 +1,16 @@
#pragma once
#include <Common/Exception.h>
#include <Parsers/ASTFunction.h>
#include <Parsers/ASTIdentifier.h>
#include <Storages/checkAndGetLiteralArgument.h>
namespace DB
{
std::pair<String, ASTPtr> getKeyValueMongoDBArgument(const ASTFunction * ast_func);
}

View File

@ -15,6 +15,7 @@
#include <TableFunctions/registerTableFunctions.h>
#include <Storages/checkAndGetLiteralArgument.h>
#include <Storages/ColumnsDescription.h>
#include <TableFunctions/TableFunctionMongoDB.h>
namespace DB
@ -97,17 +98,11 @@ void TableFunctionMongoDBPocoLegacy::parseArguments(const ASTPtr & ast_function,
{
if (const auto * ast_func = typeid_cast<const ASTFunction *>(args[i].get()))
{
const auto * args_expr = assert_cast<const ASTExpressionList *>(ast_func->arguments.get());
auto function_args = args_expr->children;
if (function_args.size() != 2)
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Expected key-value defined argument");
auto arg_name = function_args[0]->as<ASTIdentifier>()->name();
const auto & [arg_name, arg_value] = getKeyValueMongoDBArgument(ast_func);
if (arg_name == "structure")
structure = checkAndGetLiteralArgument<String>(function_args[1], "structure");
structure = checkAndGetLiteralArgument<String>(arg_value, "structure");
else if (arg_name == "options")
main_arguments.push_back(function_args[1]);
main_arguments.push_back(arg_value);
}
else if (i == 5)
{

View File

@ -0,0 +1,14 @@
-- Tags: no-fasttest
SELECT * FROM mongodb('mongodb://some-cluster:27017/?retryWrites=false', NULL, 'my_collection', 'test_user', 'password', 'x Int32'); -- { serverError BAD_ARGUMENTS }
SELECT * FROM mongodb('mongodb://some-cluster:27017/?retryWrites=false', 'test', NULL, 'test_user', 'password', 'x Int32'); -- { serverError BAD_ARGUMENTS }
SELECT * FROM mongodb('mongodb://some-cluster:27017/?retryWrites=false', 'test', 'my_collection', NULL, 'password', 'x Int32'); -- { serverError BAD_ARGUMENTS }
SELECT * FROM mongodb('mongodb://some-cluster:27017/?retryWrites=false', 'test', 'my_collection', 'test_user', NULL, 'x Int32'); -- { serverError BAD_ARGUMENTS }
SELECT * FROM mongodb('mongodb://some-cluster:27017/?retryWrites=false', 'test', 'my_collection', 'test_user', 'password', NULL); -- { serverError BAD_ARGUMENTS }
SELECT * FROM mongodb('mongodb://some-cluster:27017/?retryWrites=false', 'test', 'my_collection', 'test_user', 'password', materialize(1) + 1); -- { serverError BAD_ARGUMENTS }
SELECT * FROM mongodb('mongodb://some-cluster:27017/?retryWrites=false', 'test', 'my_collection', 'test_user', 'password', 'x Int32', NULL); -- { serverError BAD_ARGUMENTS }
SELECT * FROM mongodb('mongodb://some-cluster:27017/?retryWrites=false', 'test', 'my_collection', 'test_user', 'password', NULL, 'x Int32'); -- { serverError BAD_ARGUMENTS }
SELECT * FROM mongodb('mongodb://some-cluster:27017/?retryWrites=false', 'test', 'my_collection', 'test_user', 'password', NULL, 'x Int32'); -- { serverError BAD_ARGUMENTS }
SELECT * FROM mongodb(NULL, 'test', 'my_collection', 'test_user', 'password', 'x Int32'); -- { serverError BAD_ARGUMENTS }
CREATE TABLE IF NOT EXISTS store_version ( `_id` String ) ENGINE = MongoDB(`localhost:27017`, mongodb, storeinfo, adminUser, adminUser); -- { serverError NAMED_COLLECTION_DOESNT_EXIST }