Add properties to Session

presto-client/src/main/java/com/facebook/presto/client/PrestoHeaders.java

@@ -21,6 +21,7 @@ public final class PrestoHeaders
     public static final String PRESTO_SCHEMA = "X-Presto-Schema";
     public static final String PRESTO_TIME_ZONE = "X-Presto-Time-Zone";
     public static final String PRESTO_LANGUAGE = "X-Presto-Language";
+    public static final String PRESTO_SESSION = "X-Presto-Session";
 
     public static final String PRESTO_CURRENT_STATE = "X-Presto-Current-State";
     public static final String PRESTO_MAX_WAIT = "X-Presto-Max-Wait";

presto-main/src/main/java/com/facebook/presto/Session.java

@@ -18,12 +18,17 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Objects;
+import com.google.common.collect.ImmutableMap;
 
 import javax.annotation.Nullable;
 
+import java.util.HashMap;
 import java.util.Locale;
+import java.util.Map;
 import java.util.TimeZone;
 
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
 import static java.util.Objects.requireNonNull;
 
 public final class Session
@@ -40,6 +45,8 @@ public final class Session
     @Nullable
     private final String userAgent;
     private final long startTime;
+    private final Map<String, String> systemProperties;
+    private final Map<String, Map<String, String>> catalogProperties;
 
     @JsonCreator
     public Session(
@@ -51,7 +58,9 @@ public Session(
             @JsonProperty("locale") Locale locale,
             @JsonProperty("remoteUserAddress") @Nullable String remoteUserAddress,
             @JsonProperty("userAgent") @Nullable String userAgent,
-            @JsonProperty("startTime") long startTime)
+            @JsonProperty("startTime") long startTime,
+            @JsonProperty("systemProperties") @Nullable Map<String, String> systemProperties,
+            @JsonProperty("catalogProperties") @Nullable Map<String, Map<String, String>> catalogProperties)
     {
         this.user = requireNonNull(user, "user is null");
         this.source = source;
@@ -62,6 +71,8 @@ public Session(
         this.remoteUserAddress = remoteUserAddress;
         this.userAgent = userAgent;
         this.startTime = startTime;
+        this.systemProperties = systemProperties != null ? ImmutableMap.copyOf(systemProperties) : ImmutableMap.<String, String>of();
+        this.catalogProperties = catalogProperties != null ? ImmutableMap.copyOf(catalogProperties) : ImmutableMap.<String, Map<String, String>>of();
     }
 
     @JsonProperty
@@ -121,6 +132,18 @@ public long getStartTime()
         return startTime;
     }
 
+    @JsonProperty
+    public Map<String, String> getSystemProperties()
+    {
+        return systemProperties;
+    }
+
+    @JsonProperty
+    public Map<String, Map<String, String>> getCatalogProperties()
+    {
+        return catalogProperties;
+    }
+
     public ConnectorSession toConnectorSession()
     {
         return new ConnectorSession(user, timeZoneKey, locale, startTime);
@@ -158,6 +181,8 @@ public static class SessionBuilder
         private String remoteUserAddress;
         private String userAgent;
         private long startTime = System.currentTimeMillis();
+        private Map<String, String> systemProperties = ImmutableMap.of();
+        private final Map<String, Map<String, String>> catalogProperties = new HashMap<>();
 
         private SessionBuilder()
         {
@@ -217,9 +242,32 @@ public SessionBuilder setUserAgent(String userAgent)
             return this;
         }
 
+        /**
+         * Sets the system properties for the session.  The property names and
+         * values must only contain characters from US-ASCII and must not be for '='.
+         */
+        public SessionBuilder setSystemProperties(Map<String, String> systemProperties)
+        {
+            this.systemProperties = ImmutableMap.copyOf(systemProperties);
+            return this;
+        }
+
+        /**
+         * Sets the properties for a catalog.  The catalog name, property names, and
+         * values must only contain characters from US-ASCII and must not be for '='.
+         */
+        public SessionBuilder setCatalogProperties(String catalog, Map<String, String> properties)
+        {
+            checkNotNull(catalog, "catalog is null");
+            checkArgument(!catalog.isEmpty(), "catalog is empty");
+
+            catalogProperties.put(catalog, ImmutableMap.copyOf(properties));
+            return this;
+        }
+
         public Session build()
         {
-            return new Session(user, source, catalog, schema, timeZoneKey, locale, remoteUserAddress, userAgent, startTime);
+            return new Session(user, source, catalog, schema, timeZoneKey, locale, remoteUserAddress, userAgent, startTime, systemProperties, catalogProperties);
         }
     }
 }

