AddDefaultDatabaseVisitor support dictGet

src/Interpreters/AddDefaultDatabaseVisitor.h

@@ -10,10 +10,14 @@
 #include <Parsers/ASTSelectWithUnionQuery.h>
 #include <Parsers/ASTTablesInSelectQuery.h>
 #include <Parsers/ASTFunction.h>
+#include <Parsers/ASTCreateQuery.h>
 #include <Parsers/DumpASTNode.h>
 #include <Parsers/ASTAlterQuery.h>
 #include <Interpreters/DatabaseAndTableWithAlias.h>
 #include <Interpreters/IdentifierSemantic.h>
+#include <Interpreters/Context.h>
+#include <Interpreters/ExternalDictionariesLoader.h>
+#include <Interpreters/misc.h>
 
 namespace DB
 {
@@ -26,11 +30,12 @@ class AddDefaultDatabaseVisitor
 {
 public:
     explicit AddDefaultDatabaseVisitor(
-        const String & database_name_, bool only_replace_current_database_function_ = false, WriteBuffer * ostr_ = nullptr)
-        : database_name(database_name_)
+        ContextPtr context_,
+        const String & database_name_,
+        bool only_replace_current_database_function_ = false)
+        : context(context_)
+        , database_name(database_name_)
         , only_replace_current_database_function(only_replace_current_database_function_)
-        , visit_depth(0)
-        , ostr(ostr_)
     {}
 
     void visitDDL(ASTPtr & ast) const
@@ -64,11 +69,19 @@ public:
         visit(select, unused);
     }
 
+    void visit(ASTColumns & columns) const
+    {
+        for (auto & child : columns.children)
+            visit(child);
+    }
+
 private:
+
+    ContextPtr context;
+
     const String database_name;
+
     bool only_replace_current_database_function = false;
