mirror of
https://github.com/DeNNiiInc/dbgate.git
synced 2026-04-29 12:14:00 +00:00
feat: use straming api for duckdb driver's readQuery
This commit is contained in:
@@ -24,6 +24,20 @@ function getDuckDb() {
|
|||||||
return duckDb;
|
return duckDb;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function getReturningStatementTypes() {
|
||||||
|
const duckdb = getDuckDb();
|
||||||
|
|
||||||
|
const returningStatementTypes = [
|
||||||
|
duckdb.StatementType.SELECT,
|
||||||
|
duckdb.StatementType.EXPLAIN,
|
||||||
|
duckdb.StatementType.EXECUTE,
|
||||||
|
duckdb.StatementType.RELATION,
|
||||||
|
duckdb.StatementType.LOGICAL_PLAN,
|
||||||
|
];
|
||||||
|
|
||||||
|
return returningStatementTypes;
|
||||||
|
}
|
||||||
|
|
||||||
/** @type {import('dbgate-types').EngineDriver<import('@duckdb/node-api').DuckDBConnection>} */
|
/** @type {import('dbgate-types').EngineDriver<import('@duckdb/node-api').DuckDBConnection>} */
|
||||||
const driver = {
|
const driver = {
|
||||||
...driverBase,
|
...driverBase,
|
||||||
@@ -60,19 +74,12 @@ const driver = {
|
|||||||
|
|
||||||
try {
|
try {
|
||||||
const statements = await dbhan.client.extractStatements(sql);
|
const statements = await dbhan.client.extractStatements(sql);
|
||||||
|
const returningStatementTypes = getReturningStatementTypes();
|
||||||
const count = statements.count;
|
const count = statements.count;
|
||||||
|
|
||||||
for (let i = 0; i < count; i++) {
|
for (let i = 0; i < count; i++) {
|
||||||
const stmt = await statements.prepare(i);
|
const stmt = await statements.prepare(i);
|
||||||
|
|
||||||
const returningStatementTypes = [
|
|
||||||
duckdb.StatementType.SELECT,
|
|
||||||
duckdb.StatementType.EXPLAIN,
|
|
||||||
duckdb.StatementType.EXECUTE,
|
|
||||||
duckdb.StatementType.RELATION,
|
|
||||||
duckdb.StatementType.LOGICAL_PLAN,
|
|
||||||
];
|
|
||||||
|
|
||||||
const result = await stmt.stream();
|
const result = await stmt.stream();
|
||||||
let hasSentColumns = false;
|
let hasSentColumns = false;
|
||||||
|
|
||||||
@@ -149,33 +156,71 @@ const driver = {
|
|||||||
highWaterMark: 100,
|
highWaterMark: 100,
|
||||||
});
|
});
|
||||||
|
|
||||||
const res = await dbhan.client.runAndReadAll(sql);
|
try {
|
||||||
const rowsObjects = res.getRowObjects();
|
const statements = await dbhan.client.extractStatements(sql);
|
||||||
|
const returningStatementTypes = getReturningStatementTypes();
|
||||||
|
const count = statements.count;
|
||||||
|
|
||||||
const columnNames = res.columnNames();
|
for (let i = 0; i < count; i++) {
|
||||||
const columnTypes = res.columnTypes();
|
const stmt = await statements.prepare(i);
|
||||||
|
|
||||||
const columns = getColumnsInfo(columnNames, columnTypes).map(normalizeRow);
|
if (!returningStatementTypes.includes(stmt.statementType)) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
const rows = rowsObjects.map(normalizeRow);
|
const result = await stmt.stream();
|
||||||
|
let hasSentHeader = false;
|
||||||
|
|
||||||
pass.write({
|
while (true) {
|
||||||
__isStreamHeader: true,
|
const chunk = await result.fetchChunk();
|
||||||
...(structure || {
|
if (!chunk || chunk.rowCount === 0) {
|
||||||
columns: columns.map((col) => ({
|
break;
|
||||||
columnName: col.name,
|
}
|
||||||
dataType: col.type,
|
|
||||||
})),
|
|
||||||
}),
|
|
||||||
});
|
|
||||||
|
|
||||||
for (const row of rows) {
|
if (!hasSentHeader) {
|
||||||
pass.write(row);
|
const columnNames = result.columnNames();
|
||||||
|
const columnTypes = result.columnTypes();
|
||||||
|
const columns = getColumnsInfo(columnNames, columnTypes);
|
||||||
|
|
||||||
|
pass.write({
|
||||||
|
__isStreamHeader: true,
|
||||||
|
...(structure || {
|
||||||
|
columns: columns.map((col) => ({
|
||||||
|
columnName: col.columnName,
|
||||||
|
dataType: col.dataType,
|
||||||
|
})),
|
||||||
|
}),
|
||||||
|
});
|
||||||
|
hasSentHeader = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
const rows = chunk.getRows();
|
||||||
|
const columnNames = result.columnNames();
|
||||||
|
for (const row of rows) {
|
||||||
|
const zipped = zipObject(columnNames, row);
|
||||||
|
pass.write(normalizeRow(zipped));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pass.end();
|
||||||
|
return pass;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(extractErrorLogData(error), 'ReadQuery error');
|
||||||
|
const { message, procName } = error;
|
||||||
|
pass.write({
|
||||||
|
__isStreamInfo: true,
|
||||||
|
info: {
|
||||||
|
message,
|
||||||
|
line: 0,
|
||||||
|
procedure: procName,
|
||||||
|
time: new Date(),
|
||||||
|
severity: 'error',
|
||||||
|
},
|
||||||
|
});
|
||||||
|
pass.end();
|
||||||
|
return pass;
|
||||||
}
|
}
|
||||||
|
|
||||||
pass.end();
|
|
||||||
|
|
||||||
return pass;
|
|
||||||
},
|
},
|
||||||
async writeTable(dbhan, name, options) {
|
async writeTable(dbhan, name, options) {
|
||||||
return createBulkInsertStreamBase(this, stream, dbhan, name, options);
|
return createBulkInsertStreamBase(this, stream, dbhan, name, options);
|
||||||
|
|||||||
Reference in New Issue
Block a user