Implement SET PATH statement

prestodb · Jul 10, 2018 · e811d00 · e811d00
1 parent 6a52346
commit e811d00
Show file tree

Hide file tree

Showing 27 changed files with 282 additions and 3 deletions.
diff --git a/...ark-driver/src/main/java/com/facebook/presto/benchmark/driver/BenchmarkDriverOptions.java b/...ark-driver/src/main/java/com/facebook/presto/benchmark/driver/BenchmarkDriverOptions.java
@@ -96,6 +96,7 @@ public ClientSession getClientSession()
                 null,
                 catalog,
                 schema,
+                null,
                 TimeZone.getDefault().getID(),
                 Locale.getDefault(),
                 ImmutableMap.of(),

diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java b/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java
@@ -159,6 +159,7 @@ public ClientSession toClientSession()
                 clientInfo,
                 catalog,
                 schema,
+                null,
                 TimeZone.getDefault().getID(),
                 Locale.getDefault(),
                 toResourceEstimates(resourceEstimates),

diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/Console.java b/presto-cli/src/main/java/com/facebook/presto/cli/Console.java
@@ -320,6 +320,13 @@ private static boolean process(QueryRunner queryRunner, String sql, OutputFormat
                 schemaChanged.run();
             }
 
+            // update path if present
+            if (query.getSetPath().isPresent()) {
+                session = ClientSession.builder(session)
+                        .withPath(query.getSetPath().get())
+                        .build();
+            }
+
             // update session properties if present
             if (!query.getSetSessionProperties().isEmpty() || !query.getResetSessionProperties().isEmpty()) {
                 Map<String, String> sessionProperties = new HashMap<>(session.getProperties());

diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/Query.java b/presto-cli/src/main/java/com/facebook/presto/cli/Query.java
@@ -72,6 +72,11 @@ public Optional<String> getSetSchema()
         return client.getSetSchema();
     }
 
