Merge pull request #71562 from ClickHouse/backport/24.3/71541

Backport #71541 to 24.3: Avoid crash when using a UDF in a constraint

src/Functions/UserDefined/UserDefinedSQLFunctionVisitor.cpp

@@ -24,92 +24,7 @@ namespace ErrorCodes
 
 void UserDefinedSQLFunctionVisitor::visit(ASTPtr & ast)
 {
-    if (!ast)
-    {
-        chassert(false);
-        return;
-    }
-
-    /// FIXME: this helper should use updatePointerToChild(), but
-    /// forEachPointerToChild() is not implemented for ASTColumnDeclaration
-    /// (and also some members should be adjusted for this).
-    const auto visit_child_with_shared_ptr = [&](ASTPtr & child)
-    {
-        if (!child)
-            return;
-
-        auto * old_value = child.get();
-        visit(child);
-
-        // child did not change
-        if (old_value == child.get())
-            return;
-
-        // child changed, we need to modify it in the list of children of the parent also
-        for (auto & current_child : ast->children)
-        {
-            if (current_child.get() == old_value)
-                current_child = child;
-        }
-    };
-
-    if (auto * col_decl = ast->as<ASTColumnDeclaration>())
-    {
-        visit_child_with_shared_ptr(col_decl->default_expression);
-        visit_child_with_shared_ptr(col_decl->ttl);
-        return;
-    }
-
-    if (auto * storage = ast->as<ASTStorage>())
-    {
-        const auto visit_child = [&](IAST * & child)
-        {
-            if (!child)
-                return;
-
-            if (const auto * function = child->template as<ASTFunction>())
-            {
-                std::unordered_set<std::string> udf_in_replace_process;
-                auto replace_result = tryToReplaceFunction(*function, udf_in_replace_process);
-                if (replace_result)
-                    ast->setOrReplace(child, replace_result);
-            }
-
-            visit(child);
-        };
-
-        visit_child(storage->partition_by);
-        visit_child(storage->primary_key);
-        visit_child(storage->order_by);
-        visit_child(storage->sample_by);
-        visit_child(storage->ttl_table);
-
-        return;
-    }
-
-    if (auto * alter = ast->as<ASTAlterCommand>())
-    {
-        /// It is OK to use updatePointerToChild() because ASTAlterCommand implements forEachPointerToChild()
-        const auto visit_child_update_parent = [&](ASTPtr & child)
-        {
-            if (!child)
-                return;
-
-            auto * old_ptr = child.get();
-            visit(child);
-            auto * new_ptr = child.get();
-
-            /// Some AST classes have naked pointers to children elements as members.
-            /// We have to replace them if the child was replaced.
-            if (new_ptr != old_ptr)
-                ast->updatePointerToChild(old_ptr, new_ptr);
-        };
-
-        for (auto & children : alter->children)
-            visit_child_update_parent(children);
-
-        return;
-    }
+    chassert(ast);
 
     if (const auto * function = ast->template as<ASTFunction>())
     {
@@ -120,7 +35,19 @@ void UserDefinedSQLFunctionVisitor::visit(ASTPtr & ast)
     }
 
     for (auto & child : ast->children)
+    {
+        if (!child)
+            return;
+
+        auto * old_ptr = child.get();
         visit(child);
+        auto * new_ptr = child.get();
+
+        /// Some AST classes have naked pointers to children elements as members.
+        /// We have to replace them if the child was replaced.
+        if (new_ptr != old_ptr)
+            ast->updatePointerToChild(old_ptr, new_ptr);
+    }
 }
 
 void UserDefinedSQLFunctionVisitor::visit(IAST * ast)

src/Parsers/ASTColumnDeclaration.cpp

@@ -137,4 +137,14 @@ void ASTColumnDeclaration::formatImpl(const FormatSettings & format_settings, Fo
     }
 }
 
+void ASTColumnDeclaration::forEachPointerToChild(std::function<void(void **)> f)
+{
+    f(reinterpret_cast<void **>(&default_expression));
+    f(reinterpret_cast<void **>(&comment));
+    f(reinterpret_cast<void **>(&codec));
+    f(reinterpret_cast<void **>(&stat_type));
+    f(reinterpret_cast<void **>(&ttl));
+    f(reinterpret_cast<void **>(&collation));
+    f(reinterpret_cast<void **>(&settings));
+}
 }

src/Parsers/ASTColumnDeclaration.h

@@ -29,6 +29,9 @@ public:
 
     ASTPtr clone() const override;
     void formatImpl(const FormatSettings & format_settings, FormatState & state, FormatStateStacked frame) const override;
+
+protected:
+    void forEachPointerToChild(std::function<void(void **)> f) override;
 };
 
 }

tests/queries/0_stateless/03262_udf_in_constraint.reference

@@ -0,0 +1,2 @@
+CREATE TABLE default.t0\n(\n    `c0` Int32,\n    CONSTRAINT c1 CHECK c0 > 5\n)\nENGINE = MergeTree\nORDER BY tuple()\nSETTINGS index_granularity = 8192
+10

tests/queries/0_stateless/03262_udf_in_constraint.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+CUR_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+# shellcheck source=../shell_config.sh
+. "$CUR_DIR"/../shell_config.sh
+
+$CLICKHOUSE_CLIENT -nm -q "
+  CREATE FUNCTION ${CLICKHOUSE_DATABASE}_function AS (x) -> x > 5;
+  CREATE TABLE t0 (c0 Int, CONSTRAINT c1 CHECK ${CLICKHOUSE_DATABASE}_function(c0)) ENGINE = MergeTree() ORDER BY tuple();
+  SHOW CREATE TABLE t0;
+  INSERT INTO t0(c0) VALUES (10);
+  INSERT INTO t0(c0) VALUES (3); -- {serverError VIOLATED_CONSTRAINT}
+  SELECT * FROM t0;
+
+  DROP TABLE t0;
+  DROP FUNCTION ${CLICKHOUSE_DATABASE}_function;
+"