Execute table functions once for a query [#CLICKHOUSE-3615]

dbms/src/Common/ErrorCodes.cpp

@@ -368,6 +368,7 @@ namespace ErrorCodes
     extern const int CANNOT_GET_CREATE_TABLE_QUERY = 390;
     extern const int EXTERNAL_LIBRARY_ERROR = 391;
     extern const int QUERY_IS_PROHIBITED = 392;
+    extern const int THERE_IS_NO_QUERY = 393;
 
 
     extern const int KEEPER_EXCEPTION = 999;

dbms/src/Interpreters/Context.cpp

@@ -76,6 +76,7 @@ namespace ErrorCodes
     extern const int DATABASE_ALREADY_EXISTS;
     extern const int TABLE_METADATA_DOESNT_EXIST;
     extern const int THERE_IS_NO_SESSION;
+    extern const int THERE_IS_NO_QUERY;
     extern const int NO_ELEMENTS_IN_CONFIG;
     extern const int DDL_GUARD_IS_ACTIVE;
     extern const int TABLE_SIZE_EXCEEDS_MAX_DROP_SIZE_LIMIT;
@@ -919,30 +920,29 @@ ASTPtr Context::getCreateQuery(const String & database_name, const String & tabl
 
 Settings Context::getSettings() const
 {
-    auto lock = getLock();
     return settings;
 }
 
 
 Limits Context::getLimits() const
 {
-    auto lock = getLock();
     return settings.limits;
 }
 
 
 void Context::setSettings(const Settings & settings_)
 {
-    auto lock = getLock();
     settings = settings_;
 }
 
 
 void Context::setSetting(const String & name, const Field & value)
 {
-    auto lock = getLock();
     if (name == "profile")
+    {
+        auto lock = getLock();
         settings.setProfile(value.safeGet<String>(), *shared->users_config);
+    }
     else
         settings.set(name, value);
 }
@@ -950,9 +950,11 @@ void Context::setSetting(const String & name, const Field & value)
 
 void Context::setSetting(const String & name, const std::string & value)
 {
-    auto lock = getLock();
     if (name == "profile")
+    {
+        auto lock = getLock();
         settings.setProfile(value, *shared->users_config);
+    }
     else
         settings.set(name, value);
 }
@@ -960,14 +962,12 @@ void Context::setSetting(const String & name, const std::string & value)
 
 String Context::getCurrentDatabase() const
 {
-    auto lock = getLock();
     return current_database;
 }
 
 
 String Context::getCurrentQueryId() const
 {
-    auto lock = getLock();
     return client_info.current_query_id;
 }
 
@@ -982,8 +982,6 @@ void Context::setCurrentDatabase(const String & name)
 
 void Context::setCurrentQueryId(const String & query_id)
 {
-    auto lock = getLock();
-
     if (!client_info.current_query_id.empty())
         throw Exception("Logical error: attempt to set query_id twice", ErrorCodes::LOGICAL_ERROR);
 
@@ -1006,8 +1004,12 @@ void Context::setCurrentQueryId(const String & query_id)
             };
         } random;
 
-        random.a = shared->rng();
-        random.b = shared->rng();
+        {
+            auto lock = getLock();
+
+            random.a = shared->rng();
+            random.b = shared->rng();
+        }
 
         /// Use protected constructor.
         struct UUID : Poco::UUID
@@ -1025,14 +1027,12 @@ void Context::setCurrentQueryId(const String & query_id)
 
 String Context::getDefaultFormat() const
 {
-    auto lock = getLock();
     return default_format.empty() ? "TabSeparated" : default_format;
 }
 
 
 void Context::setDefaultFormat(const String & name)
 {
-    auto lock = getLock();
     default_format = name;
 }
 
@@ -1047,6 +1047,20 @@ void Context::setMacros(Macros && macros)
     shared->macros = macros;
 }
 
+const Context & Context::getQueryContext() const
+{
+    if (!query_context)
+        throw Exception("There is no query", ErrorCodes::THERE_IS_NO_QUERY);
+    return *query_context;
+}
+
+Context & Context::getQueryContext()
+{
+    if (!query_context)
+        throw Exception("There is no query", ErrorCodes::THERE_IS_NO_QUERY);
+    return *query_context;
+}
+
 const Context & Context::getSessionContext() const
 {
     if (!session_context)

dbms/src/Interpreters/Context.h

@@ -105,6 +105,7 @@ private:
                             /// Thus, used in HTTP interface. If not specified - then some globally default format is used.
     Tables external_tables;                 /// Temporary tables. Keyed by table name.
     Tables table_function_results;          /// Temporary tables obtained by execution of table functions. Keyed by AST tree id.
+    Context * query_context = nullptr;
     Context * session_context = nullptr;    /// Session context or nullptr. Could be equal to this.
     Context * global_context = nullptr;     /// Global context or nullptr. Could be equal to this.
     SystemLogsPtr system_logs;              /// Used to log queries and operations on parts
@@ -257,6 +258,10 @@ public:
     /// For methods below you may need to acquire a lock by yourself.
     std::unique_lock<std::recursive_mutex> getLock() const;
 
+    const Context & getQueryContext() const;
+    Context & getQueryContext();
+    bool hasQueryContext() const { return query_context != nullptr; }
+
     const Context & getSessionContext() const;
     Context & getSessionContext();
     bool hasSessionContext() const { return session_context != nullptr; }
@@ -265,6 +270,7 @@ public:
     Context & getGlobalContext();
     bool hasGlobalContext() const { return global_context != nullptr; }
 
+    void setQueryContext(Context & context_) { query_context = &context_; }
     void setSessionContext(Context & context_) { session_context = &context_; }
     void setGlobalContext(Context & context_) { global_context = &context_; }
 

dbms/src/Interpreters/InterpreterSelectQuery.cpp

@@ -103,6 +103,9 @@ void InterpreterSelectQuery::init(const Names & required_result_column_names)
 {
     ProfileEvents::increment(ProfileEvents::SelectQuery);
 
+    if (!context.hasQueryContext())
+        context.setQueryContext(context);
+
     initSettings();
     const Settings & settings = context.getSettingsRef();
 
@@ -128,7 +131,7 @@ void InterpreterSelectQuery::init(const Names & required_result_column_names)
     else if (table_expression && typeid_cast<const ASTFunction *>(table_expression.get()))
     {
         /// Read from table function.
-        storage = context.executeTableFunction(table_expression);
+        storage = context.getQueryContext().executeTableFunction(table_expression);
     }
     else
     {

dbms/src/Interpreters/InterpreterSelectWithUnionQuery.cpp

@@ -32,6 +32,9 @@ InterpreterSelectWithUnionQuery::InterpreterSelectWithUnionQuery(
     to_stage(to_stage_),
     subquery_depth(subquery_depth_)
 {
+    if (!context.hasQueryContext())
+        context.setQueryContext(context);
+
     const ASTSelectWithUnionQuery & ast = typeid_cast<const ASTSelectWithUnionQuery &>(*query_ptr);
 
     size_t num_selects = ast.list_of_selects->children.size();

dbms/src/Interpreters/executeQuery.cpp

@@ -140,6 +140,8 @@ static std::tuple<ASTPtr, BlockIO> executeQueryImpl(
     ProfileEvents::increment(ProfileEvents::Query);
     time_t current_time = time(nullptr);
 
+    context.setQueryContext(context);
+
     const Settings & settings = context.getSettingsRef();
 
     ParserQuery parser(end);