/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.drill.exec.rpc.user;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.common.collect.Lists;
import org.apache.calcite.schema.Schema;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.schema.Table;
import org.apache.calcite.tools.ValidationException;
import org.apache.drill.common.config.DrillConfig;
import org.apache.drill.common.config.DrillProperties;
import org.apache.drill.exec.planner.physical.PlannerSettings;
import org.apache.drill.exec.planner.sql.SchemaUtilites;
import org.apache.drill.exec.planner.sql.handlers.SqlHandlerUtil;
import org.apache.drill.exec.proto.UserBitShared.UserCredentials;
import org.apache.drill.exec.proto.UserProtos.UserProperties;
import org.apache.drill.exec.server.options.OptionManager;
import org.apache.drill.exec.server.options.OptionValue;
import org.apache.drill.exec.server.options.SessionOptionManager;
import com.google.common.collect.Maps;
import org.apache.drill.exec.server.options.SystemOptionManager;
import org.apache.drill.exec.store.AbstractSchema;
import org.apache.drill.exec.store.StorageStrategy;
import org.apache.drill.exec.store.dfs.DrillFileSystem;
import org.apache.drill.exec.store.dfs.WorkspaceSchemaFactory;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
public class UserSession implements AutoCloseable {
private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(UserSession.class);
private boolean supportComplexTypes = false;
private UserCredentials credentials;
private DrillProperties properties;
private OptionManager sessionOptions;
private final AtomicInteger queryCount;
private final String sessionId;
/** Stores list of temporary tables, key is original table name converted to lower case to achieve case-insensitivity,
* value is generated table name. **/
private final ConcurrentMap<String, String> temporaryTables;
/** Stores list of session temporary locations, key is path to location, value is file system associated with location. **/
private final ConcurrentMap<Path, FileSystem> temporaryLocations;
/** On session close deletes all session temporary locations recursively and clears temporary locations list. */
@Override
public void close() {
for (Map.Entry<Path, FileSystem> entry : temporaryLocations.entrySet()) {
Path path = entry.getKey();
FileSystem fs = entry.getValue();
try {
fs.delete(path, true);
logger.info("Deleted session temporary location [{}] from file system [{}]",
path.toUri().getPath(), fs.getUri());
} catch (Exception e) {
logger.warn("Error during session temporary location [{}] deletion from file system [{}]: [{}]",
path.toUri().getPath(), fs.getUri(), e.getMessage());
}
}
temporaryLocations.clear();
}
/**
* Implementations of this interface are allowed to increment queryCount.
* {@link org.apache.drill.exec.work.user.UserWorker} should have a member that implements the interface.
* No other core class should implement this interface. Test classes may implement (see ControlsInjectionUtil).
*/
public interface QueryCountIncrementer {
void increment(final UserSession session);
}
public static class Builder {
private UserSession userSession;
public static Builder newBuilder() {
return new Builder();
}
public Builder withCredentials(UserCredentials credentials) {
userSession.credentials = credentials;
return this;
}
public Builder withOptionManager(OptionManager systemOptions) {
userSession.sessionOptions = new SessionOptionManager(systemOptions, userSession);
return this;
}
public Builder withUserProperties(UserProperties properties) {
userSession.properties = DrillProperties.createFromProperties(properties, false);
return this;
}
public Builder setSupportComplexTypes(boolean supportComplexTypes) {
userSession.supportComplexTypes = supportComplexTypes;
return this;
}
public UserSession build() {
if (userSession.properties.containsKey(DrillProperties.QUOTING_IDENTIFIERS)) {
if (userSession.sessionOptions != null) {
userSession.setSessionOption(PlannerSettings.QUOTING_IDENTIFIERS_KEY,
userSession.properties.getProperty(DrillProperties.QUOTING_IDENTIFIERS));
} else {
logger.warn("User property {} can't be installed as a server option without the session option manager",
DrillProperties.QUOTING_IDENTIFIERS);
}
}
UserSession session = userSession;
userSession = null;
return session;
}
Builder() {
userSession = new UserSession();
}
}
private UserSession() {
queryCount = new AtomicInteger(0);
sessionId = UUID.randomUUID().toString();
temporaryTables = Maps.newConcurrentMap();
temporaryLocations = Maps.newConcurrentMap();
properties = DrillProperties.createEmpty();
}
public boolean isSupportComplexTypes() {
return supportComplexTypes;
}
public OptionManager getOptions() {
return sessionOptions;
}
public UserCredentials getCredentials() {
return credentials;
}
/**
* Replace current user credentials with the given user's credentials. Meant to be called only by a
* {@link InboundImpersonationManager impersonation manager}.
*
* @param impersonationManager impersonation manager making this call
* @param newCredentials user credentials to change to
*/
public void replaceUserCredentials(final InboundImpersonationManager impersonationManager,
final UserCredentials newCredentials) {
Preconditions.checkNotNull(impersonationManager, "User credentials can only be replaced by an" +
" impersonation manager.");
credentials = newCredentials;
}
public String getTargetUserName() {
return properties.getProperty(DrillProperties.IMPERSONATION_TARGET);
}
public void incrementQueryCount(final QueryCountIncrementer incrementer) {
assert incrementer != null;
queryCount.incrementAndGet();
}
public int getQueryCount() {
return queryCount.get();
}
/**
* Update the schema path for the session.
* @param newDefaultSchemaPath New default schema path to set. It could be relative to the current default schema or
* absolute schema.
* @param currentDefaultSchema Current default schema.
* @throws ValidationException If the given default schema path is invalid in current schema tree.
*/
public void setDefaultSchemaPath(String newDefaultSchemaPath, SchemaPlus currentDefaultSchema)
throws ValidationException {
final List<String> newDefaultPathAsList = Lists.newArrayList(newDefaultSchemaPath.split("\\."));
SchemaPlus newDefault;
// First try to find the given schema relative to the current default schema.
newDefault = SchemaUtilites.findSchema(currentDefaultSchema, newDefaultPathAsList);
if (newDefault == null) {
// If we fail to find the schema relative to current default schema, consider the given new default schema path as
// absolute schema path.
newDefault = SchemaUtilites.findSchema(currentDefaultSchema, newDefaultPathAsList);
}
if (newDefault == null) {
SchemaUtilites.throwSchemaNotFoundException(currentDefaultSchema, newDefaultSchemaPath);
}
properties.setProperty(DrillProperties.SCHEMA, SchemaUtilites.getSchemaPath(newDefault));
}
/**
* @return Get current default schema path.
*/
public String getDefaultSchemaPath() {
return properties.getProperty(DrillProperties.SCHEMA, "");
}
/**
* Get default schema from current default schema path and given schema tree.
* @param rootSchema root schema
* @return A {@link org.apache.calcite.schema.SchemaPlus} object.
*/
public SchemaPlus getDefaultSchema(SchemaPlus rootSchema) {
final String defaultSchemaPath = getDefaultSchemaPath();
if (Strings.isNullOrEmpty(defaultSchemaPath)) {
return null;
}
return SchemaUtilites.findSchema(rootSchema, defaultSchemaPath);
}
/**
* Set the option of a session level.
* Note: Option's kind is automatically detected if such option exists.
*
* @param name option name
* @param value option value
*/
public void setSessionOption(String name, String value) {
OptionValue.Kind optionKind = SystemOptionManager.getValidator(name).getKind();
OptionValue optionValue = OptionValue.createOption(optionKind, OptionValue.OptionType.SESSION, name, value);
sessionOptions.setOption(optionValue);
}
/**
* @return unique session identifier
*/
public String getSessionId() { return sessionId; }
/**
* Creates and adds session temporary location if absent using schema configuration.
* Before any actions, checks if passed table schema is valid default temporary workspace.
* Generates temporary table name and stores it's original name as key
* and generated name as value in session temporary tables cache.
* Original temporary name is converted to lower case to achieve case-insensitivity.
* If original table name already exists, new name is not regenerated and is reused.
* This can happen if default temporary workspace was changed (file system or location) or
* orphan temporary table name has remained (name was registered but table creation did not succeed).
*
* @param schema table schema
* @param tableName original table name
* @param config drill config
* @return generated temporary table name
* @throws IOException if error during session temporary location creation
*/
public String registerTemporaryTable(AbstractSchema schema, String tableName, DrillConfig config) throws IOException {
addTemporaryLocation(SchemaUtilites.resolveToValidTemporaryWorkspace(schema, config));
String temporaryTableName = new Path(sessionId, UUID.randomUUID().toString()).toUri().getPath();
String oldTemporaryTableName = temporaryTables.putIfAbsent(tableName.toLowerCase(), temporaryTableName);
return oldTemporaryTableName == null ? temporaryTableName : oldTemporaryTableName;
}
/**
* Returns generated temporary table name from the list of session temporary tables, null otherwise.
* Original temporary name is converted to lower case to achieve case-insensitivity.
*
* @param tableName original table name
* @return generated temporary table name
*/
public String resolveTemporaryTableName(String tableName) {
return temporaryTables.get(tableName.toLowerCase());
}
/**
* Checks if passed table is temporary, table name is case-insensitive.
* Before looking for table checks if passed schema is temporary and returns false if not
* since temporary tables are allowed to be created in temporary workspace only.
* If passed workspace is temporary, looks for temporary table.
* First checks if table name is among temporary tables, if not returns false.
* If temporary table named was resolved, checks that temporary table exists on disk,
* to ensure that temporary table actually exists and resolved table name is not orphan
* (for example, in result of unsuccessful temporary table creation).
*
* @param drillSchema table schema
* @param config drill config
* @param tableName original table name
* @return true if temporary table exists in schema, false otherwise
*/
public boolean isTemporaryTable(AbstractSchema drillSchema, DrillConfig config, String tableName) {
if (drillSchema == null || !SchemaUtilites.isTemporaryWorkspace(drillSchema.getFullSchemaName(), config)) {
return false;
}
String temporaryTableName = resolveTemporaryTableName(tableName);
if (temporaryTableName != null) {
Table temporaryTable = SqlHandlerUtil.getTableFromSchema(drillSchema, temporaryTableName);
if (temporaryTable != null && temporaryTable.getJdbcTableType() == Schema.TableType.TABLE) {
return true;
}
}
return false;
}
/**
* Removes temporary table name from the list of session temporary tables.
* Original temporary name is converted to lower case to achieve case-insensitivity.
* Before temporary table drop, checks if passed table schema is valid default temporary workspace.
*
* @param schema table schema
* @param tableName original table name
* @param config drill config
*/
public void removeTemporaryTable(AbstractSchema schema, String tableName, DrillConfig config) {
String temporaryTable = resolveTemporaryTableName(tableName);
if (temporaryTable == null) {
return;
}
SqlHandlerUtil.dropTableFromSchema(SchemaUtilites.resolveToValidTemporaryWorkspace(schema, config), temporaryTable);
temporaryTables.remove(tableName.toLowerCase());
}
/**
* Session temporary tables are stored under temporary workspace location in session folder
* defined by unique session id. These session temporary locations are deleted on session close.
* If default temporary workspace file system or location is changed at runtime,
* new session temporary location will be added with corresponding file system
* to the list of session temporary locations. If location does not exist it will be created and
* {@link StorageStrategy#TEMPORARY} storage rules will be applied to it.
*
* @param temporaryWorkspace temporary workspace
* @throws IOException in case of error during temporary location creation
*/
private void addTemporaryLocation(WorkspaceSchemaFactory.WorkspaceSchema temporaryWorkspace) throws IOException {
DrillFileSystem fs = temporaryWorkspace.getFS();
Path temporaryLocation = new Path(fs.getUri().toString(),
new Path(temporaryWorkspace.getDefaultLocation(), sessionId));
FileSystem fileSystem = temporaryLocations.putIfAbsent(temporaryLocation, fs);
if (fileSystem == null) {
StorageStrategy.TEMPORARY.createPathAndApply(fs, temporaryLocation);
Preconditions.checkArgument(fs.exists(temporaryLocation),
String.format("Temporary location should exist [%s]", temporaryLocation.toUri().getPath()));
}
}
}