diff --git a/lib/duckdb.d.ts b/lib/duckdb.d.ts index c613c097..d7696473 100644 --- a/lib/duckdb.d.ts +++ b/lib/duckdb.d.ts @@ -101,7 +101,7 @@ export class Connection { ): void; unregister_udf(name: string, callback: Callback): void; - stream(sql: any, ...args: any[]): QueryResult; + stream(sql: any, ...args: any[]): Promise; arrowIPCStream(sql: any, ...args: any[]): Promise; register_buffer(name: string, array: ArrowIterable, force: boolean, callback?: Callback): void; @@ -109,6 +109,7 @@ export class Connection { } export class QueryResult implements AsyncIterable { + columns(): ColumnInfo[]; [Symbol.asyncIterator](): AsyncIterator; } @@ -166,7 +167,7 @@ export class Database { ): void; unregister_udf(name: string, callback: Callback): void; - stream(sql: any, ...args: any[]): QueryResult; + stream(sql: any, ...args: any[]): Promise; arrowIPCStream(sql: any, ...args: any[]): Promise; serialize(done?: Callback): void; @@ -261,7 +262,7 @@ export class Statement { run(...args: [...any, Callback] | any[]): Statement; - columns(): ColumnInfo[]; + columns(): ColumnInfo[] | null; } export const ERROR: number; diff --git a/lib/duckdb.js b/lib/duckdb.js index fc760d2f..5c04ccd1 100644 --- a/lib/duckdb.js +++ b/lib/duckdb.js @@ -84,6 +84,14 @@ QueryResult.prototype.nextChunk; */ QueryResult.prototype.nextIpcBuffer; +/** + * Function to return logical types for columns + * + * @method + * @return {ColumnInfo[]} - Array of column names and types + */ +QueryResult.prototype.columns; + /** * @name asyncIterator * @memberof module:duckdb~QueryResult @@ -218,12 +226,9 @@ Connection.prototype.each = function (sql) { * @param {...*} params * @yields row chunks */ -Connection.prototype.stream = async function* (sql) { +Connection.prototype.stream = async function (sql) { const statement = new Statement(this, sql); - const queryResult = await statement.stream.apply(statement, arguments); - for await (const result of queryResult) { - yield result; - } + return statement.stream.apply(statement, arguments); } /** @@ -713,7 +718,7 @@ Statement.prototype.sql; /** * @method - * @return {ColumnInfo[]} - Array of column names and types + * @return {ColumnInfo[] | null} - Array of column names and types */ Statement.prototype.columns; diff --git a/src/duckdb_node.hpp b/src/duckdb_node.hpp index 614f32cb..aaa81937 100644 --- a/src/duckdb_node.hpp +++ b/src/duckdb_node.hpp @@ -208,6 +208,7 @@ class QueryResult : public Napi::ObjectWrap { public: Napi::Value NextChunk(const Napi::CallbackInfo &info); Napi::Value NextIpcBuffer(const Napi::CallbackInfo &info); + Napi::Value Columns(const Napi::CallbackInfo &info); duckdb::shared_ptr cschema; private: diff --git a/src/statement.cpp b/src/statement.cpp index 3d813f8a..500fa8df 100644 --- a/src/statement.cpp +++ b/src/statement.cpp @@ -636,7 +636,8 @@ Napi::FunctionReference QueryResult::Init(Napi::Env env, Napi::Object exports) { Napi::Function t = DefineClass(env, "QueryResult", {InstanceMethod("nextChunk", &QueryResult::NextChunk), - InstanceMethod("nextIpcBuffer", &QueryResult::NextIpcBuffer)}); + InstanceMethod("nextIpcBuffer", &QueryResult::NextIpcBuffer), + InstanceMethod("columns", &QueryResult::Columns)}); exports.Set("QueryResult", t); @@ -742,6 +743,26 @@ Napi::Value QueryResult::NextIpcBuffer(const Napi::CallbackInfo &info) { return deferred.Promise(); } +Napi::Value QueryResult::Columns(const Napi::CallbackInfo &info) +{ + auto env = info.Env(); + auto result = Napi::Array::New(env, this->result->ColumnCount()); + + for (duckdb::idx_t column_idx = 0; column_idx < this->result->ColumnCount(); column_idx++) + { + auto column_name = this->result->ColumnName(column_idx); + auto column_type = this->result->types[column_idx]; + + auto obj = Napi::Object::New(env); + obj.Set("name", Napi::String::New(env, column_name)); + obj.Set("type", TypeToObject(env, column_type)); + + result.Set(column_idx, obj); + } + + return result; +} + Napi::Object QueryResult::NewInstance(const Napi::Object &db) { return NodeDuckDB::GetData(db.Env())->query_result_constructor.New({db}); } diff --git a/test/columns.test.ts b/test/columns.test.ts index 55dabce2..29371cd7 100644 --- a/test/columns.test.ts +++ b/test/columns.test.ts @@ -10,7 +10,7 @@ describe('Column Types', function() { var stmt = db.prepare("SELECT * EXCLUDE(medium_enum, large_enum) FROM test_all_types()", function(err: null | Error) { if (err) throw err; - let cols = stmt.columns(); + let cols = stmt.columns() as duckdb.ColumnInfo[]; assert.equal(cols.length, 42); diff --git a/test/query_result.test.ts b/test/query_result.test.ts index 16eb231c..0dccdc29 100644 --- a/test/query_result.test.ts +++ b/test/query_result.test.ts @@ -14,7 +14,18 @@ describe('QueryResult', () => { it('streams results', async () => { let retrieved = 0; - const stream = conn.stream('SELECT * FROM range(0, ?)', total); + + const stream = await conn.stream("SELECT * FROM range(0, ?)", total); + assert.deepEqual(stream.columns(), [ + { + name: 'range', + type: { + id: 'BIGINT', + sql_type: 'BIGINT', + } + } + ]); + for await (const row of stream) { retrieved++; } diff --git a/test/typescript_decls.test.ts b/test/typescript_decls.test.ts index ffdb6125..d0fe3167 100644 --- a/test/typescript_decls.test.ts +++ b/test/typescript_decls.test.ts @@ -222,7 +222,7 @@ describe("typescript: stream and QueryResult", function () { it("streams results", async () => { let retrieved = 0; - const stream = conn.stream("SELECT * FROM range(0, ?)", total); + const stream = await conn.stream("SELECT * FROM range(0, ?)", total); for await (const row of stream) { retrieved++; }