Examples with AbstractExecutionSource org.jkiss.dbeaver.model.impl.AbstractExecutionSource used on opensource projects

Example 11 with AbstractExecutionSource

use of org.jkiss.dbeaver.model.impl.AbstractExecutionSource in project dbeaver by dbeaver.

the class DatabaseTransferProducer method transferData.

@Override
public void transferData(DBRProgressMonitor monitor, IDataTransferConsumer consumer, DatabaseProducerSettings settings) throws DBException {
    String contextTask = CoreMessages.data_transfer_wizard_job_task_export;
    DBPDataSource dataSource = getSourceObject().getDataSource();
    assert (dataSource != null);
    boolean selectiveExportFromUI = settings.isSelectedColumnsOnly() || settings.isSelectedRowsOnly();
    if (dataContainer instanceof ResultSetDataContainer) {
        ((ResultSetDataContainer) dataContainer).getOptions().setExportSelectedRows(settings.isSelectedRowsOnly());
        ((ResultSetDataContainer) dataContainer).getOptions().setExportSelectedColumns(settings.isSelectedColumnsOnly());
    }
    boolean newConnection = settings.isOpenNewConnections();
    boolean forceDataReadTransactions = Boolean.TRUE.equals(dataSource.getDataSourceFeature(FEATURE_FORCE_TRANSACTIONS));
    DBCExecutionContext context = !selectiveExportFromUI && newConnection ? dataSource.openIsolatedContext(monitor, "Data transfer producer") : dataSource.getDefaultContext(false);
    try (DBCSession session = context.openSession(monitor, DBCExecutionPurpose.UTIL, contextTask)) {
        try {
            AbstractExecutionSource transferSource = new AbstractExecutionSource(dataContainer, context, consumer);
            session.enableLogging(false);
            if (!selectiveExportFromUI && (newConnection || forceDataReadTransactions)) {
                // other complex structures only in transactional mode
                try {
                    DBCTransactionManager txnManager = DBUtils.getTransactionManager(context);
                    if (txnManager != null) {
                        txnManager.setAutoCommit(monitor, false);
                    }
                } catch (DBCException e) {
                    log.warn("Can't change auto-commit", e);
                }
            }
            long totalRows = 0;
            if (settings.isQueryRowCount() && (dataContainer.getSupportedFeatures() & DBSDataContainer.DATA_COUNT) != 0) {
                monitor.beginTask(CoreMessages.data_transfer_wizard_job_task_retrieve, 1);
                try {
                    totalRows = dataContainer.countData(transferSource, session, dataFilter);
                } catch (Throwable e) {
                    log.warn("Can't retrieve row count from '" + dataContainer.getName() + "'", e);
                    try {
                        DBCTransactionManager txnManager = DBUtils.getTransactionManager(session.getExecutionContext());
                        if (txnManager != null && !txnManager.isAutoCommit()) {
                            txnManager.rollback(session, null);
                        }
                    } catch (Throwable e1) {
                        log.warn("Error rolling back transaction", e1);
                    }
                } finally {
                    monitor.done();
                }
            }
            monitor.beginTask(CoreMessages.data_transfer_wizard_job_task_export_table_data, (int) totalRows);
            try {
                // Perform export
                if (settings.getExtractType() == DatabaseProducerSettings.ExtractType.SINGLE_QUERY) {
                    // Just do it in single query
                    dataContainer.readData(transferSource, session, consumer, dataFilter, -1, -1, DBSDataContainer.FLAG_NONE);
                } else {
                    // Read all data by segments
                    long offset = 0;
                    int segmentSize = settings.getSegmentSize();
                    for (; ; ) {
                        DBCStatistics statistics = dataContainer.readData(transferSource, session, consumer, dataFilter, offset, segmentSize, DBSDataContainer.FLAG_NONE);
                        if (statistics == null || statistics.getRowsFetched() < segmentSize) {
                            // Done
                            break;
                        }
                        offset += statistics.getRowsFetched();
                    }
                }
            } finally {
                monitor.done();
            }
        } finally {
            if (!selectiveExportFromUI && (newConnection || forceDataReadTransactions)) {
                DBCTransactionManager txnManager = DBUtils.getTransactionManager(context);
                if (txnManager != null) {
                    try {
                        txnManager.commit(session);
                    } catch (DBCException e) {
                        log.error("Can't finish transaction in data producer connection", e);
                    }
                }
            }
            if (newConnection) {
                context.close();
            }
        }
    }
}

Example 12 with AbstractExecutionSource

use of org.jkiss.dbeaver.model.impl.AbstractExecutionSource in project dbeaver by dbeaver.

the class SQLQueryJob method executeStatement.

