Support zone context variants for native properties

3 files changed, 33 insertions, 17 deletions

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