Merge pull request #68534 from yariks5s/fix_functionSQLJSON

Fix `Context has expired` while executing query
This commit is contained in:
Yarik Briukhovetskyi 2024-08-22 12:15:21 +00:00 committed by GitHub
commit 7bae85ff08
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -123,7 +123,7 @@ public:
class Executor class Executor
{ {
public: public:
static ColumnPtr run(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count, uint32_t parse_depth, uint32_t parse_backtracks, const ContextPtr & context) static ColumnPtr run(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count, uint32_t parse_depth, uint32_t parse_backtracks, bool function_json_value_return_type_allow_complex)
{ {
MutableColumnPtr to{result_type->createColumn()}; MutableColumnPtr to{result_type->createColumn()};
to->reserve(input_rows_count); to->reserve(input_rows_count);
@ -191,7 +191,7 @@ public:
{ {
/// Instead of creating a new generator for each row, we can reuse the same one. /// Instead of creating a new generator for each row, we can reuse the same one.
generator_json_path.reinitialize(); generator_json_path.reinitialize();
added_to_column = impl.insertResultToColumn(*to, document, generator_json_path, context); added_to_column = impl.insertResultToColumn(*to, document, generator_json_path, function_json_value_return_type_allow_complex);
} }
if (!added_to_column) if (!added_to_column)
{ {
@ -204,11 +204,18 @@ public:
}; };
template <typename Name, template <typename, typename> typename Impl> template <typename Name, template <typename, typename> typename Impl>
class FunctionSQLJSON : public IFunction, WithConstContext class FunctionSQLJSON : public IFunction
{ {
public: public:
static FunctionPtr create(ContextPtr context_) { return std::make_shared<FunctionSQLJSON>(context_); } static FunctionPtr create(ContextPtr context_) { return std::make_shared<FunctionSQLJSON>(context_); }
explicit FunctionSQLJSON(ContextPtr context_) : WithConstContext(context_) { } explicit FunctionSQLJSON(ContextPtr context_)
: max_parser_depth(context_->getSettingsRef().max_parser_depth),
max_parser_backtracks(context_->getSettingsRef().max_parser_backtracks),
allow_simdjson(context_->getSettingsRef().allow_simdjson),
function_json_value_return_type_allow_complex(context_->getSettingsRef().function_json_value_return_type_allow_complex),
function_json_value_return_type_allow_nullable(context_->getSettingsRef().function_json_value_return_type_allow_nullable)
{
}
static constexpr auto name = Name::name; static constexpr auto name = Name::name;
String getName() const override { return Name::name; } String getName() const override { return Name::name; }
@ -221,7 +228,7 @@ public:
DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
{ {
return Impl<DummyJSONParser, DefaultJSONStringSerializer<DummyJSONParser::Element>>::getReturnType( return Impl<DummyJSONParser, DefaultJSONStringSerializer<DummyJSONParser::Element>>::getReturnType(
Name::name, arguments, getContext()); Name::name, arguments, function_json_value_return_type_allow_nullable);
} }
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override
@ -231,19 +238,25 @@ public:
/// 2. Create ASTPtr /// 2. Create ASTPtr
/// 3. Parser(Tokens, ASTPtr) -> complete AST /// 3. Parser(Tokens, ASTPtr) -> complete AST
/// 4. Execute functions: call getNextItem on generator and handle each item /// 4. Execute functions: call getNextItem on generator and handle each item
unsigned parse_depth = static_cast<unsigned>(getContext()->getSettingsRef().max_parser_depth); unsigned parse_depth = static_cast<unsigned>(max_parser_depth);
unsigned parse_backtracks = static_cast<unsigned>(getContext()->getSettingsRef().max_parser_backtracks); unsigned parse_backtracks = static_cast<unsigned>(max_parser_backtracks);
#if USE_SIMDJSON #if USE_SIMDJSON
if (getContext()->getSettingsRef().allow_simdjson) if (allow_simdjson)
return FunctionSQLJSONHelpers::Executor< return FunctionSQLJSONHelpers::Executor<
Name, Name,
Impl<SimdJSONParser, JSONStringSerializer<SimdJSONParser::Element, SimdJSONElementFormatter>>, Impl<SimdJSONParser, JSONStringSerializer<SimdJSONParser::Element, SimdJSONElementFormatter>>,
SimdJSONParser>::run(arguments, result_type, input_rows_count, parse_depth, parse_backtracks, getContext()); SimdJSONParser>::run(arguments, result_type, input_rows_count, parse_depth, parse_backtracks, function_json_value_return_type_allow_complex);
#endif #endif
return FunctionSQLJSONHelpers:: return FunctionSQLJSONHelpers::
Executor<Name, Impl<DummyJSONParser, DefaultJSONStringSerializer<DummyJSONParser::Element>>, DummyJSONParser>::run( Executor<Name, Impl<DummyJSONParser, DefaultJSONStringSerializer<DummyJSONParser::Element>>, DummyJSONParser>::run(
arguments, result_type, input_rows_count, parse_depth, parse_backtracks, getContext()); arguments, result_type, input_rows_count, parse_depth, parse_backtracks, function_json_value_return_type_allow_complex);
} }
private:
const size_t max_parser_depth;
const size_t max_parser_backtracks;
const bool allow_simdjson;
const bool function_json_value_return_type_allow_complex;
const bool function_json_value_return_type_allow_nullable;
}; };
struct NameJSONExists struct NameJSONExists
@ -267,11 +280,11 @@ class JSONExistsImpl
public: public:
using Element = typename JSONParser::Element; using Element = typename JSONParser::Element;
static DataTypePtr getReturnType(const char *, const ColumnsWithTypeAndName &, const ContextPtr &) { return std::make_shared<DataTypeUInt8>(); } static DataTypePtr getReturnType(const char *, const ColumnsWithTypeAndName &, bool) { return std::make_shared<DataTypeUInt8>(); }
static size_t getNumberOfIndexArguments(const ColumnsWithTypeAndName & arguments) { return arguments.size() - 1; } static size_t getNumberOfIndexArguments(const ColumnsWithTypeAndName & arguments) { return arguments.size() - 1; }
static bool insertResultToColumn(IColumn & dest, const Element & root, GeneratorJSONPath<JSONParser> & generator_json_path, const ContextPtr &) static bool insertResultToColumn(IColumn & dest, const Element & root, GeneratorJSONPath<JSONParser> & generator_json_path, bool)
{ {
Element current_element = root; Element current_element = root;
VisitorStatus status; VisitorStatus status;
@ -305,9 +318,9 @@ class JSONValueImpl
public: public:
using Element = typename JSONParser::Element; using Element = typename JSONParser::Element;
static DataTypePtr getReturnType(const char *, const ColumnsWithTypeAndName &, const ContextPtr & context) static DataTypePtr getReturnType(const char *, const ColumnsWithTypeAndName &, bool function_json_value_return_type_allow_nullable)
{ {
if (context->getSettingsRef().function_json_value_return_type_allow_nullable) if (function_json_value_return_type_allow_nullable)
{ {
DataTypePtr string_type = std::make_shared<DataTypeString>(); DataTypePtr string_type = std::make_shared<DataTypeString>();
return std::make_shared<DataTypeNullable>(string_type); return std::make_shared<DataTypeNullable>(string_type);
@ -320,7 +333,7 @@ public:
static size_t getNumberOfIndexArguments(const ColumnsWithTypeAndName & arguments) { return arguments.size() - 1; } static size_t getNumberOfIndexArguments(const ColumnsWithTypeAndName & arguments) { return arguments.size() - 1; }
static bool insertResultToColumn(IColumn & dest, const Element & root, GeneratorJSONPath<JSONParser> & generator_json_path, const ContextPtr & context) static bool insertResultToColumn(IColumn & dest, const Element & root, GeneratorJSONPath<JSONParser> & generator_json_path, bool function_json_value_return_type_allow_complex)
{ {
Element current_element = root; Element current_element = root;
VisitorStatus status; VisitorStatus status;
@ -329,7 +342,7 @@ public:
{ {
if (status == VisitorStatus::Ok) if (status == VisitorStatus::Ok)
{ {
if (context->getSettingsRef().function_json_value_return_type_allow_complex) if (function_json_value_return_type_allow_complex)
{ {
break; break;
} }
@ -383,11 +396,11 @@ class JSONQueryImpl
public: public:
using Element = typename JSONParser::Element; using Element = typename JSONParser::Element;
static DataTypePtr getReturnType(const char *, const ColumnsWithTypeAndName &, const ContextPtr &) { return std::make_shared<DataTypeString>(); } static DataTypePtr getReturnType(const char *, const ColumnsWithTypeAndName &, bool) { return std::make_shared<DataTypeString>(); }
static size_t getNumberOfIndexArguments(const ColumnsWithTypeAndName & arguments) { return arguments.size() - 1; } static size_t getNumberOfIndexArguments(const ColumnsWithTypeAndName & arguments) { return arguments.size() - 1; }
static bool insertResultToColumn(IColumn & dest, const Element & root, GeneratorJSONPath<JSONParser> & generator_json_path, const ContextPtr &) static bool insertResultToColumn(IColumn & dest, const Element & root, GeneratorJSONPath<JSONParser> & generator_json_path, bool)
{ {
ColumnString & col_str = assert_cast<ColumnString &>(dest); ColumnString & col_str = assert_cast<ColumnString &>(dest);