private void executeStatement(@NotNull DBCSession session, SQLQuery sqlQuery, long startTime, SQLQueryResult curResult) throws DBCException {
    DBCExecutionSource source = new AbstractExecutionSource(dataContainer, session.getExecutionContext(), partSite.getPart(), sqlQuery);
    final DBCStatement dbcStatement = DBUtils.makeStatement(source, session, DBCStatementType.SCRIPT, sqlQuery, rsOffset, rsMaxRows);
    curStatement = dbcStatement;
    int statementTimeout = getDataSourceContainer().getPreferenceStore().getInt(DBeaverPreferences.STATEMENT_TIMEOUT);
    if (statementTimeout > 0) {
        try {
            dbcStatement.setStatementTimeout(statementTimeout);
        } catch (Throwable e) {
            log.debug("Can't set statement timeout:" + e.getMessage());
        }
    }
    // Execute statement
    try {
        session.getProgressMonitor().subTask("Execute query");
        boolean hasResultSet = dbcStatement.executeStatement();
        curResult.setHasResultSet(hasResultSet);
        statistics.addExecuteTime(System.currentTimeMillis() - startTime);
        statistics.addStatementsCount();
        long updateCount = -1;
        while (hasResultSet || resultSetNumber == 0 || updateCount >= 0) {
            // Fetch data only if we have to fetch all results or if it is rs requested
            if (fetchResultSetNumber < 0 || fetchResultSetNumber == resultSetNumber) {
                if (hasResultSet && fetchResultSets) {
                    DBCResultSet resultSet = dbcStatement.openResultSet();
                    if (resultSet == null) {
                        // Kind of bug in the driver. It says it has resultset but returns null
                        break;
                    } else {
                        DBDDataReceiver dataReceiver = resultsConsumer.getDataReceiver(sqlQuery, resultSetNumber);
                        if (dataReceiver != null) {
                            hasResultSet = fetchQueryData(session, resultSet, curResult, dataReceiver, true);
                        }
                    }
                }
            }
            if (!hasResultSet) {
                try {
                    updateCount = dbcStatement.getUpdateRowCount();
                    if (updateCount >= 0) {
                        curResult.setUpdateCount(updateCount);
                        statistics.addRowsUpdated(updateCount);
                    }
                } catch (DBCException e) {
                    // In some cases we can't read update count
                    // This is bad but we can live with it
                    // Just print a warning
                    log.warn("Can't obtain update count", e);
                }
            }
            if (hasResultSet && fetchResultSets) {
                resultSetNumber++;
                fetchResultSetNumber = resultSetNumber;
            }
            if (!hasResultSet && updateCount < 0) {
                // Nothing else to fetch
                break;
            }
            if (session.getDataSource().getInfo().supportsMultipleResults()) {
                try {
                    hasResultSet = dbcStatement.nextResults();
                } catch (DBCException e) {
                    log.error(e);
                    // #2792: Check this twice. Some drivers (e.g. Sybase jConnect)
                    // throw error on n'th result fetch - but it still can keep fetching next results
                    hasResultSet = dbcStatement.nextResults();
                }
                updateCount = hasResultSet ? -1 : 0;
            } else {
                break;
            }
        }
    } finally {
        try {
            curResult.addWarnings(dbcStatement.getStatementWarnings());
        } catch (Throwable e) {
            log.warn("Can't read execution warnings", e);
        }
        // monitor.subTask("Close query");
        if (!keepStatementOpen()) {
            closeStatement();
        }
    }
}

Example 13 with AbstractExecutionSource

use of org.jkiss.dbeaver.model.impl.AbstractExecutionSource in project dbeaver by serge-rider.

the class DatabaseTransferConsumer method fetchStart.

