From 9203007b615fd93d7d92bc25aab28b138f392293 Mon Sep 17 00:00:00 2001
From: Jan Kuipers <jan.kuipers@elastic.co>
Date: Thu, 10 Apr 2025 13:41:19 +0200
Subject: [PATCH] License check for ES|QL categorize

---
 .../esql/expression/function/grouping/Categorize.java  | 10 +++++++++-
 .../org/elasticsearch/xpack/ml/MachineLearning.java    |  2 +-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/grouping/Categorize.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/grouping/Categorize.java
index ef5275cd5f70b..f4be43b18a731 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/grouping/Categorize.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/grouping/Categorize.java
@@ -10,6 +10,8 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.license.XPackLicenseState;
+import org.elasticsearch.xpack.esql.LicenseAware;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Nullability;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
@@ -20,6 +22,7 @@
 import org.elasticsearch.xpack.esql.expression.function.FunctionType;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.ml.MachineLearning;
 
 import java.io.IOException;
 import java.util.List;
@@ -36,7 +39,7 @@
  *     For the implementation, see {@link org.elasticsearch.compute.aggregation.blockhash.CategorizeBlockHash}
  * </p>
  */
-public class Categorize extends GroupingFunction.NonEvaluatableGroupingFunction {
+public class Categorize extends GroupingFunction.NonEvaluatableGroupingFunction implements LicenseAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         Expression.class,
         "Categorize",
@@ -128,4 +131,9 @@ public Expression field() {
     public String toString() {
         return "Categorize{field=" + field + "}";
     }
+
+    @Override
+    public boolean licenseCheck(XPackLicenseState state) {
+        return MachineLearning.CATEGORIZE_TEXT_AGG_FEATURE.check(state);
+    }
 }
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearning.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearning.java
index 1894b1e59ad5c..7505e51300739 100644
--- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearning.java
+++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearning.java
@@ -525,7 +525,7 @@ public class MachineLearning extends Plugin
         License.OperationMode.PLATINUM
     );
 
-    private static final LicensedFeature.Momentary CATEGORIZE_TEXT_AGG_FEATURE = LicensedFeature.momentary(
+    public static final LicensedFeature.Momentary CATEGORIZE_TEXT_AGG_FEATURE = LicensedFeature.momentary(
         MachineLearningField.ML_FEATURE_FAMILY,
         "categorize-text-agg",
         License.OperationMode.PLATINUM