+    public Optional<String> getSetPath()
+    {
+        return client.getSetPath();
+    }
+
     public Map<String, String> getSetSessionProperties()
     {
         return client.getSetSessionProperties();

diff --git a/presto-cli/src/test/java/com/facebook/presto/cli/TestQueryRunner.java b/presto-cli/src/test/java/com/facebook/presto/cli/TestQueryRunner.java
@@ -91,6 +91,7 @@ public void testCookie()
                         "clientInfo",
                         "catalog",
                         "schema",
+                        "path",
                         "America/Los_Angeles",
                         Locale.ENGLISH,
                         ImmutableMap.of(),

diff --git a/presto-client/src/main/java/com/facebook/presto/client/ClientSession.java b/presto-client/src/main/java/com/facebook/presto/client/ClientSession.java
@@ -41,6 +41,7 @@ public class ClientSession
     private final String clientInfo;
     private final String catalog;
     private final String schema;
+    private final String path;
     private final TimeZoneKey timeZone;
     private final Locale locale;
     private final Map<String, String> resourceEstimates;
@@ -70,6 +71,7 @@ public ClientSession(
             String clientInfo,
             String catalog,
             String schema,
+            String path,
             String timeZoneId,
             Locale locale,
             Map<String, String> resourceEstimates,
@@ -86,6 +88,7 @@ public ClientSession(
         this.clientInfo = clientInfo;
         this.catalog = catalog;
         this.schema = schema;
+        this.path = path;
         this.locale = locale;
         this.timeZone = TimeZoneKey.getTimeZoneKey(timeZoneId);
         this.transactionId = transactionId;
@@ -155,6 +158,11 @@ public String getSchema()
         return schema;
     }
 
+    public String getPath()
+    {
+        return path;
+    }
+
     public TimeZoneKey getTimeZone()
     {
         return timeZone;
@@ -205,6 +213,7 @@ public String toString()
                 .add("clientInfo", clientInfo)
                 .add("catalog", catalog)
                 .add("schema", schema)
+                .add("path", path)
                 .add("traceToken", traceToken.orElse(null))
                 .add("timeZone", timeZone)
                 .add("locale", locale)
@@ -224,6 +233,7 @@ public static final class Builder
         private String clientInfo;
         private String catalog;
         private String schema;
+        private String path;
         private TimeZoneKey timeZone;
         private Locale locale;
         private Map<String, String> resourceEstimates;
@@ -243,6 +253,7 @@ private Builder(ClientSession clientSession)
             clientInfo = clientSession.getClientInfo();
             catalog = clientSession.getCatalog();
             schema = clientSession.getSchema();
+            path = clientSession.getPath();
             timeZone = clientSession.getTimeZone();
             locale = clientSession.getLocale();
             resourceEstimates = clientSession.getResourceEstimates();
@@ -264,6 +275,12 @@ public Builder withSchema(String schema)
             return this;
         }
 
+        public Builder withPath(String path)
+        {
+            this.path = requireNonNull(path, "path is null");
+            return this;
+        }
+
         public Builder withProperties(Map<String, String> properties)
         {
             this.properties = requireNonNull(properties, "properties is null");
@@ -299,6 +316,7 @@ public ClientSession build()
                     clientInfo,
                     catalog,
                     schema,
+                    path,
                     timeZone.getId(),
                     locale,
                     resourceEstimates,

diff --git a/presto-client/src/main/java/com/facebook/presto/client/PrestoHeaders.java b/presto-client/src/main/java/com/facebook/presto/client/PrestoHeaders.java
@@ -19,12 +19,14 @@ public final class PrestoHeaders
     public static final String PRESTO_SOURCE = "X-Presto-Source";
     public static final String PRESTO_CATALOG = "X-Presto-Catalog";
     public static final String PRESTO_SCHEMA = "X-Presto-Schema";
+    public static final String PRESTO_PATH = "X-Presto-Path";
     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_TRACE_TOKEN = "X-Presto-Trace-Token";
     public static final String PRESTO_SESSION = "X-Presto-Session";
     public static final String PRESTO_SET_CATALOG = "X-Presto-Set-Catalog";
     public static final String PRESTO_SET_SCHEMA = "X-Presto-Set-Schema";
+    public static final String PRESTO_SET_PATH = "X-Presto-Set-Path";
     public static final String PRESTO_SET_SESSION = "X-Presto-Set-Session";
     public static final String PRESTO_CLEAR_SESSION = "X-Presto-Clear-Session";
     public static final String PRESTO_PREPARED_STATEMENT = "X-Presto-Prepared-Statement";

diff --git a/presto-client/src/main/java/com/facebook/presto/client/StatementClient.java b/presto-client/src/main/java/com/facebook/presto/client/StatementClient.java
@@ -49,6 +49,8 @@ public interface StatementClient
 
     Optional<String> getSetSchema();
 
+    Optional<String> getSetPath();
+
     Map<String, String> getSetSessionProperties();
 
     Set<String> getResetSessionProperties();

diff --git a/presto-client/src/main/java/com/facebook/presto/client/StatementClientV1.java b/presto-client/src/main/java/com/facebook/presto/client/StatementClientV1.java
@@ -53,11 +53,13 @@
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_CLIENT_TAGS;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_DEALLOCATED_PREPARE;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_LANGUAGE;
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_PATH;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_PREPARED_STATEMENT;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_RESOURCE_ESTIMATE;
 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_SET_CATALOG;
+import static com.facebook.presto.client.PrestoHeaders.PRESTO_SET_PATH;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_SET_SCHEMA;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_SET_SESSION;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_SOURCE;
@@ -94,6 +96,7 @@ class StatementClientV1
     private final AtomicReference<QueryResults> currentResults = new AtomicReference<>();
     private final AtomicReference<String> setCatalog = new AtomicReference<>();
     private final AtomicReference<String> setSchema = new AtomicReference<>();
+    private final AtomicReference<String> setPath = new AtomicReference<>();
     private final Map<String, String> setSessionProperties = new ConcurrentHashMap<>();
     private final Set<String> resetSessionProperties = Sets.newConcurrentHashSet();
     private final Map<String, String> addedPreparedStatements = new ConcurrentHashMap<>();
@@ -158,6 +161,9 @@ private Request buildQueryRequest(ClientSession session, String query)
         if (session.getSchema() != null) {
             builder.addHeader(PRESTO_SCHEMA, session.getSchema());
         }
+        if (session.getPath() != null) {
+            builder.addHeader(PRESTO_PATH, session.getPath());
+        }
         builder.addHeader(PRESTO_TIME_ZONE, session.getTimeZone().getId());
         if (session.getLocale() != null) {
             builder.addHeader(PRESTO_LANGUAGE, session.getLocale().toLanguageTag());
@@ -254,6 +260,12 @@ public Optional<String> getSetSchema()
         return Optional.ofNullable(setSchema.get());
     }
 
+    @Override
+    public Optional<String> getSetPath()
+    {
+        return Optional.ofNullable(setPath.get());
+    }
+
     @Override
     public Map<String, String> getSetSessionProperties()
     {
@@ -372,6 +384,7 @@ private void processResponse(Headers headers, QueryResults results)
     {
         setCatalog.set(headers.get(PRESTO_SET_CATALOG));
         setSchema.set(headers.get(PRESTO_SET_SCHEMA));
+        setPath.set(headers.get(PRESTO_SET_PATH));
 
         for (String setSession : headers.values(PRESTO_SET_SESSION)) {
             List<String> keyValue = SESSION_HEADER_SPLITTER.splitToList(setSession);

diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoConnection.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoConnection.java
@@ -73,6 +73,7 @@ public class PrestoConnection
     private final AtomicBoolean readOnly = new AtomicBoolean();
     private final AtomicReference<String> catalog = new AtomicReference<>();
     private final AtomicReference<String> schema = new AtomicReference<>();
+    private final AtomicReference<String> path = new AtomicReference<>();
     private final AtomicReference<String> timeZoneId = new AtomicReference<>();
     private final AtomicReference<Locale> locale = new AtomicReference<>();
     private final AtomicReference<Integer> networkTimeoutMillis = new AtomicReference<>(Ints.saturatedCast(MINUTES.toMillis(2)));
@@ -669,6 +670,7 @@ else if (applicationName != null) {
                 clientInfo.get("ClientInfo"),
                 catalog.get(),
                 schema.get(),
+                path.get(),
                 timeZoneId.get(),
                 locale.get(),
                 ImmutableMap.of(),
@@ -690,6 +692,7 @@ void updateSession(StatementClient client)
 
         client.getSetCatalog().ifPresent(catalog::set);
         client.getSetSchema().ifPresent(schema::set);
+        client.getSetPath().ifPresent(path::set);
 
         if (client.getStartedTransactionId() != null) {
             transactionId.set(client.getStartedTransactionId());

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryInfo.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryInfo.java
@@ -54,6 +54,7 @@ public class QueryInfo
     private final QueryStats queryStats;
     private final Optional<String> setCatalog;
     private final Optional<String> setSchema;
+    private final Optional<String> setPath;
     private final Map<String, String> setSessionProperties;
     private final Set<String> resetSessionProperties;
     private final Map<String, String> addedPreparedStatements;
@@ -83,6 +84,7 @@ public QueryInfo(
             @JsonProperty("queryStats") QueryStats queryStats,
             @JsonProperty("setCatalog") Optional<String> setCatalog,
             @JsonProperty("setSchema") Optional<String> setSchema,
+            @JsonProperty("setPath") Optional<String> setPath,
             @JsonProperty("setSessionProperties") Map<String, String> setSessionProperties,
             @JsonProperty("resetSessionProperties") Set<String> resetSessionProperties,
             @JsonProperty("addedPreparedStatements") Map<String, String> addedPreparedStatements,
@@ -106,6 +108,7 @@ public QueryInfo(
         requireNonNull(queryStats, "queryStats is null");
         requireNonNull(setCatalog, "setCatalog is null");
         requireNonNull(setSchema, "setSchema is null");
+        requireNonNull(setPath, "setPath is null");
         requireNonNull(setSessionProperties, "setSessionProperties is null");
         requireNonNull(resetSessionProperties, "resetSessionProperties is null");
         requireNonNull(addedPreparedStatements, "addedPreparedStatemetns is null");
@@ -128,6 +131,7 @@ public QueryInfo(
         this.queryStats = queryStats;
         this.setCatalog = setCatalog;
         this.setSchema = setSchema;
+        this.setPath = setPath;
         this.setSessionProperties = ImmutableMap.copyOf(setSessionProperties);
         this.resetSessionProperties = ImmutableSet.copyOf(resetSessionProperties);
         this.addedPreparedStatements = ImmutableMap.copyOf(addedPreparedStatements);
@@ -211,6 +215,12 @@ public Optional<String> getSetSchema()
         return setSchema;
     }
 
+    @JsonProperty
+    public Optional<String> getSetPath()
+    {
+        return setPath;
+    }
+
     @JsonProperty
     public Map<String, String> getSetSessionProperties()
     {

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java
@@ -136,6 +136,7 @@ public class QueryStateMachine
 
     private final AtomicReference<String> setCatalog = new AtomicReference<>();
     private final AtomicReference<String> setSchema = new AtomicReference<>();
+    private final AtomicReference<String> setPath = new AtomicReference<>();
 
     private final Map<String, String> setSessionProperties = new ConcurrentHashMap<>();
     private final Set<String> resetSessionProperties = Sets.newConcurrentHashSet();
@@ -482,6 +483,7 @@ public QueryInfo getQueryInfo(Optional<StageInfo> rootStage)
                 queryStats,
                 Optional.ofNullable(setCatalog.get()),
                 Optional.ofNullable(setSchema.get()),
+                Optional.ofNullable(setPath.get()),
                 setSessionProperties,
                 resetSessionProperties,
                 addedPreparedStatements,
@@ -550,6 +552,17 @@ public void setSetSchema(String schema)
         setSchema.set(requireNonNull(schema, "schema is null"));
     }
 
+    public void setSetPath(String path)
+    {
+        requireNonNull(path, "path is null");
+        setPath.set(path);
+    }
+
+    public String getSetPath()
+    {
+        return setPath.get();
+    }
+
     public void addSetSessionProperties(String key, String value)
     {
         setSessionProperties.put(requireNonNull(key, "key is null"), requireNonNull(value, "value is null"));
@@ -864,6 +877,7 @@ public void pruneQueryInfo()
                 pruneQueryStats(queryInfo.getQueryStats()),
                 queryInfo.getSetCatalog(),
                 queryInfo.getSetSchema(),
+                queryInfo.getSetPath(),
                 queryInfo.getSetSessionProperties(),
                 queryInfo.getResetSessionProperties(),
                 queryInfo.getAddedPreparedStatements(),

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SetPathTask.java b/presto-main/src/main/java/com/facebook/presto/execution/SetPathTask.java
@@ -13,14 +13,25 @@
  */
 package com.facebook.presto.execution;
 
+import com.facebook.presto.Session;
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.security.AccessControl;
+import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.sql.SqlPath;
+import com.facebook.presto.sql.SqlPathElement;
+import com.facebook.presto.sql.analyzer.SemanticException;
 import com.facebook.presto.sql.tree.Expression;
 import com.facebook.presto.sql.tree.SetPath;
 import com.facebook.presto.transaction.TransactionManager;
 import com.google.common.util.concurrent.ListenableFuture;
 
 import java.util.List;
+import java.util.Optional;
+
+import static com.facebook.presto.spi.StandardErrorCode.NOT_FOUND;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.CATALOG_NOT_SPECIFIED;
+import static com.google.common.util.concurrent.Futures.immediateFuture;
+import static java.util.Locale.ENGLISH;
 
 public class SetPathTask
         implements DataDefinitionTask<SetPath>
@@ -34,6 +45,24 @@ public String getName()
     @Override
     public ListenableFuture<?> execute(SetPath statement, TransactionManager transactionManager, Metadata metadata, AccessControl accessControl, QueryStateMachine stateMachine, List<Expression> parameters)
     {
-        throw new UnsupportedOperationException("SET PATH not implemented yet");
+        Session session = stateMachine.getSession();
+
+        // convert to IR before setting HTTP headers - ensures that the representations of all path objects outside the parser remain consistent
+        SqlPath sqlPath = new SqlPath(Optional.of(statement.getPathSpecification().toString()));
+
+        for (SqlPathElement element : sqlPath.getParsedPath()) {
+            if (!element.getCatalog().isPresent() && !session.getCatalog().isPresent()) {
+                throw new SemanticException(CATALOG_NOT_SPECIFIED, statement, "Catalog must be specified for each path element when session catalog is not set");
+            }
+
+            element.getCatalog().ifPresent(catalog -> {
+                String catalogName = catalog.getValue().toLowerCase(ENGLISH);
+                if (!metadata.getCatalogHandle(session, catalogName).isPresent()) {
+                    throw new PrestoException(NOT_FOUND, "Catalog does not exist: " + catalogName);
+                }
+            });
+        }
+        stateMachine.setSetPath(sqlPath.toString());
+        return immediateFuture(null);
     }
 }