Support zone context variants for native properties

5 files changed, 98 insertions, 24 deletions

diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index 7fe1702be8f..9f8fb496a4d 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -5992,6 +5992,7 @@
       "public void <init>(com.yahoo.search.query.profile.compiled.CompiledQueryProfile)",
       "public void <init>(com.yahoo.search.query.profile.compiled.CompiledQueryProfile, com.yahoo.language.process.Embedder)",
       "public void <init>(com.yahoo.search.query.profile.compiled.CompiledQueryProfile, java.util.Map)",
+      "public void <init>(com.yahoo.search.query.profile.compiled.CompiledQueryProfile, java.util.Map, ai.vespa.cloud.ZoneInfo)",
       "public com.yahoo.search.query.profile.compiled.CompiledQueryProfile getQueryProfile()",
       "public java.lang.Object get(com.yahoo.processing.request.CompoundName, java.util.Map, com.yahoo.processing.request.Properties)",
       "public void set(com.yahoo.processing.request.CompoundName, java.lang.Object, java.util.Map)",
@@ -6687,7 +6688,7 @@
       "public"
     ],
     "methods": [
-      "public void <init>(java.util.Map, ai.vespa.cloud.ZoneInfo)",
+      "public void <init>(java.util.Map)",
       "public java.lang.Object get(com.yahoo.processing.request.CompoundName, java.util.Map, com.yahoo.processing.request.Properties)",
       "public void set(com.yahoo.processing.request.CompoundName, java.lang.Object, java.util.Map)",
       "public java.util.Map listProperties(com.yahoo.processing.request.CompoundName, java.util.Map, com.yahoo.processing.request.Properties)"
diff --git a/container-search/src/main/java/com/yahoo/search/Query.java b/container-search/src/main/java/com/yahoo/search/Query.java
index dcbda80ecab..dc731669a48 100644
--- a/container-search/src/main/java/com/yahoo/search/Query.java
+++ b/container-search/src/main/java/com/yahoo/search/Query.java
@@ -369,7 +369,7 @@ public class Query extends com.yahoo.processing.Request implements Cloneable {
         startTime = httpRequest.getJDiscRequest().creationTime(TimeUnit.MILLISECONDS);
         if (queryProfile != null) {
             // Move all request parameters to the query profile
-            Properties queryProfileProperties = new QueryProfileProperties(queryProfile, embedders);
+            Properties queryProfileProperties = new QueryProfileProperties(queryProfile, embedders, zoneInfo);
             properties().chain(queryProfileProperties);
             setPropertiesFromRequestMap(requestMap, properties(), true);
 
@@ -377,11 +377,11 @@ public class Query extends com.yahoo.processing.Request implements Cloneable {
             properties().chain(new RankProfileInputProperties(schemaInfo, this, embedders))
                         .chain(new QueryProperties(this, queryProfile.getRegistry(), embedders))
                         .chain(new ModelObjectMap())
-                        .chain(new RequestContextProperties(requestMap, zoneInfo))
+                        .chain(new RequestContextProperties(requestMap))
                         .chain(queryProfileProperties)
                         .chain(new DefaultProperties());
 
-            // Pass the values from the query profile which maps through a field in the Query object model
+            // Pass values from the query profile which maps to a field in the Query object model
             // through the property chain to cause those values to be set in the Query object model with
             // the right types according to query profiles
             setFieldsFrom(queryProfileProperties, requestMap);
diff --git a/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfileProperties.java b/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfileProperties.java
index abd23c1822d..8c1a0ac1d25 100644
--- a/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfileProperties.java
+++ b/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfileProperties.java
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.search.query.profile;
 
+import ai.vespa.cloud.ZoneInfo;
 import com.yahoo.collections.Pair;
 import com.yahoo.language.process.Embedder;
 import com.yahoo.processing.IllegalInputException;
@@ -32,6 +33,7 @@ public class QueryProfileProperties extends Properties {
 
     private final CompiledQueryProfile profile;
     private final Map<String, Embedder> embedders;
+    private final ZoneInfo zoneInfo;
 
     // Note: The priority order is: values has precedence over references
 
@@ -46,18 +48,26 @@ public class QueryProfileProperties extends Properties {
     private List<Pair<CompoundName, CompiledQueryProfile>> references = null;
 
     public QueryProfileProperties(CompiledQueryProfile profile) {
-        this(profile, Embedder.throwsOnUse.asMap());
+        this(profile, Embedder.throwsOnUse.asMap(), ZoneInfo.defaultInfo());
     }
 
+    @Deprecated // TODO: Remove on Vespa 9
     public QueryProfileProperties(CompiledQueryProfile profile, Embedder embedder) {
-        this(profile, Map.of(Embedder.defaultEmbedderId, embedder));
+        this(profile, Map.of(Embedder.defaultEmbedderId, embedder), ZoneInfo.defaultInfo());
     }
 
     /** Creates an instance from a profile, throws an exception if the given profile is null */
+    @Deprecated // TODO: Remove on Vespa 9
     public QueryProfileProperties(CompiledQueryProfile profile, Map<String, Embedder> embedders) {
+        this(profile, embedders, ZoneInfo.defaultInfo());
+    }
+
+    /** Creates an instance from a profile, throws an exception if the given profile is null */
+    public QueryProfileProperties(CompiledQueryProfile profile, Map<String, Embedder> embedders, ZoneInfo zoneInfo) {
         Validator.ensureNotNull("The profile wrapped by this cannot be null", profile);
         this.profile = profile;
         this.embedders = embedders;
+        this.zoneInfo = zoneInfo;
     }
 
     /** Returns the query profile backing this, or null if none */
@@ -67,6 +77,7 @@ public class QueryProfileProperties extends Properties {
     @Override
     public Object get(CompoundName name, Map<String, String> context,
                       com.yahoo.processing.request.Properties substitution) {
+        context = contextWithZoneInfo(context);
         name = unalias(name, context);
         if (values != null && values.containsKey(name))
             return values.get(name); // Returns this value, even if null
@@ -92,11 +103,13 @@ public class QueryProfileProperties extends Properties {
      */
     @Override
     public void set(CompoundName name, Object value, Map<String, String> context) {
+        context = contextWithZoneInfo(context);
         setOrCheckSettable(name, value, context, true);
     }
 
     @Override
     public void requireSettable(CompoundName name, Object value, Map<String, String> context) {
+        context = contextWithZoneInfo(context);
         setOrCheckSettable(name, value, context, false);
     }
 
@@ -210,6 +223,8 @@ public class QueryProfileProperties extends Properties {
     @Override
     public Map<String, Object> listProperties(CompoundName path, Map<String, String> context,
                                               com.yahoo.processing.request.Properties substitution) {
+        context = contextWithZoneInfo(context);
+
         path = unalias(path, context);
         if (context == null) context = Collections.emptyMap();
 
@@ -257,7 +272,7 @@ public class QueryProfileProperties extends Properties {
         return properties;
     }
 
-    public boolean isComplete(StringBuilder firstMissingName, Map<String,String> context) {
+    public boolean isComplete(StringBuilder firstMissingName, Map<String, String> context) {
         // Are all types reachable from this complete?
         if ( ! reachableTypesAreComplete(CompoundName.empty, profile, firstMissingName, context))
             return false;
@@ -272,6 +287,16 @@ public class QueryProfileProperties extends Properties {
         return true;
     }
 
+    private Map<String, String> contextWithZoneInfo(Map<String, String> context) {
+        if (zoneInfo == ZoneInfo.defaultInfo()) return context;
+
+        Map<String, String> contextWithZoneInfo = context == null ? new HashMap<>() : new HashMap<>(context);
+        contextWithZoneInfo.putIfAbsent("environment", zoneInfo.zone().environment().name());
+        contextWithZoneInfo.putIfAbsent("region", zoneInfo.zone().region());
+        contextWithZoneInfo.putIfAbsent("instance", zoneInfo.application().instance());
+        return Collections.unmodifiableMap(contextWithZoneInfo);
+    }
+
     private boolean reachableTypesAreComplete(CompoundName prefix, CompiledQueryProfile profile, StringBuilder firstMissingName, Map<String,String> context) {
         for (Map.Entry<CompoundName, DimensionalValue<QueryProfileType>> typeEntry : profile.getTypes().entrySet()) {
             QueryProfileType type = typeEntry.getValue().get(context);
diff --git a/container-search/src/main/java/com/yahoo/search/query/properties/RequestContextProperties.java b/container-search/src/main/java/com/yahoo/search/query/properties/RequestContextProperties.java
index 46aafcb11e3..ad212d05780 100644
--- a/container-search/src/main/java/com/yahoo/search/query/properties/RequestContextProperties.java
+++ b/container-search/src/main/java/com/yahoo/search/query/properties/RequestContextProperties.java
@@ -21,17 +21,8 @@ public class RequestContextProperties extends Properties {
 
     private final Map<String, String> context;
 
-    public RequestContextProperties(Map<String, String> properties, ZoneInfo zoneInfo) {
-        if (zoneInfo == ZoneInfo.defaultInfo()) {
-            context = properties;
-        }
-        else {
-            Map<String, String> context = new HashMap<>(properties);
-            context.putIfAbsent("environment", zoneInfo.zone().environment().name());
-            context.putIfAbsent("region", zoneInfo.zone().region());
-            context.putIfAbsent("instance", zoneInfo.application().instance());
-            this.context = Collections.unmodifiableMap(context);
-        }
+    public RequestContextProperties(Map<String, String> properties) {
+        this.context = Collections.unmodifiableMap(properties);
     }
 
     @Override
diff --git a/container-search/src/test/java/com/yahoo/search/query/profile/test/QueryProfileVariantsTestCase.java b/container-search/src/test/java/com/yahoo/search/query/profile/test/QueryProfileVariantsTestCase.java
index 3542e1413eb..3cf2949f33c 100644
--- a/container-search/src/test/java/com/yahoo/search/query/profile/test/QueryProfileVariantsTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/query/profile/test/QueryProfileVariantsTestCase.java
@@ -1359,33 +1359,90 @@ public class QueryProfileVariantsTestCase {
         CompiledQueryProfileRegistry cRegistry = registry.compile();
         CompiledQueryProfile cTest = cRegistry.findQueryProfile("test");
 
-        assertValueForZone("default", ZoneInfo.defaultInfo(), cTest);
+        assertValueForZone("default", ZoneInfo.defaultInfo(), null, cTest);
         assertValueForZone("prod-region1-instance1",
                            new ZoneInfo(new ApplicationId("tenant1", "application1", "instance1"),
                                         new Zone(Environment.prod, "region1")),
+                           null,
                            cTest);
         assertValueForZone("prod-instance2",
                            new ZoneInfo(new ApplicationId("tenant2", "application2", "instance2"),
                                         new Zone(Environment.prod, "region1")),
+                           null,
                            cTest);
         assertValueForZone("prod-region3",
                            new ZoneInfo(new ApplicationId("tenant3", "application3", "instance3"),
                                         new Zone(Environment.prod, "region3")),
+                           null,
                            cTest);
         assertValueForZone("dev",
                            new ZoneInfo(new ApplicationId("tenant4", "application4", "instance4"),
                                         new Zone(Environment.dev, "region4")),
+                           null,
                            cTest);
     }
 
-    private void assertValueForZone(String expected, ZoneInfo zoneInfo, CompiledQueryProfile cTest) {
-        assertEquals(expected,
-                     new Query.Builder().setQueryProfile(cTest).setZoneInfo(zoneInfo).build().properties().get("value"));
+    @Test
+    public void testZoneInfoInContextWithUnoverridability() {
+        QueryProfileRegistry registry = new QueryProfileRegistry();
+        QueryProfile profile = new QueryProfile("test");
+        profile.setDimensions(new String[] { "instance", "environment", "region" });
+        profile.set("value", "default", registry);
+        profile.set("value", "prod-beta",
+                    toMap("environment=prod", "instance=beta"),
+                    registry);
+        profile.setOverridable("value", false, toMap("environment=prod", "instance=beta"));
+        registry.register(profile);
+
+        CompiledQueryProfileRegistry cRegistry = registry.compile();
+        CompiledQueryProfile cTest = cRegistry.findQueryProfile("test");
+
+        assertValueForZone("prod-beta",
+                           new ZoneInfo(new ApplicationId("tenant1", "application1", "beta"),
+                                        new Zone(Environment.prod, "region1")),
+                           "fromRequest",
+                           cTest);
+    }
+
+
+    private void assertValueForZone(String expected, ZoneInfo zoneInfo, String requestValue, CompiledQueryProfile cTest) {
+        var builder = new Query.Builder().setQueryProfile(cTest).setZoneInfo(zoneInfo);
+        if (requestValue != null)
+            builder.setRequestMap(Map.of("value", requestValue));
+        assertEquals(expected, builder.build().properties().get("value"));
+    }
+
+    @Test
+    public void testZoneInfoInContextSettingNativeProperty() {
+        QueryProfileRegistry registry = new QueryProfileRegistry();
+        QueryProfile profile = new QueryProfile("test");
+        profile.setDimensions(new String[] { "instance", "environment", "region" });
+        profile.set("timeout", "0.3",
+                    toMap("environment=prod", "instance=beta"),
+                    registry);
+        registry.register(profile);
+
+        CompiledQueryProfileRegistry cRegistry = registry.compile();
+        CompiledQueryProfile cTest = cRegistry.findQueryProfile("test");
+
+        assertTimeoutForZone(300,
+                           new ZoneInfo(new ApplicationId("tenant1", "application1", "beta"),
+                                        new Zone(Environment.prod, "region1")),
+                           null,
+                           cTest);
+    }
+
+    private void assertTimeoutForZone(int expected, ZoneInfo zoneInfo, String requestValue, CompiledQueryProfile cTest) {
+        var builder = new Query.Builder().setQueryProfile(cTest).setZoneInfo(zoneInfo);
+        if (requestValue != null)
+            builder.setRequestMap(Map.of("timeout", requestValue));
+        assertEquals(expected, builder.build().getTimeout());
     }
 
     private void assertGet(String expectedValue, String parameter, String[] dimensionValues, QueryProfile profile, CompiledQueryProfile cprofile) {
-        Map<String,String> context=toMap(profile,dimensionValues);
-        assertEquals("Looking up '" + parameter + "' for '" + Arrays.toString(dimensionValues) + "'",expectedValue,cprofile.get(parameter,context));
+        Map<String, String> context = toMap(profile,dimensionValues);
+        assertEquals("Looking up '" + parameter + "' for '" + Arrays.toString(dimensionValues) + "'",
+                     expectedValue, cprofile.get(parameter,context));
     }
 
     public static Map<String,String> toMap(QueryProfile profile, String[] dimensionValues) {