1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
| public ListenableFuture<?> createQuery(QueryId queryId, Slug slug, SessionContext sessionContext, String query) { requireNonNull(queryId, "queryId is null"); requireNonNull(sessionContext, "sessionContext is null"); requireNonNull(query, "query is null"); checkArgument(!query.isEmpty(), "query must not be empty string"); checkArgument(queryTracker.tryGetQuery(queryId).isEmpty(), "query %s already exists", queryId);
DispatchQueryCreationFuture queryCreationFuture = new DispatchQueryCreationFuture(); dispatchExecutor.execute(() -> { try { createQueryInternal(queryId, slug, sessionContext, query, resourceGroupManager); } finally { queryCreationFuture.set(null); } }); return queryCreationFuture; }
private <C> void createQueryInternal(QueryId queryId, Slug slug, SessionContext sessionContext, String query, ResourceGroupManager<C> resourceGroupManager) { Session session = null; PreparedQuery preparedQuery = null; try { if (query.length() > maxQueryLength) { int queryLength = query.length(); query = query.substring(0, maxQueryLength); throw new TrinoException(QUERY_TEXT_TOO_LARGE, format("Query text length (%s) exceeds the maximum length (%s)", queryLength, maxQueryLength)); }
session = sessionSupplier.createSession(queryId, sessionContext);
accessControl.checkCanExecuteQuery(sessionContext.getIdentity());
preparedQuery = queryPreparer.prepareQuery(session, query);
Optional<String> queryType = getQueryType(preparedQuery.getStatement().getClass()).map(Enum::name); SelectionContext<C> selectionContext = resourceGroupManager.selectGroup(new SelectionCriteria( sessionContext.getIdentity().getPrincipal().isPresent(), sessionContext.getIdentity().getUser(), sessionContext.getIdentity().getGroups(), Optional.ofNullable(sessionContext.getSource()), sessionContext.getClientTags(), sessionContext.getResourceEstimates(), queryType));
session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, queryType, selectionContext.getResourceGroupId());
transactionManager.activateTransaction(session, isTransactionControlStatement(preparedQuery.getStatement()), accessControl); DispatchQuery dispatchQuery = dispatchQueryFactory.createDispatchQuery( session, query, preparedQuery, slug, selectionContext.getResourceGroupId()); boolean queryAdded = queryCreated(dispatchQuery); if (queryAdded && !dispatchQuery.isDone()) { try { resourceGroupManager.submit(dispatchQuery, selectionContext, dispatchExecutor); } catch (Throwable e) { dispatchQuery.fail(e); } } } catch (Throwable throwable) { if (session == null) { session = Session.builder(new SessionPropertyManager()) .setQueryId(queryId) .setIdentity(sessionContext.getIdentity()) .setSource(sessionContext.getSource()) .build(); } Optional<String> preparedSql = Optional.ofNullable(preparedQuery).flatMap(PreparedQuery::getPrepareSql); DispatchQuery failedDispatchQuery = failedDispatchQueryFactory.createFailedDispatchQuery(session, query, preparedSql, Optional.empty(), throwable); queryCreated(failedDispatchQuery); } }
|