diff --git a/src/bun.js/bindings/SQLClient.cpp b/src/bun.js/bindings/SQLClient.cpp index 514af1b664..2077cb29b5 100644 --- a/src/bun.js/bindings/SQLClient.cpp +++ b/src/bun.js/bindings/SQLClient.cpp @@ -74,8 +74,28 @@ typedef struct DataCell { DataCellTag tag; DataCellValue value; uint8_t freeValue; + uint8_t _indexedColumnFlag; + uint32_t index; + + bool isIndexedColumn() const { return _indexedColumnFlag == 1; } + bool isNamedColumn() const { return _indexedColumnFlag == 0; } + bool isDuplicateColumn() const { return _indexedColumnFlag == 2; } } DataCell; +class BunStructureFlags { +public: + uint32_t flags; + + BunStructureFlags(uint32_t flags) + : flags(flags) + { + } + + bool hasIndexedColumns() const { return flags & (1 << 0); } + bool hasNamedColumns() const { return flags & (1 << 1); } + bool hasDuplicateColumns() const { return flags & (1 << 2); } +}; + static JSC::JSValue toJS(JSC::VM& vm, JSC::JSGlobalObject* globalObject, DataCell& cell) { switch (cell.tag) { @@ -230,25 +250,79 @@ static JSC::JSValue toJS(JSC::VM& vm, JSC::JSGlobalObject* globalObject, DataCel } } -static JSC::JSValue toJS(JSC::Structure* structure, DataCell* cells, unsigned count, JSC::JSGlobalObject* globalObject) +static JSC::JSValue toJS(JSC::Structure* structure, DataCell* cells, unsigned count, JSC::JSGlobalObject* globalObject, Bun::BunStructureFlags flags) { auto& vm = globalObject->vm(); auto* object = JSC::constructEmptyObject(vm, structure); auto scope = DECLARE_THROW_SCOPE(vm); - for (unsigned i = 0; i < count; i++) { - auto& cell = cells[i]; - JSValue value = toJS(vm, globalObject, cell); - RETURN_IF_EXCEPTION(scope, {}); - object->putDirectOffset(vm, i, value); + // TODO: once we have more tests for this, let's add another branch for + // "only mixed names and mixed indexed columns, no duplicates" + // then we cna remove this sort and instead do two passes. + if (flags.hasIndexedColumns() && flags.hasNamedColumns()) { + // sort the cells by if they're named or indexed, put named first. + // this is to conform to the Structure offsets from earlier. + std::sort(cells, cells + count, [](DataCell& a, DataCell& b) { + return a.isNamedColumn() && !b.isNamedColumn(); + }); } + // Fast path: named columns only, no duplicate columns + if (flags.hasNamedColumns() && !flags.hasDuplicateColumns() && !flags.hasIndexedColumns()) { + for (unsigned i = 0; i < count; i++) { + auto& cell = cells[i]; + JSValue value = toJS(vm, globalObject, cell); + RETURN_IF_EXCEPTION(scope, {}); + ASSERT(!cell.isDuplicateColumn()); + ASSERT(!cell.isIndexedColumn()); + ASSERT(cell.isNamedColumn()); + object->putDirectOffset(vm, i, value); + } + } else if (flags.hasIndexedColumns() && !flags.hasNamedColumns() && !flags.hasDuplicateColumns()) { + for (unsigned i = 0; i < count; i++) { + auto& cell = cells[i]; + JSValue value = toJS(vm, globalObject, cell); + RETURN_IF_EXCEPTION(scope, {}); + ASSERT(!cell.isDuplicateColumn()); + ASSERT(cell.isIndexedColumn()); + ASSERT(!cell.isNamedColumn()); + // cell.index can be > count + // for example: + // select 1 as "8", 2 as "2", 3 as "3" + // -> { "8": 1, "2": 2, "3": 3 } + // 8 > count + object->putDirectIndex(globalObject, cell.index, value); + } + } else { + unsigned structureOffsetIndex = 0; + // slow path: named columns with duplicate columns or indexed columns + for (unsigned i = 0; i < count; i++) { + auto& cell = cells[i]; + if (cell.isIndexedColumn()) { + JSValue value = toJS(vm, globalObject, cell); + RETURN_IF_EXCEPTION(scope, {}); + ASSERT(cell.index < count); + ASSERT(!cell.isNamedColumn()); + ASSERT(!cell.isDuplicateColumn()); + object->putDirectIndex(globalObject, cell.index, value); + } else if (cell.isNamedColumn()) { + JSValue value = toJS(vm, globalObject, cell); + RETURN_IF_EXCEPTION(scope, {}); + ASSERT(!cell.isIndexedColumn()); + ASSERT(!cell.isDuplicateColumn()); + ASSERT(cell.index < count); + object->putDirectOffset(vm, structureOffsetIndex++, value); + } else if (cell.isDuplicateColumn()) { + // skip it! + } + } + } return object; } -static JSC::JSValue toJS(JSC::JSArray* array, JSC::Structure* structure, DataCell* cells, unsigned count, JSC::JSGlobalObject* globalObject) +static JSC::JSValue toJS(JSC::JSArray* array, JSC::Structure* structure, DataCell* cells, unsigned count, JSC::JSGlobalObject* globalObject, Bun::BunStructureFlags flags) { - JSValue value = toJS(structure, cells, count, globalObject); + JSValue value = toJS(structure, cells, count, globalObject, flags); if (value.isEmpty()) return {}; @@ -268,20 +342,44 @@ static JSC::JSValue toJS(JSC::JSArray* array, JSC::Structure* structure, DataCel extern "C" EncodedJSValue JSC__constructObjectFromDataCell( JSC::JSGlobalObject* globalObject, EncodedJSValue encodedArrayValue, - EncodedJSValue encodedStructureValue, DataCell* cells, unsigned count) + EncodedJSValue encodedStructureValue, DataCell* cells, unsigned count, unsigned flags) { JSValue arrayValue = JSValue::decode(encodedArrayValue); JSValue structureValue = JSValue::decode(encodedStructureValue); auto* array = arrayValue ? jsDynamicCast(arrayValue) : nullptr; auto* structure = jsDynamicCast(structureValue); - return JSValue::encode(toJS(array, structure, cells, count, globalObject)); + return JSValue::encode(toJS(array, structure, cells, count, globalObject, Bun::BunStructureFlags(flags))); } -extern "C" EncodedJSValue JSC__createStructure(JSC::JSGlobalObject* globalObject, JSC::JSCell* owner, unsigned int inlineCapacity, BunString* names) +typedef struct ExternColumnIdentifier { + uint8_t tag; + union { + uint32_t index; + BunString name; + }; + + bool isIndexedColumn() const { return tag == 1; } + bool isNamedColumn() const { return tag == 2; } + bool isDuplicateColumn() const { return tag == 0; } +} ExternColumnIdentifier; + +extern "C" EncodedJSValue JSC__createStructure(JSC::JSGlobalObject* globalObject, JSC::JSCell* owner, unsigned int inlineCapacity, ExternColumnIdentifier* namesPtr) { auto& vm = globalObject->vm(); - Structure* structure = globalObject->structureCache().emptyObjectStructureForPrototype(globalObject, globalObject->objectPrototype(), inlineCapacity); + + PropertyNameArray propertyNames(vm, PropertyNameMode::Strings, PrivateSymbolMode::Exclude); + std::span names(namesPtr, inlineCapacity); + unsigned nonDuplicateCount = 0; + for (unsigned i = 0; i < inlineCapacity; i++) { + ExternColumnIdentifier& name = names[i]; + if (name.isNamedColumn()) { + propertyNames.add(Identifier::fromString(vm, name.name.toWTFString())); + } + nonDuplicateCount += !name.isDuplicateColumn(); + } + + Structure* structure = globalObject->structureCache().emptyObjectStructureForPrototype(globalObject, globalObject->objectPrototype(), std::min(nonDuplicateCount, JSFinalObject::maxInlineCapacity)); if (owner) { vm.writeBarrier(owner, structure); } else { @@ -289,14 +387,15 @@ extern "C" EncodedJSValue JSC__createStructure(JSC::JSGlobalObject* globalObject } ensureStillAliveHere(structure); - PropertyNameArray propertyNames(vm, PropertyNameMode::Strings, PrivateSymbolMode::Exclude); - for (unsigned i = 0; i < inlineCapacity; i++) { - propertyNames.add(Identifier::fromString(vm, names[i].toWTFString())); - } - - PropertyOffset offset = 0; - for (unsigned i = 0; i < inlineCapacity; i++) { - structure = structure->addPropertyTransition(vm, structure, propertyNames[i], 0, offset); + if (names.size() > 0) { + PropertyOffset offset = 0; + unsigned indexInPropertyNamesArray = 0; + for (unsigned i = 0; i < inlineCapacity; i++) { + ExternColumnIdentifier& name = names[i]; + if (name.isNamedColumn()) { + structure = structure->addPropertyTransition(vm, structure, propertyNames[indexInPropertyNamesArray++], 0, offset); + } + } } return JSValue::encode(structure); @@ -317,5 +416,4 @@ extern "C" void JSC__putDirectOffset(JSC::VM* vm, JSC::EncodedJSValue object, un { JSValue::decode(object).getObject()->putDirectOffset(*vm, offset, JSValue::decode(value)); } - } diff --git a/src/bun.js/bindings/bindings.zig b/src/bun.js/bindings/bindings.zig index da43353ea4..eecc319413 100644 --- a/src/bun.js/bindings/bindings.zig +++ b/src/bun.js/bindings/bindings.zig @@ -102,11 +102,31 @@ pub const JSObject = extern struct { } } - extern fn JSC__createStructure(*JSC.JSGlobalObject, *JSC.JSCell, u32, names: [*]bun.String) JSC.JSValue; + extern fn JSC__createStructure(*JSC.JSGlobalObject, *JSC.JSCell, u32, names: [*]ExternColumnIdentifier, flags: u32) JSC.JSValue; - pub fn createStructure(global: *JSGlobalObject, owner: JSC.JSValue, length: u32, names: [*]bun.String) JSValue { + pub const ExternColumnIdentifier = extern struct { + tag: u8 = 0, + value: extern union { + index: u32, + name: bun.String, + }, + + pub fn string(this: *ExternColumnIdentifier) ?*bun.String { + return switch (this.tag) { + 2 => &this.value.name, + else => null, + }; + } + + pub fn deinit(this: *ExternColumnIdentifier) void { + if (this.string()) |str| { + str.deref(); + } + } + }; + pub fn createStructure(global: *JSGlobalObject, owner: JSC.JSValue, length: u32, names: [*]ExternColumnIdentifier, flags: u32) JSValue { JSC.markBinding(@src()); - return JSC__createStructure(global, owner.asCell(), length, names); + return JSC__createStructure(global, owner.asCell(), length, names, flags); } const InitializeCallback = *const fn (ctx: *anyopaque, obj: *JSObject, global: *JSGlobalObject) callconv(.C) void; diff --git a/src/sql/postgres.zig b/src/sql/postgres.zig index 296ce7ee1c..c0f2bbef84 100644 --- a/src/sql/postgres.zig +++ b/src/sql/postgres.zig @@ -1398,7 +1398,9 @@ pub const PostgresSQLConnection = struct { pub fn onClose(this: *PostgresSQLConnection) void { var vm = this.globalObject.bunVM(); - defer vm.drainMicrotasks(); + const loop = vm.eventLoop(); + loop.enter(); + defer loop.exit(); this.fail("Connection closed", error.ConnectionClosed); } @@ -1987,6 +1989,8 @@ pub const PostgresSQLConnection = struct { value: Value, free_value: u8 = 0, + isIndexedColumn: u8 = 0, + index: u32 = 0, pub const Tag = enum(u8) { null = 0, @@ -2280,6 +2284,13 @@ pub const PostgresSQLConnection = struct { } } + pub const Flags = packed struct(u32) { + has_indexed_columns: bool = false, + has_named_columns: bool = false, + has_duplicate_columns: bool = false, + _: u29 = 0, + }; + pub const Putter = struct { list: []DataCell, fields: []const protocol.FieldDescription, @@ -2287,16 +2298,25 @@ pub const PostgresSQLConnection = struct { count: usize = 0, globalObject: *JSC.JSGlobalObject, - extern fn JSC__constructObjectFromDataCell(*JSC.JSGlobalObject, JSValue, JSValue, [*]DataCell, u32) JSValue; - pub fn toJS(this: *Putter, globalObject: *JSC.JSGlobalObject, array: JSValue, structure: JSValue) JSValue { - return JSC__constructObjectFromDataCell(globalObject, array, structure, this.list.ptr, @truncate(this.fields.len)); + extern fn JSC__constructObjectFromDataCell( + *JSC.JSGlobalObject, + JSValue, + JSValue, + [*]DataCell, + u32, + Flags, + ) JSValue; + + pub fn toJS(this: *Putter, globalObject: *JSC.JSGlobalObject, array: JSValue, structure: JSValue, flags: Flags) JSValue { + return JSC__constructObjectFromDataCell(globalObject, array, structure, this.list.ptr, @truncate(this.fields.len), flags); } pub fn put(this: *Putter, index: u32, optional_bytes: ?*Data) !bool { - const oid = this.fields[index].type_oid; + const field = &this.fields[index]; + const oid = field.type_oid; debug("index: {d}, oid: {d}", .{ index, oid }); - - this.list[index] = if (optional_bytes) |data| + const cell: *DataCell = &this.list[index]; + cell.* = if (optional_bytes) |data| try DataCell.fromBytes(this.binary, oid, data.slice(), this.globalObject) else DataCell{ @@ -2306,6 +2326,21 @@ pub const PostgresSQLConnection = struct { }, }; this.count += 1; + cell.index = switch (field.name_or_index) { + // The indexed columns can be out of order. + .index => |i| i, + + else => @intCast(index), + }; + + // TODO: when duplicate and we know the result will be an object + // and not a .values() array, we can discard the data + // immediately. + cell.isIndexedColumn = switch (field.name_or_index) { + .duplicate => 2, + .index => 1, + .name => 0, + }; return true; } }; @@ -2380,6 +2415,7 @@ pub const PostgresSQLConnection = struct { .DataRow => { const request = this.current() orelse return error.ExpectedRequest; var statement = request.statement orelse return error.ExpectedStatement; + statement.checkForDuplicateFields(); var putter = DataCell.Putter{ .list = &.{}, @@ -2413,7 +2449,7 @@ pub const PostgresSQLConnection = struct { const pending_value = PostgresSQLQuery.pendingValueGetCached(request.thisValue) orelse .zero; pending_value.ensureStillAlive(); - const result = putter.toJS(this.globalObject, pending_value, statement.structure(this.js_value, this.globalObject)); + const result = putter.toJS(this.globalObject, pending_value, statement.structure(this.js_value, this.globalObject), statement.fields_flags); if (pending_value == .zero) { PostgresSQLQuery.pendingValueSetCached(request.thisValue, this.globalObject, result); @@ -2802,11 +2838,13 @@ pub const PostgresSQLConnection = struct { pub const PostgresSQLStatement = struct { cached_structure: JSC.Strong = .{}, ref_count: u32 = 1, - fields: []const protocol.FieldDescription = &[_]protocol.FieldDescription{}, + fields: []protocol.FieldDescription = &[_]protocol.FieldDescription{}, parameters: []const int4 = &[_]int4{}, signature: Signature, status: Status = Status.parsing, error_response: protocol.ErrorResponse = .{}, + needs_duplicate_check: bool = true, + fields_flags: PostgresSQLConnection.DataCell.Flags = .{}, pub const Status = enum { parsing, @@ -2827,13 +2865,58 @@ pub const PostgresSQLStatement = struct { } } + pub fn checkForDuplicateFields(this: *PostgresSQLStatement) void { + if (!this.needs_duplicate_check) return; + this.needs_duplicate_check = false; + + var seen_numbers = std.ArrayList(u32).init(bun.default_allocator); + defer seen_numbers.deinit(); + var seen_fields = bun.StringHashMap(void).init(bun.default_allocator); + seen_fields.ensureUnusedCapacity(@intCast(this.fields.len)) catch bun.outOfMemory(); + defer seen_fields.deinit(); + + // iterate backwards + var remaining = this.fields.len; + var flags: PostgresSQLConnection.DataCell.Flags = .{}; + while (remaining > 0) { + remaining -= 1; + const field: *protocol.FieldDescription = &this.fields[remaining]; + switch (field.name_or_index) { + .name => |*name| { + const seen = seen_fields.getOrPut(name.slice()) catch unreachable; + if (seen.found_existing) { + field.name_or_index = .duplicate; + flags.has_duplicate_columns = true; + } + + flags.has_named_columns = true; + }, + .index => |index| { + if (std.mem.indexOfScalar(u32, seen_numbers.items, index) != null) { + field.name_or_index = .duplicate; + flags.has_duplicate_columns = true; + } else { + seen_numbers.append(index) catch bun.outOfMemory(); + } + + flags.has_indexed_columns = true; + }, + .duplicate => { + flags.has_duplicate_columns = true; + }, + } + } + + this.fields_flags = flags; + } + pub fn deinit(this: *PostgresSQLStatement) void { debug("PostgresSQLStatement deinit", .{}); bun.assert(this.ref_count == 0); for (this.fields) |*field| { - @constCast(field).deinit(); + field.deinit(); } bun.default_allocator.free(this.fields); bun.default_allocator.free(this.parameters); @@ -2845,21 +2928,37 @@ pub const PostgresSQLStatement = struct { pub fn structure(this: *PostgresSQLStatement, owner: JSValue, globalObject: *JSC.JSGlobalObject) JSValue { return this.cached_structure.get() orelse { - const names = bun.default_allocator.alloc(bun.String, this.fields.len) catch return .undefined; + const ids = bun.default_allocator.alloc(JSC.JSObject.ExternColumnIdentifier, this.fields.len) catch return .undefined; + this.checkForDuplicateFields(); defer { - for (names) |*name| { - name.deref(); + for (ids) |*name| { + name.deinit(); } - bun.default_allocator.free(names); + bun.default_allocator.free(ids); } - for (this.fields, names) |*field, *name| { - name.* = String.fromUTF8(field.name.slice()); + + for (this.fields, ids) |*field, *id| { + id.tag = switch (field.name_or_index) { + .name => 2, + .index => 1, + .duplicate => 0, + }; + switch (field.name_or_index) { + .name => |name| { + id.value.name = String.createUTF8(name.slice()); + }, + .index => |index| { + id.value.index = index; + }, + .duplicate => {}, + } } const structure_ = JSC.JSObject.createStructure( globalObject, owner, - @truncate(this.fields.len), - names.ptr, + @truncate(ids.len), + ids.ptr, + @bitCast(this.fields_flags), ); this.cached_structure.set(globalObject, structure_); return structure_; diff --git a/src/sql/postgres/postgres_protocol.zig b/src/sql/postgres/postgres_protocol.zig index 60eeaf9f9d..2abeff0787 100644 --- a/src/sql/postgres/postgres_protocol.zig +++ b/src/sql/postgres/postgres_protocol.zig @@ -963,8 +963,47 @@ pub const DataRow = struct { pub const BindComplete = [_]u8{'2'} ++ toBytes(Int32(4)); +pub const ColumnIdentifier = union(enum) { + name: Data, + index: u32, + duplicate: void, + + pub fn init(name: Data) !@This() { + if (switch (name.slice().len) { + 1..."4294967295".len => true, + 0 => return .{ .name = .{ .empty = {} } }, + else => false, + }) might_be_int: { + // use a u64 to avoid overflow + var int: u64 = 0; + for (name.slice()) |byte| { + int = int * 10 + switch (byte) { + '0'...'9' => @as(u64, byte - '0'), + else => break :might_be_int, + }; + } + + // JSC only supports indexed property names up to 2^32 + if (int < std.math.maxInt(u32)) + return .{ .index = @intCast(int) }; + } + + return .{ .name = .{ .owned = try name.toOwned() } }; + } + + pub fn deinit(this: *@This()) void { + switch (this.*) { + .name => |*name| name.deinit(), + else => {}, + } + } +}; pub const FieldDescription = struct { - name: Data = .{ .empty = {} }, + /// JavaScriptCore treats numeric property names differently than string property names. + /// so we do the work to figure out if the property name is a number ahead of time. + name_or_index: ColumnIdentifier = .{ + .name = .{ .empty = {} }, + }, table_oid: int4 = 0, column_index: short = 0, type_oid: int4 = 0, @@ -974,7 +1013,7 @@ pub const FieldDescription = struct { } pub fn deinit(this: *@This()) void { - this.name.deinit(); + this.name_or_index.deinit(); } pub fn decodeInternal(this: *@This(), comptime Container: type, reader: NewReader(Container)) AnyPostgresError!void { @@ -997,7 +1036,7 @@ pub const FieldDescription = struct { .table_oid = try reader.int4(), .column_index = try reader.short(), .type_oid = try reader.int4(), - .name = .{ .owned = try name.toOwned() }, + .name_or_index = try ColumnIdentifier.init(name), }; try reader.skip(2 + 4 + 2); @@ -1007,10 +1046,10 @@ pub const FieldDescription = struct { }; pub const RowDescription = struct { - fields: []const FieldDescription = &[_]FieldDescription{}, + fields: []FieldDescription = &[_]FieldDescription{}, pub fn deinit(this: *@This()) void { for (this.fields) |*field| { - @constCast(field).deinit(); + field.deinit(); } bun.default_allocator.free(this.fields); diff --git a/test/js/sql/sql.test.ts b/test/js/sql/sql.test.ts index 92fd82931b..4f16d46073 100644 --- a/test/js/sql/sql.test.ts +++ b/test/js/sql/sql.test.ts @@ -157,6 +157,62 @@ if (!isCI) { expect(error.code).toBe(`ERR_POSTGRES_LIFETIME_TIMEOUT`); }); + // Last one wins. + test("Handles duplicate string column names", async () => { + const result = await sql`select 1 as x, 2 as x, 3 as x`; + expect(result).toEqual([{ x: 3 }]); + }); + + test("Handles numeric column names", async () => { + // deliberately out of order + const result = await sql`select 1 as "1", 2 as "2", 3 as "3", 0 as "0"`; + expect(result).toEqual([{ "1": 1, "2": 2, "3": 3, "0": 0 }]); + + expect(Object.keys(result[0])).toEqual(["0", "1", "2", "3"]); + // Sanity check: ensure iterating through the properties doesn't crash. + Bun.inspect(result); + }); + + // Last one wins. + test("Handles duplicate numeric column names", async () => { + const result = await sql`select 1 as "1", 2 as "1", 3 as "1"`; + expect(result).toEqual([{ "1": 3 }]); + // Sanity check: ensure iterating through the properties doesn't crash. + Bun.inspect(result); + }); + + test("Handles mixed column names", async () => { + const result = await sql`select 1 as "1", 2 as "2", 3 as "3", 4 as x`; + expect(result).toEqual([{ "1": 1, "2": 2, "3": 3, x: 4 }]); + // Sanity check: ensure iterating through the properties doesn't crash. + Bun.inspect(result); + }); + + test("Handles mixed column names with duplicates", async () => { + const result = await sql`select 1 as "1", 2 as "2", 3 as "3", 4 as "1", 1 as x, 2 as x`; + expect(result).toEqual([{ "1": 4, "2": 2, "3": 3, x: 2 }]); + // Sanity check: ensure iterating through the properties doesn't crash. + Bun.inspect(result); + + // Named columns are inserted first, but they appear from JS as last. + expect(Object.keys(result[0])).toEqual(["1", "2", "3", "x"]); + }); + + test("Handles mixed column names with duplicates at the end", async () => { + const result = await sql`select 1 as "1", 2 as "2", 3 as "3", 4 as "1", 1 as x, 2 as x, 3 as x, 4 as "y"`; + expect(result).toEqual([{ "1": 4, "2": 2, "3": 3, x: 3, y: 4 }]); + + // Sanity check: ensure iterating through the properties doesn't crash. + Bun.inspect(result); + }); + + test("Handles mixed column names with duplicates at the start", async () => { + const result = await sql`select 1 as "1", 2 as "1", 3 as "2", 4 as "3", 1 as x, 2 as x, 3 as x`; + expect(result).toEqual([{ "1": 2, "2": 3, "3": 4, x: 3 }]); + // Sanity check: ensure iterating through the properties doesn't crash. + Bun.inspect(result); + }); + test("Uses default database without slash", async () => { const sql = postgres("postgres://localhost"); expect(sql.options.username).toBe(sql.options.database);