@Override
public void fetchStart(DBCSession session, DBCResultSet resultSet, long offset, long maxRows) throws DBCException {
    try {
        initExporter(session.getProgressMonitor());
    } catch (DBException e) {
        throw new DBCException("Error initializing exporter", e);
    }
    if (containerMapping == null) {
        throw new DBCException("Internal error: consumer mappings not set");
    }
    AbstractExecutionSource executionSource = new AbstractExecutionSource(containerMapping.getSource(), targetContext, this);
    DBSDataManipulator targetObject = getTargetObject();
    if (!isPreview && offset <= 0 && settings.isTruncateBeforeLoad() && (containerMapping == null || containerMapping.getMappingType() == DatabaseMappingType.existing)) {
        // Truncate target tables
        if ((targetObject.getSupportedFeatures() & DBSDataManipulator.DATA_TRUNCATE) != 0) {
            targetObject.truncateData(targetSession, executionSource);
        } else {
            log.error("Table '" + targetObject.getName() + "' doesn't support truncate operation");
        }
    }
    DBDAttributeBinding[] rsAttributes;
    boolean dynamicTarget = targetContext.getDataSource().getInfo().isDynamicMetadata();
    DBSDataContainer sourceObject = getSourceObject();
    if (dynamicTarget) {
        // Document-based datasource
        rsAttributes = DBUtils.getAttributeBindings(session, sourceObject, resultSet.getMeta());
    } else {
        rsAttributes = DBUtils.makeLeafAttributeBindings(session, sourceObject, resultSet);
    }
    columnMappings = new ColumnMapping[rsAttributes.length];
    sourceBindings = rsAttributes;
    targetAttributes = new ArrayList<>(columnMappings.length);
    for (int i = 0; i < rsAttributes.length; i++) {
        if (isSkipColumn(rsAttributes[i])) {
            continue;
        }
        ColumnMapping columnMapping = new ColumnMapping(rsAttributes[i]);
        if (containerMapping == null) {
            // Map all attributes directly.
            if (targetObject instanceof DBSEntity) {
                try {
                    DBSEntityAttribute attribute = ((DBSEntity) targetObject).getAttribute(session.getProgressMonitor(), columnMapping.sourceAttr.getName());
                    if (attribute != null) {
                        columnMapping.targetAttr = new DatabaseMappingAttribute(null, columnMapping.sourceAttr);
                        columnMapping.targetAttr.setTarget(attribute);
                        columnMapping.targetAttr.setMappingType(DatabaseMappingType.existing);
                    }
                } catch (DBException e) {
                    log.error("Error getting target attribute");
                }
            }
            if (columnMapping.targetAttr == null) {
                throw new DBCException("Can't resolve target attribute for [" + columnMapping.sourceAttr.getName() + "]");
            }
        } else if (sourceObject instanceof DBSDocumentContainer && dynamicTarget) {
            try {
                DBSDocumentContainer docContainer = (DBSDocumentContainer) (targetObject instanceof DBSDocumentContainer ? targetObject : sourceObject);
                DBSEntityAttribute docAttribute = docContainer.getDocumentAttribute(session.getProgressMonitor());
                if (docAttribute != null) {
                    columnMapping.targetAttr = new DatabaseMappingAttribute(containerMapping, columnMapping.sourceAttr);
                    columnMapping.targetAttr.setTarget(docAttribute);
                    columnMapping.targetAttr.setMappingType(DatabaseMappingType.existing);
                }
            } catch (DBException e) {
                throw new DBCException("Error getting document attribute", e);
            }
        } else {
            columnMapping.targetAttr = containerMapping.getAttributeMapping(columnMapping.sourceAttr);
            if (columnMapping.targetAttr == null) {
                throw new DBCException("Can't find target attribute [" + columnMapping.sourceAttr.getName() + "]");
            }
        }
        if (columnMapping.targetAttr.getMappingType() == DatabaseMappingType.skip) {
            continue;
        }
        DBSEntityAttribute targetAttr = columnMapping.targetAttr.getTarget();
        if (targetAttr == null) {
            if (isPreview) {
                targetAttr = new PreviewColumnInfo(null, columnMapping.sourceAttr, columnMapping.targetIndex);
            } else if (columnMapping.targetAttr.getSource() instanceof DBSEntityAttribute) {
                // Use source attr. Some datasource (e.g. document oriented do not have strict set of attributes)
                targetAttr = (DBSEntityAttribute) columnMapping.targetAttr.getSource();
            } else {
                throw new DBCException("Target attribute for [" + columnMapping.sourceAttr.getName() + "] wasn't resolved");
            }
        }
        columnMapping.sourceValueHandler = columnMapping.sourceAttr.getValueHandler();
        columnMapping.targetValueHandler = DBUtils.findValueHandler(targetContext.getDataSource(), targetAttr);
        columnMapping.targetIndex = targetAttributes.size();
        columnMappings[i] = columnMapping;
        targetAttributes.add(targetAttr);
    }
    DBSAttributeBase[] attributes = targetAttributes.toArray(new DBSAttributeBase[0]);
    if (!isPreview) {
        if (targetObject instanceof DBSDataManipulatorExt) {
            ((DBSDataManipulatorExt) targetObject).beforeDataChange(targetSession, DBSManipulationType.INSERT, attributes, executionSource);
        }
        executeBatch = targetObject.insertData(targetSession, attributes, null, executionSource);
    } else {
        previewRows = new ArrayList<>();
        executeBatch = new PreviewBatch();
    }
}

Example 14 with AbstractExecutionSource

