/**
* Get tablespace for the given URI. If uri is null, the default tablespace will be returned
*
* @param uri Table or Table Fragment URI.
* @param <T> Tablespace class type
* @return Tablespace. If uri is null, the default tablespace will be returned.
*/
public static <T extends Tablespace> T get(@Nullable String uri) {
if (uri == null || uri.isEmpty()) {
return getDefault();
}
Tablespace lastOne = null;
// Find the longest matched one. For example, assume that the caller tries to find /x/y/z, and
// there are /x and /x/y. In this case, /x/y will be chosen because it is more specific.
for (Map.Entry<URI, Tablespace> entry :
TABLE_SPACES.headMap(URI.create(uri), true).entrySet()) {
if (uri.startsWith(entry.getKey().toString())) {
lastOne = entry.getValue();
}
}
if (lastOne == null) {
lastOne =
initializeTableSpace(UUID.randomUUID().toString(), URI.create(uri), new JSONObject());
try {
lastOne.init(systemConf);
} catch (IOException e) {
throw new RuntimeException(e);
}
}
return (T) lastOne;
}
public void executeDistributedQuery(
QueryContext queryContext,
Session session,
LogicalPlan plan,
String sql,
String jsonExpr,
SubmitQueryResponse.Builder responseBuilder)
throws Exception {
LogicalRootNode rootNode = plan.getRootBlock().getRoot();
TableDesc tableDesc = PlannerUtil.getTableDesc(catalog, plan.getRootBlock().getRoot());
if (tableDesc != null) {
Tablespace space = TablespaceManager.get(tableDesc.getUri()).get();
FormatProperty formatProperty = space.getFormatProperty(tableDesc.getMeta());
if (!formatProperty.isInsertable()) {
throw new UnsupportedException(
String.format("INSERT operation on %s tablespace", tableDesc.getUri().toString()));
}
space.prepareTable(rootNode.getChild());
}
hookManager.doHooks(queryContext, plan);
QueryManager queryManager = this.context.getQueryJobManager();
QueryInfo queryInfo;
queryInfo = queryManager.scheduleQuery(session, queryContext, sql, jsonExpr, rootNode);
responseBuilder.setState(OK);
responseBuilder.setQueryId(queryInfo.getQueryId().getProto());
responseBuilder.setResultType(ResultType.FETCH);
if (queryInfo.getQueryMasterHost() != null) {
responseBuilder.setQueryMasterHost(queryInfo.getQueryMasterHost());
}
responseBuilder.setQueryMasterPort(queryInfo.getQueryMasterClientPort());
LOG.info(
"Query "
+ queryInfo.getQueryId().toString()
+ ","
+ queryInfo.getSql()
+ ","
+ " is forwarded to "
+ queryInfo.getQueryMasterHost()
+ ":"
+ queryInfo.getQueryMasterPort());
}
public MasterPlan compileMasterPlan(LogicalPlan plan, QueryContext context, GlobalPlanner planner)
throws Exception {
LogicalRootNode rootNode = plan.getRootBlock().getRoot();
TableDesc tableDesc = PlannerUtil.getTableDesc(planner.getCatalog(), rootNode.getChild());
if (tableDesc != null) {
Tablespace space = TablespaceManager.get(tableDesc.getUri()).get();
space.rewritePlan(context, plan);
}
MasterPlan masterPlan = new MasterPlan(QueryIdFactory.NULL_QUERY_ID, context, plan);
planner.build(context, masterPlan);
return masterPlan;
}
@VisibleForTesting
public static Optional<Tablespace> addTableSpaceForTest(Tablespace space) {
Tablespace existing;
synchronized (SPACES_URIS_MAP) {
String scheme = UriUtil.getScheme(space.getUri());
if (!TABLE_SPACE_HANDLERS.containsKey(scheme)) {
TABLE_SPACE_HANDLERS.put(scheme, space.getClass());
}
// Remove existing one
SPACES_URIS_MAP.remove(space.getName());
existing = TABLE_SPACES.remove(space.getUri());
// Add anotherone for test
registerTableSpace(space.name, space.uri, space.getConfig(), true, true);
}
// if there is an existing one, return it.
return Optional.fromNullable(existing);
}
private static void registerTableSpace(
String spaceName, URI uri, JSONObject spaceDesc, boolean visible, boolean override) {
Tablespace tableSpace = initializeTableSpace(spaceName, uri, spaceDesc);
tableSpace.setVisible(visible);
try {
tableSpace.init(systemConf);
} catch (IOException e) {
throw new RuntimeException(e);
}
putTablespace(tableSpace, override);
// If the arbitrary path is allowed, root uri is also added as a tablespace
if (tableSpace.getProperty().isArbitraryPathAllowed()) {
URI rootUri = tableSpace.getRootUri();
// if there already exists or the rootUri is 'file:/', it won't overwrite the tablespace.
if (!TABLE_SPACES.containsKey(rootUri)
&& !rootUri.toString().startsWith(LOCAL_FS_URI.toString())) {
String tmpName = UUID.randomUUID().toString();
registerTableSpace(tmpName, rootUri, spaceDesc, false, override);
}
}
}
private static void putTablespace(Tablespace space, boolean override) {
// It is a device to keep the relationship among name, URI, and tablespace 1:1:1.
boolean nameExist = SPACES_URIS_MAP.containsKey(space.getName());
boolean uriExist = TABLE_SPACES.containsKey(space.uri);
boolean mismatch = nameExist && !SPACES_URIS_MAP.get(space.getName()).equals(space.getUri());
mismatch = mismatch || uriExist && TABLE_SPACES.get(space.uri).equals(space);
if (!override && mismatch) {
throw new RuntimeException("Name or URI of Tablespace must be unique.");
}
SPACES_URIS_MAP.put(space.getName(), space.getUri());
// We must guarantee that the same uri results in the same tablespace instance.
TABLE_SPACES.put(space.getUri(), space);
}
/** Insert row values */
private void insertRowValues(
QueryContext queryContext,
InsertNode insertNode,
SubmitQueryResponse.Builder responseBuilder) {
try {
String nodeUniqName =
insertNode.getTableName() == null
? new Path(insertNode.getUri()).getName()
: insertNode.getTableName();
String queryId = nodeUniqName + "_" + System.currentTimeMillis();
URI finalOutputUri = insertNode.getUri();
Tablespace space = TablespaceManager.get(finalOutputUri).get();
TableMeta tableMeta = new TableMeta(insertNode.getStorageType(), insertNode.getOptions());
tableMeta.putOption(StorageConstants.INSERT_DIRECTLY, Boolean.TRUE.toString());
FormatProperty formatProperty = space.getFormatProperty(tableMeta);
TaskAttemptContext taskAttemptContext;
if (formatProperty
.directInsertSupported()) { // if this format and storage supports direct insertion
taskAttemptContext = new TaskAttemptContext(queryContext, null, null, null, null);
taskAttemptContext.setOutputPath(new Path(finalOutputUri));
EvalExprExec evalExprExec =
new EvalExprExec(taskAttemptContext, (EvalExprNode) insertNode.getChild());
InsertRowsExec exec = new InsertRowsExec(taskAttemptContext, insertNode, evalExprExec);
try {
exec.init();
exec.next();
} finally {
exec.close();
}
} else {
URI stagingSpaceUri =
space.prepareStagingSpace(context.getConf(), queryId, queryContext, tableMeta);
Path stagingDir = new Path(stagingSpaceUri);
Path stagingResultDir = new Path(stagingDir, TajoConstants.RESULT_DIR_NAME);
taskAttemptContext = new TaskAttemptContext(queryContext, null, null, null, stagingDir);
taskAttemptContext.setOutputPath(new Path(stagingResultDir, "part-01-000000"));
insertRowsThroughStaging(
taskAttemptContext, insertNode, new Path(finalOutputUri), stagingDir, stagingResultDir);
}
// set insert stats (how many rows and bytes)
TableStats stats = new TableStats();
stats.setNumBytes(taskAttemptContext.getResultStats().getNumBytes());
stats.setNumRows(taskAttemptContext.getResultStats().getNumRows());
if (insertNode.hasTargetTable()) {
CatalogProtos.UpdateTableStatsProto.Builder builder =
CatalogProtos.UpdateTableStatsProto.newBuilder();
builder.setTableName(insertNode.getTableName());
builder.setStats(stats.getProto());
catalog.updateTableStats(builder.build());
TableDesc desc =
new TableDesc(
insertNode.getTableName(), insertNode.getTargetSchema(), tableMeta, finalOutputUri);
responseBuilder.setTableDesc(desc.getProto());
} else { // If INSERT INTO LOCATION
// Empty TableDesc
List<CatalogProtos.ColumnProto> columns = new ArrayList<CatalogProtos.ColumnProto>();
CatalogProtos.TableDescProto tableDescProto =
CatalogProtos.TableDescProto.newBuilder()
.setTableName(nodeUniqName)
.setMeta(
CatalogProtos.TableProto.newBuilder()
.setStoreType(BuiltinStorages.TEXT)
.build())
.setSchema(CatalogProtos.SchemaProto.newBuilder().addAllFields(columns).build())
.setStats(stats.getProto())
.build();
responseBuilder.setTableDesc(tableDescProto);
}
// If queryId == NULL_QUERY_ID and MaxRowNum == -1, TajoCli prints only number of inserted
// rows.
responseBuilder.setMaxRowNum(-1);
responseBuilder.setQueryId(QueryIdFactory.NULL_QUERY_ID.getProto());
responseBuilder.setResultType(ResultType.NO_RESULT);
responseBuilder.setState(OK);
} catch (Throwable t) {
throw new RuntimeException(t);
}
}
@Override
public URI getTableURI(@Nullable String spaceName, String databaseName, String tableName) {
Tablespace space = spaceName == null ? getDefault() : getByName(spaceName);
return space.getTableUri(databaseName, tableName);
}