-    mutable size_t visit_depth;
-    WriteBuffer * ostr;
 
     void visit(ASTSelectWithUnionQuery & select, ASTPtr &) const
     {
@@ -117,15 +130,8 @@ private:
 
     void visit(ASTFunction & function, ASTPtr &) const
     {
-        bool is_operator_in = false;
-        for (const auto * name : {"in", "notIn", "globalIn", "globalNotIn"})
-        {
-            if (function.name == name)
-            {
-                is_operator_in = true;
-                break;
-            }
-        }
+        bool is_operator_in = functionIsInOrGlobalInOperator(function.name);
+        bool is_dict_get = functionIsDictGet(function.name);
 
         for (auto & child : function.children)
         {
@@ -133,7 +139,38 @@ private:
             {
                 for (size_t i = 0; i < child->children.size(); ++i)
                 {
-                    if (is_operator_in && i == 1)
+                    if (is_dict_get && i == 0)
+                    {
+                        if (auto * identifier = child->children[i]->as<ASTIdentifier>())
+                        {
+                            if (identifier->compound())
+                                continue;
+
+                            auto storage_id = context->getExternalDictionariesLoader().getStorageID(identifier->name(), context);
+
+                            if (!storage_id.database_name.empty())
+                            {
+                                std::vector<std::string> name_parts = {storage_id.database_name, storage_id.table_name};
+                                child->children[i] = std::make_shared<ASTIdentifier>(std::move(name_parts));
+                            }
+                            else
+                            {
+                                std::vector<std::string> name_parts = {storage_id.table_name};
+                                child->children[i] = std::make_shared<ASTIdentifier>(std::move(name_parts));
+                            }
+                        }
+                        else if (auto * literal = child->children[i]->as<ASTLiteral>())
+                        {
+                            auto & literal_value = literal->value;
+
+                            if (literal_value.getType() != Field::Types::String)
+                                continue;
+
+                            auto dictionary_name = literal_value.get<String>();
+                            literal_value = context->getExternalDictionariesLoader().getStorageID(dictionary_name, context).getFullTableName();
+                        }
+                    }
+                    else if (is_operator_in && i == 1)
                     {
                         /// XXX: for some unknown reason this place assumes that argument can't be an alias,
                         ///      like in the similar code in `MarkTableIdentifierVisitor`.
@@ -151,11 +188,15 @@ private:
                             visit(child->children[i]);
                     }
                     else
+                    {
                         visit(child->children[i]);
+                    }
                 }
             }
             else
+            {
                 visit(child);
+            }
         }
     }
 
@@ -170,7 +211,6 @@ private:
     {
         if (T * t = typeid_cast<T *>(ast.get()))
         {
-            DumpASTNode dump(*ast, ostr, visit_depth, "addDefaultDatabaseName");
             visit(*t, ast);
             return true;
         }

src/Interpreters/ExternalDictionariesLoader.cpp

@@ -87,6 +87,21 @@ DictionaryStructure ExternalDictionariesLoader::getDictionaryStructure(const std
     return ExternalDictionariesLoader::getDictionaryStructure(*load_result.config);
 }
 
+StorageID ExternalDictionariesLoader::getStorageID(const std::string & dictionary_name, ContextPtr context) const
+{
+    if (has(dictionary_name))
+        return StorageID("", dictionary_name);
+
+    auto qualified_name = QualifiedTableName::tryParseFromString(dictionary_name);
+    if (!qualified_name)
+        return StorageID("", dictionary_name);
+
+    if (qualified_name->database.empty())
+        return StorageID(context->getCurrentDatabase(), dictionary_name);
+
+    return StorageID("", dictionary_name);
+}
+
 std::string ExternalDictionariesLoader::resolveDictionaryName(const std::string & dictionary_name, const std::string & current_database_name) const
 {
     if (has(dictionary_name))

src/Interpreters/ExternalDictionariesLoader.h

@@ -27,6 +27,8 @@ public:
 
     void reloadDictionary(const std::string & dictionary_name, ContextPtr context) const;
 
+    StorageID getStorageID(const std::string & dictionary_name, ContextPtr context) const;
+
     DictionaryStructure getDictionaryStructure(const std::string & dictionary_name, ContextPtr context) const;
 
     static DictionaryStructure getDictionaryStructure(const Poco::Util::AbstractConfiguration & config, const std::string & key_in_config = "dictionary");

src/Interpreters/InterpreterAlterQuery.cpp

@@ -80,7 +80,7 @@ BlockIO InterpreterAlterQuery::executeToTable(const ASTAlterQuery & alter)
     auto metadata_snapshot = table->getInMemoryMetadataPtr();
 
     /// Add default database to table identifiers that we can encounter in e.g. default expressions, mutation expression, etc.
-    AddDefaultDatabaseVisitor visitor(table_id.getDatabaseName());
+    AddDefaultDatabaseVisitor visitor(getContext(), table_id.getDatabaseName());
     ASTPtr command_list_ptr = alter.command_list->ptr();
     visitor.visit(command_list_ptr);
 

src/Interpreters/InterpreterCreateQuery.cpp

@@ -920,10 +920,17 @@ BlockIO InterpreterCreateQuery::createTable(ASTCreateQuery & create)
     {
         // Expand CTE before filling default database
         ApplyWithSubqueryVisitor().visit(*create.select);
-        AddDefaultDatabaseVisitor visitor(current_database);
+        AddDefaultDatabaseVisitor visitor(getContext(), current_database);
         visitor.visit(*create.select);
     }
 
+
+    if (create.columns_list)
+    {
+        AddDefaultDatabaseVisitor visitor(getContext(), current_database);
+        visitor.visit(*create.columns_list);
+    }
+
     /// Set and retrieve list of columns, indices and constraints. Set table engine if needed. Rewrite query in canonical way.
     TableProperties properties = getTablePropertiesAndNormalizeCreateQuery(create);
 

src/Interpreters/executeDDLQueryOnCluster.cpp

@@ -152,7 +152,7 @@ BlockIO executeDDLQueryOnCluster(const ASTPtr & query_ptr_, ContextPtr context,
         }
     }
 
-    AddDefaultDatabaseVisitor visitor(current_database, !use_local_default_database);
+    AddDefaultDatabaseVisitor visitor(context, current_database, !use_local_default_database);
     visitor.visitDDL(query_ptr);
 
     /// Check access rights, assume that all servers have the same users config

src/Storages/LiveView/StorageLiveView.cpp

@@ -70,7 +70,7 @@ static StorageID extractDependentTable(ASTPtr & query, ContextPtr context, const
         if (db_and_table->database.empty())
         {
             db_and_table->database = select_database_name;
-            AddDefaultDatabaseVisitor visitor(select_database_name);
+            AddDefaultDatabaseVisitor visitor(context, select_database_name);
             visitor.visit(select_query);
         }
         else

src/Storages/SelectQueryDescription.cpp

@@ -48,7 +48,7 @@ StorageID extractDependentTableFromSelectQuery(ASTSelectQuery & query, ContextPt
 {
     if (add_default_db)
     {
-        AddDefaultDatabaseVisitor visitor(context->getCurrentDatabase(), false, nullptr);
+        AddDefaultDatabaseVisitor visitor(context, context->getCurrentDatabase());
         visitor.visit(query);
     }
 

src/Storages/StorageMaterializedView.cpp

@@ -9,7 +9,6 @@
 #include <Interpreters/InterpreterDropQuery.h>
 #include <Interpreters/InterpreterRenameQuery.h>
 #include <Interpreters/getTableExpressions.h>
-#include <Interpreters/AddDefaultDatabaseVisitor.h>
 #include <Interpreters/getHeaderForProcessingStage.h>
 #include <Access/AccessFlags.h>
 #include <DataStreams/IBlockOutputStream.h>