feature: add execution timeout for remote UDFs (#827)

* feature: add execution timeout for remote UDFs

* small adjustments

* change the way to store the props and add more unit tests

* Fix: disallow duplicate key in remote function settings

* add some comments on timeout

* add some logics about dealing with illegal queries

Author: yuanhang zhao

src/Common/ErrorCodes.cpp

@@ -694,6 +694,7 @@
     M(2611, UNKNOWN_FORMAT_SCHEMA) \
     M(2612, AMBIGUOUS_FORMAT_SCHEMA) \
     M(2613, UNKNOWN_FORMAT_SCHEMA_TYPE) \
+    M(2631, DUPLICATE_KEY) \
 /* See END */
 
 namespace DB

src/Common/sendRequest.cpp

@@ -33,10 +33,10 @@ std::pair<String, Int32> sendRequest(
     const String & password,
     const String & payload,
     const std::vector<std::pair<String, String>> & headers,
+    /// Timeout second for connect/send/receive
+    ConnectionTimeouts timeouts,
     Poco::Logger * log)
 {
-    /// One second for connect/send/receive
-    ConnectionTimeouts timeouts({2, 0}, {5, 0}, {10, 0});
 
     PooledHTTPSessionPtr session;
     try
@@ -109,6 +109,18 @@ std::pair<String, Int32> sendRequest(
         {
             session->attachSessionData(e.message());
         }
+        if (e.code() == 1000){
+            LOG_ERROR(
+                log,
+                "Execution timeout from uri={} method={} payload={} query_id={} error={} exception={}",
+                uri.toString(),
+                method,
+                payload,
+                query_id,
+                e.message(),
+                getCurrentExceptionMessage(true, true));
+            return {"Execution timeout", toHTTPCode(e)};   
+        }
 
         LOG_ERROR(
             log,

src/Common/sendRequest.h

@@ -2,6 +2,8 @@
 
 #include <base/types.h>
 
+#include <IO/ConnectionTimeouts.h>
+
 #include <Poco/URI.h>
 
 #include  <utility>
@@ -23,5 +25,6 @@ std::pair<String, Int32> sendRequest(
     const String & password,
     const String & payload,
     const std::vector<std::pair<String, String>> & headers,
+    ConnectionTimeouts timeouts,
     Poco::Logger * log);
 }

src/Coordination/MetaStoreConnection.cpp

@@ -215,7 +215,7 @@ std::pair<String, Int32> MetaStoreConnection::forwardRequest(
     for (auto it = active_servers.begin(); it != active_servers.end();)
     {
         Poco::URI uri{fmt::format({METASTORE_URL}, it->getHost(), it->getPort(), uri_parameter)};
-        auto [response, http_status] = sendRequest(uri, method, query_id, user, password, body, {}, log);
+        auto [response, http_status] = sendRequest(uri, method, query_id, user, password, body, {}, ConnectionTimeouts({2, 0}/* connect timeout */, {5, 0}/* send timeout */, {10, 0}/* receive timeout */) , log);
 
         if (http_status != Poco::Net::HTTPResponse::HTTP_INTERNAL_SERVER_ERROR)
         {

src/Functions/UserDefined/RemoteUserDefinedFunction.h

@@ -2,6 +2,7 @@
 
 #include <Functions/UserDefined/UserDefinedFunctionBase.h>
 #include <IO/ReadBufferFromString.h>
+#include <IO/ConnectionTimeouts.h>
 #include <Processors/Formats/IOutputFormat.h>
 #include <Processors/Formats/IRowInputFormat.h>
 #include <Common/sendRequest.h>
@@ -53,6 +54,7 @@ class RemoteUserDefinedFunction final : public UserDefinedFunctionBase
             "",
             out,
             {{config.auth_context.key_name, config.auth_context.key_value}, {"", context->getCurrentQueryId()}},
+            ConnectionTimeouts({2, 0}/* connect timeout */, {5, 0} /* send timeout */, {static_cast<long>(config.command_execution_timeout_milliseconds / 1000), static_cast<long>((config.command_execution_timeout_milliseconds % 1000u) * 1000u)/* receive timeout */}), ///  timeout and limit for connect/send/receive ...
             &Poco::Logger::get("UserDefinedFunction"));
 
         if (http_status != Poco::Net::HTTPResponse::HTTP_OK)

src/Functions/UserDefined/UDFHelper.cpp

@@ -255,6 +255,7 @@ createUserDefinedExecutableFunction(ContextPtr context, const std::string & name
             throw Exception(ErrorCodes::BAD_ARGUMENTS, "Invalid url for remote UDF, msg: {}", e.message());
         }
         cfg->command_read_timeout_milliseconds = config.getUInt64(key_in_config + ".command_read_timeout", 10000);
+        cfg->command_execution_timeout_milliseconds = config.getUInt64(key_in_config + ".command_execution_timeout", 10000);
         cfg->auth_method = std::move(auth_method);
         cfg->auth_context = std::move(auth_ctx);
     };

src/Functions/UserDefined/UserDefinedFunctionConfiguration.h

@@ -80,6 +80,9 @@ struct RemoteUserDefinedFunctionConfiguration : public UserDefinedFunctionConfig
     /// Timeout for reading data from input format
     size_t command_read_timeout_milliseconds = 10000;
 
+    /// Timeout for receiving response from remote endpoint
+    size_t command_execution_timeout_milliseconds = 10000;
+
     /// url of remote endpoint, only available when 'type' is 'remote'
     Poco::URI url;
     enum AuthMethod

src/Parsers/ASTCreateFunctionQuery.cpp

@@ -5,6 +5,8 @@
 #include <Parsers/ASTFunction.h>
 
 /// proton: starts
+#include <optional>
+#include <Parsers/ASTFunctionWithKeyValueArguments.h>
 #include <Parsers/ASTLiteral.h>
 #include <Parsers/ASTNameTypePair.h>
 #include <Parsers/formatAST.h>
@@ -72,15 +74,8 @@ void ASTCreateFunctionQuery::formatImpl(const IAST::FormatSettings & settings, I
     /// proton: starts
     if (is_remote)
     {
-        settings.ostr << fmt::format("\nURL '{}'\n", function_core->as<ASTLiteral>()->value.safeGet<String>());
-        auto auth_method
-            = !function_core->children.empty() ? function_core->children[0]->as<ASTLiteral>()->value.safeGet<String>() : "none";
-        settings.ostr << fmt::format("AUTH_METHOD '{}'\n", auth_method);
-        if (auth_method != "none")
-        {
-            settings.ostr << fmt::format("AUTH_HEADER '{}'\n", function_core->children[1]->as<ASTLiteral>()->value.safeGet<String>());
-            settings.ostr << fmt::format("AUTH_KEY '{}'\n", function_core->children[2]->as<ASTLiteral>()->value.safeGet<String>());
-        }
+        settings.ostr << ' ';
+        function_core->formatImpl(settings, state, frame);
         return;
     }
     /// proton: ends
@@ -153,21 +148,50 @@ Poco::JSON::Object::Ptr ASTCreateFunctionQuery::toJSON() const
     /// remote function
     if (is_remote)
     {
-        assert(function_core != nullptr);
-        inner_func->set("url", function_core->as<ASTLiteral>()->value.safeGet<String>());
-        // auth
-        if (!function_core->children.empty())
+        assert(function_core != nullptr && function_core->as<ASTExpressionList>());
+        auto keyvalue_list = function_core->as<ASTExpressionList>();
+        std::optional<String> url;
+        std::optional<String> auth_method;
+        std::optional<String> auth_header;
+        std::optional<String> auth_key;
+        std::optional<UInt64> execution_timeout;
+        for (ASTPtr child : keyvalue_list->children)
         {
-            auto auth_method = function_core->children[0]->as<ASTLiteral>()->value.safeGet<String>();
-            inner_func->set("auth_method", auth_method);
-            if (auth_method == "auth_header")
+            auto pair = child->as<ASTPair>();
+            if (pair != nullptr){
+                if (pair->first == "url")
+                    url = pair->second->as<ASTLiteral>()->value.safeGet<String>();
+                else if (pair->first == "auth_method")
+                    auth_method = pair->second->as<ASTLiteral>()->value.safeGet<String>();
+                else if (pair->first == "auth_header")
+                    auth_header = pair->second->as<ASTLiteral>()->value.safeGet<String>();
+                else if (pair->first == "auth_key")
+                    auth_key = pair->second->as<ASTLiteral>()->value.safeGet<String>();
+                else if (pair->first == "execution_timeout")
+                    execution_timeout = pair->second->as<ASTLiteral>()->value.safeGet<UInt64>();
+            }
+        }        
+
+        inner_func->set("url", url.value());
+        if (auth_method.has_value())
+        {   
+            inner_func->set("auth_method", auth_method.value());
+            if (auth_method.value() == "auth_header")
             {
                 Poco::JSON::Object::Ptr auth_context = new Poco::JSON::Object();
-                auth_context->set("key_name", function_core->children[1]->as<ASTLiteral>()->value.safeGet<String>());
-                auth_context->set("key_value", function_core->children[2]->as<ASTLiteral>()->value.safeGet<String>());
+                auth_context->set("key_name", auth_header.value_or(""));
+                auth_context->set("key_value", auth_key.value_or(""));
                 inner_func->set("auth_context", auth_context);
             }
         }
+        else
+        {
+            inner_func->set("auth_method", "none");
+        }
+        if (execution_timeout.has_value())
+        {
+            inner_func->set("command_execution_timeout", execution_timeout.value());
+        }
         func->set("function", inner_func);
         /// Remote function don't have source, return early.
         return func;

src/Parsers/ParserCreateFunctionQuery.cpp

@@ -11,6 +11,7 @@
 #include <Parsers/ASTLiteral.h>
 #include <Parsers/ParserCreateQuery.h>
 #include <Parsers/Streaming/ParserArguments.h>
+#include <Parsers/ParserKeyValuePairsSet.h>
 
 #include <Poco/JSON/Object.h>
 /// proton: ends
@@ -40,11 +41,14 @@ bool ParserCreateFunctionQuery::parseImpl(IParser::Pos & pos, ASTPtr & node, Exp
     ParserKeyword s_auth_method("AUTH_METHOD");
     ParserKeyword s_auth_header("AUTH_HEADER");
     ParserKeyword s_auth_key("AUTH_KEY");
+    ParserKeyword s_execution_timeout("EXECUTION_TIMEOUT");
     ParserLiteral value;
+    ASTPtr kv_list;
     ASTPtr url;
     ASTPtr auth_method;
     ASTPtr auth_header;
     ASTPtr auth_key;
+    ASTPtr execution_timeout;
     ParserArguments arguments_p;
     ParserDataType return_p;
     ParserStringLiteral js_src_p;
@@ -139,35 +143,74 @@ bool ParserCreateFunctionQuery::parseImpl(IParser::Pos & pos, ASTPtr & node, Exp
         {
             throw Exception("Remote udf can not be an aggregate function", ErrorCodes::AGGREGATE_FUNCTION_NOT_APPLICABLE);
         }
-        if (!s_url.ignore(pos, expected))
+        ParserKeyValuePairsSet kv_pairs_list;
+        if (!kv_pairs_list.parse(pos, kv_list, expected))
             return false;
-        if (!value.parse(pos, url, expected))
-            return false;
-        if (s_auth_method.ignore(pos, expected))
+        
+        /// check if the parameters are valid and no unsupported or unknown parameters.
+        std::optional<String> ast_url;
+        std::optional<String> ast_auth_method;
+        std::optional<String> ast_auth_header;
+        std::optional<String> ast_auth_key;
+        std::optional<UInt64> ast_execution_timeout;
+        for (const auto & kv : kv_list->children)
+        {
+            auto * kv_pair = kv->as<ASTPair>();
+            auto key = kv_pair->first;
+            auto pair_value = kv_pair->second->as<ASTLiteral>()->value;
+            if (!kv_pair)
+                throw Exception("Key-value pair expected", ErrorCodes::UNKNOWN_FUNCTION);
+            
+            if (key == "url")
+            {
+                ast_url = pair_value.safeGet<String>();
+            }
+            else if (key == "auth_method")
+            {
+                ast_auth_method = pair_value.safeGet<String>();
+                if (ast_auth_method.value() != "none" && ast_auth_method.value() != "auth_header")
+                    throw Exception("Unknown auth method", ErrorCodes::UNKNOWN_FUNCTION);
+            }
+            else if (key == "auth_header")
+            {
+                ast_auth_header = pair_value.safeGet<String>();
+            }
+            else if (key == "auth_key")
+            {
+                ast_auth_key = pair_value.safeGet<String>();
+            }
+            else if (key == "execution_timeout")
+            {
+                ast_execution_timeout = pair_value.safeGet<UInt64>();
+            }
+        }
+        /// check if URL is set
+        if (!ast_url)
+            throw Exception("URL is required for remote function", ErrorCodes::UNKNOWN_FUNCTION);
+        /// check if auth_method is "auth_header" or "none"
+        if (ast_auth_method)
         {
-            if (!value.parse(pos, auth_method, expected))
-                return false;
-            auto method_str = auth_method->as<ASTLiteral>()->value.safeGet<String>();
-            url->children.push_back(std::move(auth_method));
-            if (method_str == "auth_header")
+            if (ast_auth_method.value() == "auth_header")
             {
-                if (!s_auth_header.ignore(pos, expected))
-                    return false;
-                if (!value.parse(pos, auth_header, expected))
-                    return false;
-                if (!s_auth_key.ignore(pos, expected))
-                    return false;
-                if (!value.parse(pos, auth_key, expected))
-                    return false;
-                url->children.push_back(std::move(auth_header));
-                url->children.push_back(std::move(auth_key));
+                if (!ast_auth_header  || !ast_auth_key)
+                    throw Exception("Auth header and auth key are required for auth_header auth method", ErrorCodes::UNKNOWN_FUNCTION);
             }
-            else if (method_str != "none")
+            else if (ast_auth_method.value() == "none")
             {
-                throw Exception("AUTH_METHOD must be 'none' or 'auth_header'", ErrorCodes::UNKNOWN_FUNCTION);
+                if (ast_auth_header || ast_auth_key)
+                    throw Exception("Auth method is 'none', but auth header or auth key is set.", ErrorCodes::UNKNOWN_FUNCTION);
             }
+            else
+            {
+                throw Exception("Unknown auth method " + ast_auth_method.value(), ErrorCodes::UNKNOWN_FUNCTION);
+            }
+        }
+        else
+        {
+            if (ast_auth_header || ast_auth_key)
+                throw Exception("Auth method is 'none', but auth header or auth key is set.", ErrorCodes::UNKNOWN_FUNCTION);
         }
-        function_core = std::move(url);
+        function_core = std::move(kv_list);
     }
     /// proton: ends
 

src/Parsers/ParserKeyValuePairsSet.cpp

@@ -0,0 +1,61 @@
+#include <Parsers/ParserKeyValuePairsSet.h>
+#include <Parsers/ASTExpressionList.h>
+#include <Common/ErrorCodes.h>
+#include <unordered_map>
+
+
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int DUPLICATE_KEY;
+}
+
+bool ParserKeyValuePairsSet::parseImpl(Pos & pos, ASTPtr & node, Expected & expected, [[ maybe_unused ]] bool hint)
+{
+    ASTs elements;
+    std::unordered_map<String, size_t> exists_keys;
+
+    auto parse_element = [&]
+    {
+        ASTPtr element;
+        if (!elem_parser->parse(pos, element, expected))
+            return false;
+        auto key_value = element->as<ASTPair>();
+        if (exists_keys.find(key_value->first) == exists_keys.end())
+        {
+            exists_keys[key_value->first] = elements.size();
+            elements.push_back(element);
+        }
+        else
+        {
+            if (allow_duplicate)
+            {
+                size_t index = exists_keys[key_value->first];
+                elements.erase(elements.begin() + index);
+                exists_keys[key_value->first] = elements.size();
+                elements.push_back(element);
+            }
+            else
+            {
+                throw Exception("Duplicate key \"" + key_value->first + "\" has existed previously", ErrorCodes::DUPLICATE_KEY);
+            }
+        }
+        
+        return true;
+    };
+
+    if (!ParserList::parseUtil(pos, expected, parse_element, *separator_parser, allow_empty))
+        return false;
+
+    auto list = std::make_shared<ASTExpressionList>(result_separator);
+    list->children = std::move(elements);
+    node = list;
+
+    return true;
+}
+
+}
+