presto-main/src/main/java/com/facebook/presto/server/ResourceUtil.java

@@ -17,18 +17,29 @@
 import com.facebook.presto.Session.SessionBuilder;
 import com.facebook.presto.spi.type.TimeZoneKey;
 import com.facebook.presto.spi.type.TimeZoneNotSupportedException;
+import com.google.common.base.Splitter;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Multimap;
 
 import javax.servlet.http.HttpServletRequest;
 import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.core.Response.Status;
 
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
 import java.util.Locale;
+import java.util.Map;
+import java.util.Map.Entry;
 
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_CATALOG;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_LANGUAGE;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_SCHEMA;
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_SESSION;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_SOURCE;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_TIME_ZONE;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_USER;
@@ -62,9 +73,57 @@ public static Session createSessionForRequest(HttpServletRequest servletRequest)
             sessionBuilder.setLocale(Locale.forLanguageTag(language));
         }
 
+        // parse session properties
+        Multimap<String, Entry<String, String>> sessionPropertiesByCatalog = HashMultimap.create();
+        for (String sessionHeader : Collections.list(servletRequest.getHeaders(PRESTO_SESSION))) {
+            parseSessionHeader(sessionHeader, sessionPropertiesByCatalog);
+        }
+        sessionBuilder.setSystemProperties(toMap(sessionPropertiesByCatalog.get(null)));
+        for (Entry<String, Collection<Entry<String, String>>> entry : sessionPropertiesByCatalog.asMap().entrySet()) {
+            if (entry.getKey() != null) {
+                sessionBuilder.setCatalogProperties(entry.getKey(), toMap(entry.getValue()));
+            }
+        }
+
         return sessionBuilder.build();
     }
 
