提交 c824e773 编写于 作者: S Serge Rider 提交者: GitHub

Merge pull request #10449 from dbeaver/datatransfer-task-context#10311

#10311 Datatransfer task context

Former-commit-id: 4c252254
......@@ -32,6 +32,8 @@ import org.jkiss.dbeaver.model.DBPObject;
import org.jkiss.dbeaver.model.DBUtils;
import org.jkiss.dbeaver.model.app.DBPProject;
import org.jkiss.dbeaver.model.data.DBDCollection;
import org.jkiss.dbeaver.model.exec.DBCExecutionContext;
import org.jkiss.dbeaver.model.exec.DBCExecutionContextDefaults;
import org.jkiss.dbeaver.model.exec.DBExecUtils;
import org.jkiss.dbeaver.model.impl.DataSourceContextProvider;
import org.jkiss.dbeaver.model.navigator.*;
......@@ -220,7 +222,8 @@ public class DataTransferTaskConfigurator implements DBTTaskConfigurator {
newInstanceName = dataSourceObject.getName();
newObjectName = null;
} else if (dataSourceObject instanceof DBSSchema) {
newInstanceName = DBUtils.getObjectOwnerInstance(dataSourceObject).getName();
DBSObject parentObject = dataSourceObject.getParentObject();
newInstanceName = parentObject instanceof DBSCatalog ? parentObject.getName() : null;
newObjectName = dataSourceObject.getName();
} else {
// Use default database and schema
......@@ -229,16 +232,25 @@ public class DataTransferTaskConfigurator implements DBTTaskConfigurator {
}
DataSourceContextProvider contextProvider = new DataSourceContextProvider(dataSourceObject);
DBCExecutionContext executionContext = contextProvider.getExecutionContext();
String oldInstanceName = null;
String oldObjectName = null;
if (executionContext instanceof DBCExecutionContextDefaults) {
DBCExecutionContextDefaults<?, ?> contextDefaults = ((DBCExecutionContextDefaults<?, ?>) executionContext);
DBSCatalog defaultCatalog = contextDefaults.getDefaultCatalog();
if (defaultCatalog != null) {
oldInstanceName = defaultCatalog.getName();
}
DBSSchema defaultSchema = contextDefaults.getDefaultSchema();
if (defaultSchema != null) {
oldObjectName = defaultSchema.getName();
}
}
try {
DBExecUtils.setExecutionContextDefaults(
new VoidProgressMonitor(),
dataSource,
contextProvider.getExecutionContext(),
newInstanceName,
null,
newObjectName
);
DBExecUtils.setExecutionContextDefaults(new VoidProgressMonitor(), dataSource, executionContext, newInstanceName, null, newObjectName);
} catch (DBException ex) {
log.error("Error setting context defaults", ex);
return;
......@@ -250,11 +262,20 @@ public class DataTransferTaskConfigurator implements DBTTaskConfigurator {
if (query != null) {
SQLScriptContext scriptContext = new SQLScriptContext(null, contextProvider, null, new PrintWriter(System.err, true), null);
SQLQueryDataContainer container = new SQLQueryDataContainer(contextProvider, new SQLQuery(dataSource, query), scriptContext, log);
DataTransferPipe pipe = new DataTransferPipe(new DatabaseTransferProducer(container), null);
DatabaseTransferProducer producer = new DatabaseTransferProducer(container);
producer.setDefaultCatalog(newInstanceName);
producer.setDefaultSchema(newObjectName);
DataTransferPipe pipe = new DataTransferPipe(producer, null);
addPipeToTable(pipe);
updateSettings(propertyChangeListener);
}
}
try {
DBExecUtils.setExecutionContextDefaults(new VoidProgressMonitor(), dataSource, executionContext, oldInstanceName, null, oldObjectName);
} catch (DBException ex) {
log.error("Error setting context defaults", ex);
}
}
}
});
......
......@@ -27,6 +27,7 @@ import org.jkiss.dbeaver.model.app.DBPProject;
import org.jkiss.dbeaver.model.data.DBDDataFilter;
import org.jkiss.dbeaver.model.exec.*;
import org.jkiss.dbeaver.model.impl.AbstractExecutionSource;
import org.jkiss.dbeaver.model.impl.DataSourceContextProvider;
import org.jkiss.dbeaver.model.meta.DBSerializable;
import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
import org.jkiss.dbeaver.model.runtime.DBRRunnableContext;
......@@ -67,6 +68,10 @@ public class DatabaseTransferProducer implements IDataTransferProducer<DatabaseP
private DBSDataContainer dataContainer;
@Nullable
private DBDDataFilter dataFilter;
@Nullable
private String defaultCatalog;
@Nullable
private String defaultSchema;
public DatabaseTransferProducer() {
}
......@@ -120,6 +125,24 @@ public class DatabaseTransferProducer implements IDataTransferProducer<DatabaseP
return null;
}
@Nullable
public String getDefaultCatalog() {
return defaultCatalog;
}
public void setDefaultCatalog(@Nullable String defaultCatalog) {
this.defaultCatalog = defaultCatalog;
}
@Nullable
public String getDefaultSchema() {
return defaultSchema;
}
public void setDefaultSchema(@Nullable String defaultSchema) {
this.defaultSchema = defaultSchema;
}
@Override
public void transferData(
@NotNull DBRProgressMonitor monitor1,
......@@ -162,6 +185,7 @@ public class DatabaseTransferProducer implements IDataTransferProducer<DatabaseP
}
if (!selectiveExportFromUI && newConnection) {
context = DBUtils.getObjectOwnerInstance(getDatabaseObject()).openIsolatedContext(monitor, "Data transfer producer", context);
DBExecUtils.setExecutionContextDefaults(monitor, dataSource, context, defaultCatalog, null, defaultSchema);
}
if (task != null) {
DBTaskUtils.initFromContext(monitor, task, context);
......@@ -291,20 +315,14 @@ public class DatabaseTransferProducer implements IDataTransferProducer<DatabaseP
state.put("type", "query");
SQLQueryContainer queryContainer = (SQLQueryContainer) dataContainer;
DBPDataSourceContainer dataSource = queryContainer.getDataSourceContainer();
DBCExecutionContext executionContext = queryContainer.getExecutionContext();
if (dataSource != null) {
state.put("project", dataSource.getProject().getName());
state.put("dataSource", dataSource.getId());
if (executionContext instanceof DBCExecutionContextDefaults) {
DBCExecutionContextDefaults contextDefaults = ((DBCExecutionContextDefaults) executionContext);
DBSCatalog defaultCatalog = contextDefaults.getDefaultCatalog();
if (defaultCatalog != null) {
state.put("defaultCatalog", defaultCatalog.getName());
}
DBSSchema defaultSchema = contextDefaults.getDefaultSchema();
if (defaultSchema != null) {
state.put("defaultSchema", defaultSchema.getName());
}
if (object.defaultCatalog != null) {
state.put("defaultCatalog", object.defaultCatalog);
}
if (object.defaultSchema != null) {
state.put("defaultSchema", object.defaultSchema);
}
}
SQLScriptElement query = queryContainer.getQuery();
......@@ -353,18 +371,12 @@ public class DatabaseTransferProducer implements IDataTransferProducer<DatabaseP
}
DBPDataSource dataSource = ds.getDataSource();
SQLQuery query = new SQLQuery(dataSource, queryText);
TaskContextProvider taskContextProvider = new TaskContextProvider(runnableContext, dataSource, objectContext);
DBExecUtils.setExecutionContextDefaults(
monitor,
dataSource,
taskContextProvider.getExecutionContext(),
CommonUtils.toString(state.get("defaultCatalog"), null),
null,
CommonUtils.toString(state.get("defaultSchema"), null)
);
DataSourceContextProvider taskContextProvider = new DataSourceContextProvider(dataSource);
SQLScriptContext scriptContext = new SQLScriptContext(null,
taskContextProvider, null, new PrintWriter(System.err, true), null);
scriptContext.setVariables(DBTaskUtils.getVariables(objectContext));
producer.defaultCatalog = CommonUtils.toString(state.get("defaultCatalog"), null);
producer.defaultSchema = CommonUtils.toString(state.get("defaultSchema"), null);
producer.dataContainer = new SQLQueryDataContainer(
taskContextProvider,
query,
......@@ -388,37 +400,4 @@ public class DatabaseTransferProducer implements IDataTransferProducer<DatabaseP
return producer;
}
}
public static class TaskContextProvider implements DBPContextProvider {
private final DBRRunnableContext runnableContext;
private final DBPDataSource dataSource;
private final DBTTask task;
private DBCExecutionContext executionContext;
TaskContextProvider(DBRRunnableContext runnableContext, DBPDataSource dataSource, DBTTask task) {
this.runnableContext = runnableContext;
this.dataSource = dataSource;
this.task = task;
}
@Override
public DBCExecutionContext getExecutionContext() {
if (executionContext == null) {
executionContext = DBUtils.getDefaultContext(dataSource, false);
try {
runnableContext.run(true, true, monitor -> {
try {
DBTaskUtils.initFromContext(monitor, task, executionContext);
} catch (DBException e) {
throw new InvocationTargetException(e);
}
});
} catch (Exception e) {
log.error("Error initializing context", e);
}
}
return executionContext;
}
}
}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册