use of org.jkiss.dbeaver.model.impl.AbstractExecutionSource in project dbeaver by serge-rider.

the class ResultSetViewer method readRowCount.

/**
 * Reads row count and sets value in status label
 */
private long readRowCount(DBRProgressMonitor monitor) throws DBException {
    final DBCExecutionContext executionContext = getExecutionContext();
    DBSDataContainer dataContainer = getDataContainer();
    if (executionContext == null || dataContainer == null) {
        throw new DBException(ModelMessages.error_not_connected_to_database);
    }
    long[] result = new long[1];
    DBExecUtils.tryExecuteRecover(monitor, executionContext.getDataSource(), param -> {
        try (DBCSession session = executionContext.openSession(monitor, DBCExecutionPurpose.USER, "Read total row count")) {
            long rowCount = dataContainer.countData(new AbstractExecutionSource(dataContainer, executionContext, this), session, model.getDataFilter(), DBSDataContainer.FLAG_NONE);
            model.setTotalRowCount(rowCount);
            result[0] = rowCount;
        } catch (DBCException e) {
            throw new InvocationTargetException(e);
        }
    });
    return result[0];
}

Example 15 with AbstractExecutionSource

use of org.jkiss.dbeaver.model.impl.AbstractExecutionSource in project dbeaver by serge-rider.

the class SQLScriptProcessor method executeStatement.

private void executeStatement(@NotNull DBCSession session, SQLQuery sqlQuery, long startTime) throws DBCException {
    SQLQueryDataContainer dataContainer = new SQLQueryDataContainer(() -> executionContext, sqlQuery, scriptContext, log);
    DBCExecutionSource source = new AbstractExecutionSource(dataContainer, session.getExecutionContext(), this, sqlQuery);
    final DBCStatement statement = DBUtils.makeStatement(source, session, DBCStatementType.SCRIPT, sqlQuery, 0, 0);
    DBExecUtils.setStatementFetchSize(statement, 0, 0, fetchSize);
    // Execute statement
    try {
        DBRProgressMonitor monitor = session.getProgressMonitor();
        log.debug(STAT_LOG_PREFIX + "Execute query\n" + sqlQuery.getText());
        boolean hasResultSet = statement.executeStatement();
        statistics.addExecuteTime(System.currentTimeMillis() - startTime);
        statistics.addStatementsCount();
        long updateCount = -1;
        while (true) {
            // Fetch data only if we have to fetch all results or if it is rs requested
            {
                if (hasResultSet) {
                    DBCResultSet resultSet = statement.openResultSet();
                    if (resultSet == null) {
                        // Kind of bug in the driver. It says it has resultset but returns null
                        break;
                    } else {
                        hasResultSet = fetchQueryData(session, resultSet, dataReceiver);
                    }
                }
            }
            if (!hasResultSet) {
                try {
                    updateCount = statement.getUpdateRowCount();
                    if (updateCount >= 0) {
                        statistics.addRowsUpdated(updateCount);
                    }
                } catch (DBCException e) {
                    // In some cases we can't read update count
                    // This is bad but we can live with it
                    // Just print a warning
                    log.warn("Can't obtain update count", e);
                }
            }
            if (!hasResultSet && updateCount < 0) {
                // Nothing else to fetch
                break;
            }
            if (session.getDataSource().getInfo().supportsMultipleResults()) {
                try {
                    hasResultSet = statement.nextResults();
                } catch (DBCException e) {
                    if (session.getDataSource().getInfo().isMultipleResultsFetchBroken()) {
                        log.error(e);
                        // #2792: Check this twice. Some drivers (e.g. Sybase jConnect)
                        // throw error on n'th result fetch - but it still can keep fetching next results
                        hasResultSet = statement.nextResults();
                    } else {
                        throw e;
                    }
                }
                updateCount = hasResultSet ? -1 : 0;
            } else {
                break;
            }
        }
    } finally {
        try {
            Throwable[] warnings = statement.getStatementWarnings();
            if (warnings != null) {
                for (Throwable warning : warnings) {
                    scriptContext.getOutputWriter().println(warning.getMessage());
                }
            }
        } catch (Throwable e) {
            log.warn("Can't read execution warnings", e);
        }
        try {
            statement.close();
        } catch (Throwable e) {
            log.error("Error closing statement", e);
        }
        log.debug(STAT_LOG_PREFIX + "Time: " + RuntimeUtils.formatExecutionTime(statistics.getExecuteTime()) + (statistics.getRowsFetched() >= 0 ? ", fetched " + statistics.getRowsFetched() + " row(s)" : "") + (statistics.getRowsUpdated() >= 0 ? ", updated " + statistics.getRowsUpdated() + " row(s)" : ""));
    }
}