diff --git a/docs/table-introspect-api.md b/docs/table-introspect-api.md index dc1028d38..073fa31c8 100644 --- a/docs/table-introspect-api.md +++ b/docs/table-introspect-api.md @@ -30,3 +30,17 @@ const table = pgTable('table', { const columns/*: { id: ..., name: ... } */ = getTableColumns(table); ``` + +```ts +import { pgTable, pickTableColumns } from 'drizzle-orm/pg-core'; + +const table = pgTable('table', { + id: integer('id').primaryKey(), + name: text('name'), +}); + +const colums = pickTableColumns(table, { + id: true, + name: false, +}); +``` \ No newline at end of file diff --git a/drizzle-orm/src/utils.ts b/drizzle-orm/src/utils.ts index 51d30e97c..ec3471a98 100644 --- a/drizzle-orm/src/utils.ts +++ b/drizzle-orm/src/utils.ts @@ -188,6 +188,16 @@ export function getTableColumns(table: T): T['_']['columns'] { return table[Table.Symbol.Columns]; } +export function pickTableColumns( + table: T, + includedColumns: Partial> +): Pick { + const allColumns = table[Table.Symbol.Columns]; + return Object.fromEntries( + Object.entries(allColumns).filter(([key]) => key in includedColumns) + ) as any; +} + export function getViewSelectedFields(view: T): T['_']['selectedFields'] { return view[ViewBaseConfig].selectedFields; } diff --git a/integration-tests/tests/utils/table-columns.test.ts b/integration-tests/tests/utils/table-columns.test.ts new file mode 100644 index 000000000..ae84dede5 --- /dev/null +++ b/integration-tests/tests/utils/table-columns.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, test } from 'vitest'; +import { integer, sqliteTable, text } from 'drizzle-orm/sqlite-core'; +import { getTableColumns, pickTableColumns } from 'drizzle-orm'; + +describe('table columns utils', () => { + const users = sqliteTable('users', { + id: integer('id').primaryKey(), + name: text('name'), + email: text('email'), + }); + + test('getTableColumns should return all columns', () => { + const columns = getTableColumns(users); + + expect(Object.keys(columns)).toEqual(['id', 'name', 'email']); + expect(columns.id.name).toBe('id'); + expect(columns.name.name).toBe('name'); + expect(columns.email.name).toBe('email'); + }); + + test('pickTableColumns should return only selected columns', () => { + const columns = pickTableColumns(users, { + id: true, + name: true + }); + + expect(Object.keys(columns)).toEqual(['id', 'name']); + expect(columns.id.name).toBe('id'); + expect(columns.name.name).toBe('name'); + expect(columns.email).toBeUndefined(); + }); + + test('pickTableColumns should return empty object if no columns selected', () => { + const columns = pickTableColumns(users, {}); + expect(Object.keys(columns)).toEqual([]); + }); +});