Examples with QueryCanceledQueryException datawave.webservice.query.exception.QueryCanceledQueryException used on opensource projects

Example 1 with QueryCanceledQueryException

use of datawave.webservice.query.exception.QueryCanceledQueryException in project datawave by NationalSecurityAgency.

the class CachedResultsBean method load.

protected GenericResponse<String> load(@Required("queryId") String queryId, String alias, String nameBase) {
    GenericResponse<String> response = new GenericResponse<>();
    // Find out who/what called this method
    Principal p = ctx.getCallerPrincipal();
    String owner = getOwnerFromPrincipal(p);
    String userDn = getDNFromPrincipal(p);
    Collection<Collection<String>> cbAuths = new HashSet<>();
    if (p instanceof DatawavePrincipal) {
        DatawavePrincipal dp = (DatawavePrincipal) p;
        cbAuths.addAll(dp.getAuthorizations());
    } else {
        QueryException qe = new QueryException(DatawaveErrorCode.UNEXPECTED_PRINCIPAL_ERROR, MessageFormat.format("Class: {0}", p.getClass().getName()));
        response.addException(qe);
        throw new DatawaveWebApplicationException(qe, response);
    }
    AccumuloConnectionFactory.Priority priority;
    Connector connector = null;
    RunningQuery query = null;
    String tableName = "t" + nameBase;
    String viewName = "v" + nameBase;
    Connection con = null;
    PreparedStatement ps = null;
    boolean tableCreated = false;
    boolean viewCreated = false;
    CachedRunningQuery crq = null;
    Span span = null;
    boolean queryLockedException = false;
    int rowsPerBatch = cachedResultsConfiguration.getRowsPerBatch();
    try {
        // This RunningQuery may be in use. Make a copy using the defined Query.
        RunningQuery rq = null;
        QueryLogic<?> logic = null;
        Query q = null;
        BaseQueryMetric queryMetric = null;
        TInfo traceInfo = null;
        try {
            rq = getQueryById(queryId);
            // prevent duplicate calls to load with the same queryId
            if (CachedResultsBean.loadingQueries.contains(queryId)) {
                // if a different thread is using rq, we don't want to modify it in the finally block
                rq = null;
                // this is used in the inside & outside finally block to bypass cleanup that would adversely affect the loading query
                queryLockedException = true;
                throw new QueryException(DatawaveErrorCode.QUERY_LOCKED_ERROR);
            } else {
                CachedResultsBean.loadingQueries.add(queryId);
            }
            rq.setActiveCall(true);
            Query originalQuery = rq.getSettings();
            q = originalQuery.duplicate(originalQuery.getQueryName());
            q.setId(originalQuery.getId());
            q.setUncaughtExceptionHandler(new QueryUncaughtExceptionHandler());
            Thread.currentThread().setUncaughtExceptionHandler(q.getUncaughtExceptionHandler());
            queryMetric = rq.getMetric().duplicate();
            // clear page times
            queryMetric.setPageTimes(new ArrayList<>());
            // will throw IllegalArgumentException if not defined
            logic = rq.getLogic();
            // need to clone the logic here because the QueryExpirationBean will call close on
            // rq and RunningQuery.close will call close on the logic. This is causing the batch scanner to
            // be closed after 15 minutes
            logic = (QueryLogic<?>) logic.clone();
            if (rq.getTraceInfo() != null) {
                traceInfo = rq.getTraceInfo().deepCopy();
            }
        } finally {
            if (rq != null) {
                // the original query was cloned including the queryId
                // remove original query from the cache to avoid duplicate metrics
                // when it is expired by the QueryExpirationBean
                rq.setActiveCall(false);
                if (rq.getConnection() != null) {
                    connectionFactory.returnConnection(rq.getConnection());
                }
                runningQueryCache.remove(queryId);
            }
        }
        try {
            persistByQueryId(viewName, alias, owner, CachedRunningQuery.Status.LOADING, "", false);
        } catch (IOException e2) {
            PreConditionFailedQueryException e = new PreConditionFailedQueryException(DatawaveErrorCode.CACHED_RESULTS_IMPORT_ERROR, e2);
            response.addException(e);
            response.setResult("Error loading results into cache");
            throw new PreConditionFailedException(e, response);
        }
        // Get a accumulo connection
        priority = logic.getConnectionPriority();
        Map<String, String> trackingMap = connectionFactory.getTrackingMap(Thread.currentThread().getStackTrace());
        addQueryToTrackingMap(trackingMap, q);
        accumuloConnectionRequestBean.requestBegin(queryId);
        try {
            connector = connectionFactory.getConnection(priority, trackingMap);
        } finally {
            accumuloConnectionRequestBean.requestEnd(queryId);
        }
        CacheableLogic cacheableLogic;
        Transformer t = logic.getTransformer(q);
        // Audit the query. This may be duplicative if the caller called
        // QueryExecutorBean.create() or QueryExecutorBean.reset() first.
        AuditType auditType = logic.getAuditType(q);
        if (!auditType.equals(AuditType.NONE)) {
            try {
                MultivaluedMap<String, String> queryMap = new MultivaluedMapImpl<>();
                queryMap.putAll(q.toMap());
                marking.validate(queryMap);
                queryMap.putSingle(PrivateAuditConstants.COLUMN_VISIBILITY, marking.toColumnVisibilityString());
                queryMap.putSingle(PrivateAuditConstants.AUDIT_TYPE, auditType.name());
                queryMap.putSingle(PrivateAuditConstants.USER_DN, q.getUserDN());
                queryMap.putSingle(PrivateAuditConstants.LOGIC_CLASS, logic.getLogicName());
                try {
                    List<String> selectors = logic.getSelectors(q);
                    if (selectors != null && !selectors.isEmpty()) {
                        queryMap.put(PrivateAuditConstants.SELECTORS, selectors);
                    }
                } catch (Exception e) {
                    log.error(e.getMessage());
                }
                // if the user didn't set an audit id, use the query id
                if (!queryMap.containsKey(AuditParameters.AUDIT_ID)) {
                    queryMap.putSingle(AuditParameters.AUDIT_ID, q.getId().toString());
                }
                auditor.audit(queryMap);
            } catch (Exception e) {
                QueryException qe = new QueryException(DatawaveErrorCode.QUERY_AUDITING_ERROR, e);
                log.error(qe);
                response.addException(qe.getBottomQueryException());
                throw new DatawaveWebApplicationException(qe, response);
            }
        }
        if (t instanceof CacheableLogic) {
            // hold on to a reference of the query logic so we cancel it if need be.
            qlCache.add(q.getId().toString(), owner, logic, connector);
            try {
                query = new RunningQuery(null, null, logic.getConnectionPriority(), logic, q, q.getQueryAuthorizations(), p, new RunningQueryTimingImpl(queryExpirationConf, q.getPageTimeout()), executor, predictor, metricFactory);
                query.setActiveCall(true);
                // queryMetric was duplicated from the original earlier
                query.setMetric(queryMetric);
                query.setQueryMetrics(metrics);
                query.setConnection(connector);
                // Copy trace info from a clone of the original query
                query.setTraceInfo(traceInfo);
            } finally {
                qlCache.poll(q.getId().toString());
            }
            cacheableLogic = (CacheableLogic) t;
            CachedResultsBean.loadingQueryMap.put(queryId, query);
        } else {
            throw new IllegalArgumentException(logic.getLogicName() + " does not support CachedResults calls");
        }
        try {
            con = ds.getConnection();
            // Create the result table for this query
            Statement s = con.createStatement();
            String createTable = cachedResultsConfiguration.getParameters().get("CREATE_TABLE");
            createTable = createTable.replace(TABLE, tableName);
            s.execute(createTable);
            s.close();
            tableCreated = true;
            // Parse the PreparedStatement
            String insert = cachedResultsConfiguration.getParameters().get("INSERT");
            insert = insert.replace(TABLE, tableName);
            ps = con.prepareStatement(insert);
        } catch (SQLException sqle) {
            throw new QueryException(DatawaveErrorCode.CACHED_RESULTS_TABLE_CREATE_ERROR, sqle);
        }
        // Object for keeping track of which fields are placed in which
        // table columns
        // Key is fieldName, value is column number
        Map<String, Integer> fieldMap = new HashMap<>();
        // Loop over the results and put them into the database.
        ResultsPage results = null;
        // If we're tracing this query, then continue the trace for the next call.
        if (traceInfo != null) {
            span = Trace.trace(traceInfo, "cachedresults:load");
        }
        int rowsWritten = 0;
        boolean go = true;
        while (go) {
            if (query.isCanceled()) {
                throw new QueryCanceledQueryException(DatawaveErrorCode.QUERY_CANCELED);
            }
            Span nextSpan = (span == null) ? null : Trace.start("cachedresults:next");
            try {
                if (nextSpan != null)
                    nextSpan.data("pageNumber", Long.toString(query.getLastPageNumber() + 1));
                results = query.next();
            } finally {
                if (nextSpan != null)
                    nextSpan.stop();
            }
            if (results.getResults().isEmpty()) {
                go = false;
                break;
            }
            int maxLength = 0;
            for (Object o : results.getResults()) {
                List<CacheableQueryRow> cacheableQueryRowList = cacheableLogic.writeToCache(o);
                for (CacheableQueryRow cacheableQueryObject : cacheableQueryRowList) {
                    Collection<String> values = ((CacheableQueryRow) cacheableQueryObject).getColumnValues().values();
                    int maxValueLength = 0;
                    for (String s : values) {
                        if (s.length() > maxValueLength) {
                            maxValueLength = s.length();
                        }
                    }
                    boolean dataWritten = false;
                    // If a successful maxLength has been determined, then don't change it.
                    if (maxLength == 0)
                        maxLength = maxValueLength + 1;
                    else if (maxValueLength > maxLength) {
                        maxLength = maxValueLength;
                    }
                    int attempt = 0;
                    // exception;
                    SQLException loadBatchException = null;
                    while (dataWritten == false && attempt < 10) {
                        try {
                            loadBatch(ps, owner, queryId, logic.getLogicName(), fieldMap, cacheableQueryObject, maxLength);
                            dataWritten = true;
                            rowsWritten++;
                        } catch (SQLException e) {
                            loadBatchException = e;
                            String msg = e.getMessage();
                            if (msg.startsWith("Table") && msg.endsWith("doesn't exist")) {
                                throw new QueryException(DatawaveErrorCode.CACHE_TABLE_MISSING, MessageFormat.format("message: {0}", msg));
                            } else {
                                log.info("Caught other SQLException:" + msg + " writing batch with maxLength:" + maxLength);
                                maxLength = maxLength / 2;
                            }
                        }
                        attempt++;
                    }
                    if (dataWritten == false) {
                        String message = (loadBatchException == null) ? "unknown" : loadBatchException.getMessage();
                        log.error("Batch write FAILED - last exception = " + message + "record = " + cacheableQueryObject.getColumnValues().entrySet(), loadBatchException);
                    } else if (rowsWritten >= rowsPerBatch) {
                        persistBatch(ps);
                        ps.clearBatch();
                        rowsWritten = 0;
                    }
                }
            }
        }
        // commit the last batch
        if (rowsWritten > 0) {
            persistBatch(ps);
            ps.clearBatch();
            rowsWritten = 0;
        }
        // Dump the fieldMap for debugging
        if (log.isTraceEnabled()) {
            for (Entry<String, Integer> e : fieldMap.entrySet()) {
                log.trace("Field mapping: " + e.getKey() + " -> " + e.getValue());
            }
        }
        // Create the view of the table
        viewCreated = createView(tableName, viewName, con, viewCreated, fieldMap);
        // create the CachedRunningQuery and store it under the originalQueryName, but do not activate it
        crq = new CachedRunningQuery(q, logic, viewName, alias, owner, viewName, cachedResultsConfiguration.getDefaultPageSize(), queryId, fieldMap.keySet(), null, metricFactory);
        crq.setOriginalQueryId(queryId);
        crq.setTableName(tableName);
        crq.setStatus(CachedRunningQuery.Status.LOADED);
        crq.setPrincipal(ctx.getCallerPrincipal());
        persist(crq, owner);
        crq.getMetric().setLifecycle(QueryMetric.Lifecycle.INITIALIZED);
        response.setResult(viewName);
        if (fieldMap.isEmpty()) {
            throw new NoResultsQueryException("Field map is empty.", "204-4");
        } else {
            return response;
        }
    } catch (NoResultsQueryException e) {
        crq.getMetric().setLifecycle(QueryMetric.Lifecycle.DEFINED);
        try {
            persistByQueryId(viewName, alias, owner, CachedRunningQuery.Status.LOADED, "", false);
        } catch (IOException e1) {
            response.addException(new PreConditionFailedQueryException(DatawaveErrorCode.CACHED_RESULTS_IMPORT_ERROR, e1));
        }
        response.addException(e.getBottomQueryException());
        throw new NoResultsException(e, response.getResult());
    } catch (QueryCanceledQueryException | InterruptedException e) {
        log.info("Query " + queryId + " canceled on request");
        if (crq != null) {
            crq.getMetric().setLifecycle(QueryMetric.Lifecycle.CANCELLED);
        }
        try {
            persistByQueryId(viewName, alias, owner, CachedRunningQuery.Status.CANCELED, "query canceled", false);
        } catch (IOException e1) {
            response.addException(new PreConditionFailedQueryException(DatawaveErrorCode.CACHED_RESULTS_IMPORT_ERROR, e1));
        }
        QueryException qe = new QueryException(DatawaveErrorCode.QUERY_CANCELED, e);
        response.addException(qe.getBottomQueryException());
        throw new QueryCanceledException(qe, response);
    } catch (Throwable t) {
        if (crq != null && crq.getQueryLogic().getCollectQueryMetrics() == true) {
            try {
                crq.getMetric().setError(t);
                crq.getMetric().setLifecycle(QueryMetric.Lifecycle.DEFINED);
                metrics.updateMetric(crq.getMetric());
            } catch (Exception e1) {
                log.error(e1.getMessage(), e1);
            }
        }
        String statusMessage = t.getMessage();
        if (null == statusMessage) {
            statusMessage = t.getClass().getName();
        }
        try {
            persistByQueryId(viewName, alias, owner, CachedRunningQuery.Status.ERROR, statusMessage, false);
        } catch (IOException e2) {
            response.addException(new PreConditionFailedQueryException(DatawaveErrorCode.CACHED_RESULTS_IMPORT_ERROR, e2));
        }
        // don't log stack trace of parse errors and other IllegalArgumentExceptions
        if (t instanceof IllegalArgumentException || t instanceof TokenMgrError) {
            log.info(t.getMessage());
        } else {
            log.error(t.getMessage(), t);
        }
        if (con != null) {
            Statement s = null;
            try {
                s = con.createStatement();
                if (tableCreated) {
                    // Drop the result table and view for this query
                    String dropTable = cachedResultsConfiguration.getParameters().get("DROP_TABLE");
                    dropTable = dropTable.replace(TABLE, tableName);
                    s.execute(dropTable);
                }
                if (viewCreated) {
                    String dropView = cachedResultsConfiguration.getParameters().get("DROP_VIEW");
                    dropView = dropView.replace(TABLE, viewName);
                    s.execute(dropView);
                }
                s.close();
            } catch (Exception e1) {
                log.error(e1.getMessage(), e1);
                response.addException(new QueryException(DatawaveErrorCode.FAILURE_CLEANUP_ERROR, e1).getBottomQueryException());
            } finally {
                DbUtils.closeQuietly(s);
            }
        }
        if (t instanceof Error && (t instanceof TokenMgrError) == false) {
            throw (Error) t;
        }
        // default status code
        int statusCode = Response.Status.INTERNAL_SERVER_ERROR.getStatusCode();
        if (t instanceof QueryException) {
            response.addException(((QueryException) t).getBottomQueryException());
            statusCode = ((QueryException) t).getBottomQueryException().getStatusCode();
        } else {
            QueryException qe = new QueryException(statusMessage, t);
            response.addException(qe);
        }
        throw new DatawaveWebApplicationException(t, response, statusCode);
    } finally {
        DbUtils.closeQuietly(con, ps, null);
        if (queryLockedException == false) {
            CachedResultsBean.loadingQueryMap.remove(queryId);
            CachedResultsBean.loadingQueries.remove(queryId);
        }
        if (span != null) {
            span.stop();
            span = Trace.trace(query.getTraceInfo(), "query:close");
            span.data("closedAt", new Date().toString());
            // couple milliseconds just to ensure we get something saved.
            try {
                Thread.sleep(2);
            } catch (InterruptedException e) {
            // ignore
            }
            span.stop();
        // TODO: 1.8.1: no longer done?
        // Tracer.getInstance().flush();
        }
        if (null != query) {
            query.setActiveCall(false);
            try {
                query.closeConnection(connectionFactory);
            } catch (Exception e) {
                response.addException(new QueryException(DatawaveErrorCode.QUERY_CLOSE_ERROR, e).getBottomQueryException());
            }
        } else if (connector != null) {
            try {
                connectionFactory.returnConnection(connector);
            } catch (Exception e) {
                log.error(new QueryException(DatawaveErrorCode.CONNECTOR_RETURN_ERROR, e));
            }
        }
    }
}