Adjust cartesian license levels (only centroid-shape) (#91625)

* Adjust cartesian license levels (only centroid-shape)

* Added SpatialPluginTests for cartesian centroid and bounds

Also refactored the tests to make it easier to add additional aggregations

* Moved a private method to the end

* Removed debugging code

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/SpatialPlugin.java

@@ -91,6 +91,11 @@ public class SpatialPlugin extends Plugin implements ActionPlugin, MapperPlugin,
         "geo-centroid-agg",
         License.OperationMode.GOLD
     );
+    private final LicensedFeature.Momentary CARTESIAN_CENTROID_AGG_FEATURE = LicensedFeature.momentary(
+        "spatial",
+        "cartesian-centroid-agg",
+        License.OperationMode.GOLD
+    );
     private final LicensedFeature.Momentary GEO_GRID_AGG_FEATURE = LicensedFeature.momentary(
         "spatial",
         "geo-grid-agg",
@@ -172,10 +177,7 @@ public class SpatialPlugin extends Plugin implements ActionPlugin, MapperPlugin,
             new AggregationSpec(
                 CartesianCentroidAggregationBuilder.NAME,
                 CartesianCentroidAggregationBuilder::new,
-                usage.track(
-                    SpatialStatsAction.Item.CARTESIANCENTROID,
-                    checkLicense(CartesianCentroidAggregationBuilder.PARSER, GEO_CENTROID_AGG_FEATURE)
-                )
+                usage.track(SpatialStatsAction.Item.CARTESIANCENTROID, CartesianCentroidAggregationBuilder.PARSER)
             ).addResultReader(InternalCartesianCentroid::new).setAggregatorRegistrar(this::registerCartesianCentroidAggregator),
             new AggregationSpec(
                 CartesianBoundsAggregationBuilder.NAME,
@@ -229,26 +231,23 @@ public class SpatialPlugin extends Plugin implements ActionPlugin, MapperPlugin,
     }
 
     private void registerCartesianCentroidAggregator(ValuesSourceRegistry.Builder builder) {
+        // Only aggregations over the shape type are licensed at Gold/Platinum level
         builder.register(
             CartesianCentroidAggregationBuilder.REGISTRY_KEY,
             CartesianShapeValuesSourceType.instance(),
             (name, valuesSourceConfig, context, parent, metadata) -> {
-                if (GEO_CENTROID_AGG_FEATURE.check(getLicenseState())) {
+                if (CARTESIAN_CENTROID_AGG_FEATURE.check(getLicenseState())) {
                     return new CartesianShapeCentroidAggregator(name, context, parent, valuesSourceConfig, metadata);
                 }
                 throw LicenseUtils.newComplianceException(CartesianCentroidAggregationBuilder.NAME + " aggregation on shape fields");
             },
             true
         );
+        // Points are licensed at the default level (Basic)
         builder.register(
             CartesianCentroidAggregationBuilder.REGISTRY_KEY,
             CartesianPointValuesSourceType.instance(),
-            (name, valuesSourceConfig, context, parent, metadata) -> {
-                if (GEO_CENTROID_AGG_FEATURE.check(getLicenseState())) {
-                    return new CartesianCentroidAggregator(name, valuesSourceConfig, context, parent, metadata);
-                }
-                throw LicenseUtils.newComplianceException(CartesianCentroidAggregationBuilder.NAME + " aggregation on point fields");
-            },
+            CartesianCentroidAggregator::new,
             true
         );
     }

x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/SpatialPluginTests.java

@@ -17,121 +17,189 @@ import org.elasticsearch.search.aggregations.bucket.geogrid.GeoHashGridAggregati
 import org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileGridAggregationBuilder;
 import org.elasticsearch.search.aggregations.metrics.GeoCentroidAggregationBuilder;
 import org.elasticsearch.search.aggregations.metrics.GeoGridAggregatorSupplier;
-import org.elasticsearch.search.aggregations.metrics.MetricAggregatorSupplier;
 import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
 import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
+import org.elasticsearch.search.aggregations.support.ValuesSourceType;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoHexGridAggregationBuilder;
+import org.elasticsearch.xpack.spatial.search.aggregations.metrics.CartesianBoundsAggregationBuilder;
+import org.elasticsearch.xpack.spatial.search.aggregations.metrics.CartesianCentroidAggregationBuilder;
+import org.elasticsearch.xpack.spatial.search.aggregations.support.CartesianPointValuesSourceType;
+import org.elasticsearch.xpack.spatial.search.aggregations.support.CartesianShapeValuesSourceType;
 import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSourceType;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
 import java.util.function.Consumer;
 
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 
 public class SpatialPluginTests extends ESTestCase {
 
-    public void testGeoCentroidLicenseCheck() {
+    public void testGeoShapeCentroidLicenseCheck() {
+        checkLicenseRequired(GeoShapeValuesSourceType.instance(), GeoCentroidAggregationBuilder.REGISTRY_KEY, (agg) -> {
+            try {
+                agg.build(null, null, null, null, null);
+            } catch (IOException e) {
+                fail("Unexpected exception: " + e.getMessage());
+            }
+        }, "geo_centroid", "geo_shape");
+    }
+
+    public void testGeoHexLicenseCheck() {
+        checkLicenseRequired(CoreValuesSourceType.GEOPOINT, GeoHexGridAggregationBuilder.REGISTRY_KEY, (agg) -> {
+            try {
+                agg.build(null, AggregatorFactories.EMPTY, null, 0, null, 0, 0, null, null, CardinalityUpperBound.NONE, null);
+            } catch (IOException e) {
+                fail("Unexpected exception: " + e.getMessage());
+            }
+        }, "geohex_grid", "geo_point");
+    }
+
+    public void testGeoGridLicenseCheck() {
+        for (ValuesSourceRegistry.RegistryKey<GeoGridAggregatorSupplier> registryKey : Arrays.asList(
+            GeoHashGridAggregationBuilder.REGISTRY_KEY,
+            GeoTileGridAggregationBuilder.REGISTRY_KEY
+        )) {
+            checkLicenseRequired(GeoShapeValuesSourceType.instance(), registryKey, (agg) -> {
+                try {
+                    agg.build(null, null, null, 0, null, 0, 0, null, null, CardinalityUpperBound.NONE, null);
+                } catch (IOException e) {
+                    fail("Unexpected exception: " + e.getMessage());
+                }
+            }, registryKey.getName(), "geo_shape");
+        }
+    }
+
+    public void testCartesianShapeCentroidLicenseCheck() {
+        checkLicenseRequired(CartesianShapeValuesSourceType.instance(), CartesianCentroidAggregationBuilder.REGISTRY_KEY, (agg) -> {
+            try {
+                agg.build(null, null, null, null, null);
+            } catch (IOException e) {
+                fail("Unexpected exception: " + e.getMessage());
+            }
+        }, "cartesian_centroid", "shape");
+    }
+
+    public void testCartesianPointCentroidLicenseCheck() {
+        checkLicenseNotRequired(CartesianPointValuesSourceType.instance(), CartesianCentroidAggregationBuilder.REGISTRY_KEY, (agg) -> {
+            try {
+                agg.build(null, null, null, null, null);
+            } catch (IOException e) {
+                fail("Unexpected exception: " + e.getMessage());
+            }
+        }, "cartesian_centroid", "point");
+    }
+
+    public void testCartesianPointBoundsLicenseCheck() {
+        CartesianPointValuesSourceType sourceType = CartesianPointValuesSourceType.instance();
+        TestValuesSourceConfig sourceConfig = new TestValuesSourceConfig(sourceType);
+        checkLicenseNotRequired(sourceType, CartesianBoundsAggregationBuilder.REGISTRY_KEY, (agg) -> {
+            try {
+                agg.build(null, null, null, sourceConfig, null);
+            } catch (IOException e) {
+                fail("Unexpected exception: " + e.getMessage());
+            }
+        }, "cartesian_bounds", "point");
+    }
+
+    public void testCartesianShapeBoundsLicenseCheck() {
+        CartesianShapeValuesSourceType sourceType = CartesianShapeValuesSourceType.instance();
+        TestValuesSourceConfig sourceConfig = new TestValuesSourceConfig(sourceType);
+        checkLicenseNotRequired(sourceType, CartesianBoundsAggregationBuilder.REGISTRY_KEY, (agg) -> {
+            try {
+                agg.build(null, null, null, sourceConfig, null);
+            } catch (IOException e) {
+                fail("Unexpected exception: " + e.getMessage());
+            }
+        }, "cartesian_bounds", "shape");
+    }
+
+    private SpatialPlugin getPluginWithOperationMode(License.OperationMode operationMode) {
+        return new SpatialPlugin() {
+            protected XPackLicenseState getLicenseState() {
+                TestUtils.UpdatableLicenseState licenseState = new TestUtils.UpdatableLicenseState();
+                licenseState.update(operationMode, true, null);
+                return licenseState;
+            }
+        };
+    }
+
+    private <T> void checkLicenseNotRequired(
+        ValuesSourceType sourceType,
+        ValuesSourceRegistry.RegistryKey<T> registryKey,
+        Consumer<T> builder,
+        String aggName,
+        String fieldTypeName
+    ) {
         for (License.OperationMode operationMode : License.OperationMode.values()) {
             SpatialPlugin plugin = getPluginWithOperationMode(operationMode);
             ValuesSourceRegistry.Builder registryBuilder = new ValuesSourceRegistry.Builder();
             List<Consumer<ValuesSourceRegistry.Builder>> registrar = plugin.getAggregationExtentions();
             registrar.forEach(c -> c.accept(registryBuilder));
+            List<SearchPlugin.AggregationSpec> specs = plugin.getAggregations();
+            specs.forEach(c -> c.getAggregatorRegistrar().accept(registryBuilder));
             ValuesSourceRegistry registry = registryBuilder.build();
-            MetricAggregatorSupplier centroidSupplier = registry.getAggregator(
-                GeoCentroidAggregationBuilder.REGISTRY_KEY,
-                new ValuesSourceConfig(GeoShapeValuesSourceType.instance(), null, true, null, null, null, null, null, null)
+            T aggregator = registry.getAggregator(
+                registryKey,
+                new ValuesSourceConfig(sourceType, null, true, null, null, null, null, null, null)
+            );
+            NullPointerException exception = expectThrows(NullPointerException.class, () -> builder.accept(aggregator));
+            assertThat(
+                "Incorrect exception testing " + aggName + " on field " + fieldTypeName,
+                exception.getMessage(),
+                containsString("because \"context\" is null")
             );
-            if (License.OperationMode.TRIAL != operationMode
-                && License.OperationMode.compare(operationMode, License.OperationMode.GOLD) < 0) {
-                ElasticsearchSecurityException exception = expectThrows(
-                    ElasticsearchSecurityException.class,
-                    () -> centroidSupplier.build(null, null, null, null, null)
-                );
-                assertThat(
-                    exception.getMessage(),
-                    equalTo("current license is non-compliant for [geo_centroid aggregation on geo_shape fields]")
-                );
-            }
         }
     }
 
-    public void testGeoHexLicenseCheck() {
+    private <T> void checkLicenseRequired(
+        ValuesSourceType sourceType,
+        ValuesSourceRegistry.RegistryKey<T> registryKey,
+        Consumer<T> builder,
+        String aggName,
+        String fieldTypeName
+    ) {
         for (License.OperationMode operationMode : License.OperationMode.values()) {
             SpatialPlugin plugin = getPluginWithOperationMode(operationMode);
             ValuesSourceRegistry.Builder registryBuilder = new ValuesSourceRegistry.Builder();
+            List<Consumer<ValuesSourceRegistry.Builder>> registrar = plugin.getAggregationExtentions();
+            registrar.forEach(c -> c.accept(registryBuilder));
             List<SearchPlugin.AggregationSpec> specs = plugin.getAggregations();
             specs.forEach(c -> c.getAggregatorRegistrar().accept(registryBuilder));
             ValuesSourceRegistry registry = registryBuilder.build();
-            GeoGridAggregatorSupplier hexSupplier = registry.getAggregator(
-                GeoHexGridAggregationBuilder.REGISTRY_KEY,
-                new ValuesSourceConfig(CoreValuesSourceType.GEOPOINT, null, true, null, null, null, null, null, null)
+            T aggregator = registry.getAggregator(
+                registryKey,
+                new ValuesSourceConfig(sourceType, null, true, null, null, null, null, null, null)
             );
             if (License.OperationMode.TRIAL != operationMode
                 && License.OperationMode.compare(operationMode, License.OperationMode.GOLD) < 0) {
                 ElasticsearchSecurityException exception = expectThrows(
                     ElasticsearchSecurityException.class,
-                    () -> hexSupplier.build(
-                        null,
-                        AggregatorFactories.EMPTY,
-                        null,
-                        0,
-                        null,
-                        0,
-                        0,
-                        null,
-                        null,
-                        CardinalityUpperBound.NONE,
-                        null
-                    )
+                    () -> builder.accept(aggregator)
                 );
                 assertThat(
                     exception.getMessage(),
-                    equalTo("current license is non-compliant for [geohex_grid aggregation on geo_point fields]")
-                );
-            }
-        }
-    }
-
-    public void testGeoGridLicenseCheck() {
-        for (ValuesSourceRegistry.RegistryKey<GeoGridAggregatorSupplier> registryKey : Arrays.asList(
-            GeoHashGridAggregationBuilder.REGISTRY_KEY,
-            GeoTileGridAggregationBuilder.REGISTRY_KEY
-        )) {
-            for (License.OperationMode operationMode : License.OperationMode.values()) {
-                SpatialPlugin plugin = getPluginWithOperationMode(operationMode);
-                ValuesSourceRegistry.Builder registryBuilder = new ValuesSourceRegistry.Builder();
-                List<Consumer<ValuesSourceRegistry.Builder>> registrar = plugin.getAggregationExtentions();
-                registrar.forEach(c -> c.accept(registryBuilder));
-                ValuesSourceRegistry registry = registryBuilder.build();
-                GeoGridAggregatorSupplier supplier = registry.getAggregator(
-                    registryKey,
-                    new ValuesSourceConfig(GeoShapeValuesSourceType.instance(), null, true, null, null, null, null, null, null)
+                    equalTo("current license is non-compliant for [" + aggName + " aggregation on " + fieldTypeName + " fields]")
                 );
-                if (License.OperationMode.TRIAL != operationMode
-                    && License.OperationMode.compare(operationMode, License.OperationMode.GOLD) < 0) {
-                    ElasticsearchSecurityException exception = expectThrows(
-                        ElasticsearchSecurityException.class,
-                        () -> supplier.build(null, null, null, 0, null, 0, 0, null, null, CardinalityUpperBound.NONE, null)
-                    );
-                    assertThat(
-                        exception.getMessage(),
-                        equalTo("current license is non-compliant for [" + registryKey.getName() + " aggregation on geo_shape fields]")
-                    );
+            } else {
+                try {
+                    builder.accept(aggregator);
+                } catch (NullPointerException e) {
+                    // Expected exception from passing null aggregation context
+                } catch (Exception e) {
+                    fail("Unexpected exception testing " + aggName + " at license level " + operationMode + ": " + e.getMessage());
                 }
             }
         }
     }
 
-    private SpatialPlugin getPluginWithOperationMode(License.OperationMode operationMode) {
-        return new SpatialPlugin() {
-            protected XPackLicenseState getLicenseState() {
-                TestUtils.UpdatableLicenseState licenseState = new TestUtils.UpdatableLicenseState();
-                licenseState.update(operationMode, true, null);
-                return licenseState;
-            }
-        };
+    private static class TestValuesSourceConfig extends ValuesSourceConfig {
+        private TestValuesSourceConfig(ValuesSourceType sourceType) {
+            super(sourceType, null, true, null, null, null, null, null, null);
+        }
     }
 }