+    private static void parseSessionHeader(String header, Multimap<String, Entry<String, String>> sessionPropertiesByCatalog)
+    {
+        List<String> nameValue = Splitter.on('=').limit(2).splitToList(header);
+        assertRequest(nameValue.size() == 2, "Invalid %s header", PRESTO_SESSION);
+
+        String catalog;
+        String name;
+        List<String> nameParts = Splitter.on('.').splitToList(nameValue.get(0));
+        if (nameParts.size() == 1) {
+            catalog = null;
+            name = nameParts.get(0);
+        }
+        else if (nameParts.size() == 2) {
+            catalog = nameParts.get(0);
+            name = nameParts.get(1);
+        }
+        else {
+            throw badRequest(format("Invalid %s header", PRESTO_SESSION));
+        }
+        assertRequest(catalog == null || !catalog.isEmpty(), "Invalid %s header", PRESTO_SESSION);
+        assertRequest(!name.isEmpty(), "Invalid %s header", PRESTO_SESSION);
+
+        String value = nameValue.get(1);
+
+        sessionPropertiesByCatalog.put(catalog, Maps.immutableEntry(name, value));
+    }
+
+    private static <K, V> Map<K, V> toMap(Iterable<? extends Entry<K, V>> entries)
+    {
+        ImmutableMap.Builder<K, V> builder = ImmutableMap.builder();
+        for (Entry<K, V> entry : entries) {
+            builder.put(entry);
+        }
+        return builder.build();
+    }
+
     private static String getRequiredHeader(HttpServletRequest servletRequest, String name, String description)
     {
         String user = servletRequest.getHeader(name);

presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java

@@ -14,6 +14,7 @@
 package com.facebook.presto.server.testing;
 
 import com.facebook.presto.connector.ConnectorManager;
+import com.facebook.presto.execution.QueryManager;
 import com.facebook.presto.metadata.AllNodes;
 import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.metadata.Metadata;
@@ -78,6 +79,7 @@ public class TestingPrestoServer
     private final InternalNodeManager nodeManager;
     private final ServiceSelectorManager serviceSelectorManager;
     private final Announcer announcer;
+    private QueryManager queryManager;
 
     public TestingPrestoServer()
             throws Exception
@@ -147,6 +149,8 @@ public TestingPrestoServer(boolean coordinator, Map<String, String> properties,
 
         lifeCycleManager = injector.getInstance(LifeCycleManager.class);
 
+        queryManager = injector.getInstance(QueryManager.class);
+
         pluginManager = injector.getInstance(PluginManager.class);
 
         connectorManager = injector.getInstance(ConnectorManager.class);
@@ -183,6 +187,11 @@ public void installPlugin(Plugin plugin)
         pluginManager.installPlugin(plugin);
     }
 
+    public QueryManager getQueryManager()
+    {
+        return queryManager;
+    }
+
     public void createCatalog(String catalogName, String connectorName)
     {
         createCatalog(catalogName, connectorName, ImmutableMap.<String, String>of());

presto-main/src/test/java/com/facebook/presto/metadata/TestJsonTableHandle.java

@@ -60,6 +60,8 @@ public class TestJsonTableHandle
                     .put("remoteUserAddress", TEST_SESSION.getRemoteUserAddress())
                     .put("userAgent", TEST_SESSION.getUserAgent())
                     .put("startTime", TEST_SESSION.getStartTime())
+                    .put("systemProperties", ImmutableMap.of())
+                    .put("catalogProperties", ImmutableMap.of())
                     .build(),
             "catalogName", "information_schema_catalog",
             "schemaName", "information_schema_schema",

presto-main/src/test/java/com/facebook/presto/server/TestServer.java

@@ -13,17 +13,37 @@
  */
 package com.facebook.presto.server;
 
+import com.facebook.presto.client.QueryResults;
+import com.facebook.presto.execution.QueryId;
+import com.facebook.presto.execution.QueryInfo;
 import com.facebook.presto.server.testing.TestingPrestoServer;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import io.airlift.http.client.HttpClient;
+import io.airlift.http.client.HttpUriBuilder;
+import io.airlift.http.client.Request;
 import io.airlift.http.client.StatusResponseHandler;
 import io.airlift.http.client.jetty.JettyHttpClient;
 import io.airlift.testing.Closeables;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
+import java.net.URI;
+import java.util.List;
+
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_CATALOG;
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_SCHEMA;
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_SESSION;
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_SOURCE;
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_USER;
+import static com.google.common.base.Charsets.UTF_8;
+import static io.airlift.http.client.JsonResponseHandler.createJsonResponseHandler;
 import static io.airlift.http.client.Request.Builder.prepareGet;
+import static io.airlift.http.client.Request.Builder.preparePost;
+import static io.airlift.http.client.StaticBodyGenerator.createStaticBodyGenerator;
 import static io.airlift.http.client.StatusResponseHandler.createStatusResponseHandler;
+import static io.airlift.json.JsonCodec.jsonCodec;
 import static javax.ws.rs.core.Response.Status.OK;
 import static org.testng.Assert.assertEquals;
 
@@ -59,4 +79,55 @@ public void testServerStarts()
 
         assertEquals(response.getStatusCode(), OK.getStatusCode());
     }
+
+    @Test
+    public void testQuery()
+            throws Exception
+    {
+        // start query
+        Request request = preparePost()
+                .setUri(uriFor("/v1/statement"))
+                .setBodyGenerator(createStaticBodyGenerator("show catalogs", UTF_8))
+                .setHeader(PRESTO_USER, "user")
+                .setHeader(PRESTO_SOURCE, "source")
+                .setHeader(PRESTO_CATALOG, "catalog")
+                .setHeader(PRESTO_SCHEMA, "schema")
+                .addHeader(PRESTO_SESSION, "system-name=system-value")
+                .addHeader(PRESTO_SESSION, "catalog.name=catalog-value")
+                .addHeader(PRESTO_SESSION, "catalog2.name=catalog2-value")
+                .build();
+
+        QueryResults queryResults = client.execute(request, createJsonResponseHandler(jsonCodec(QueryResults.class)));
+
+        // get the query info
+        QueryInfo queryInfo = server.getQueryManager().getQueryInfo(new QueryId(queryResults.getId()));
+
+        // verify session properties
+        assertEquals(queryInfo.getSession().getSystemProperties(), ImmutableMap.of("system-name", "system-value"));
+        assertEquals(queryInfo.getSession().getCatalogProperties(), ImmutableMap.of(
+                "catalog", ImmutableMap.of("name", "catalog-value"),
+                "catalog2", ImmutableMap.of("name", "catalog2-value")));
+
+        ImmutableList.Builder<List<Object>> data = ImmutableList.builder();
+        if (queryResults.getData() != null) {
+            data.addAll(queryResults.getData());
+        }
+
+        while (queryResults.getNextUri() != null) {
+            queryResults = client.execute(prepareGet().setUri(queryResults.getNextUri()).build(), createJsonResponseHandler(jsonCodec(QueryResults.class)));
+
+            if (queryResults.getData() != null) {
+                data.addAll(queryResults.getData());
+            }
+        }
+
+        // no catalogs in a the server by default
+        List<List<Object>> rows = data.build();
+        assertEquals(rows.size(), 0);
+    }
+
+    public URI uriFor(String path)
+    {
+        return HttpUriBuilder.uriBuilderFrom(server.getBaseUrl()).replacePath(path).build();
+    }
 }