Pick Drizzle ORM columns from a table with type safety

drizzle-pick-table-columns.ts

import { getTableColumns, type Table } from 'drizzle-orm';
import { integer, sqliteTable } from 'drizzle-orm/sqlite-core';
import { db } from './client.js';

export const pickTableColumns = <
  T extends Table, //
  const ColumnName extends keyof T['_']['columns']
>(
  table: T,
  columnNames: ColumnName[]
) => {
  type Picked = Pick<T['_']['columns'], ColumnName>;
  const picked: Partial<Picked> = {};
  const columns = getTableColumns(table);
  for (const columnName of columnNames) {
    picked[columnName] = columns[columnName];
  }
  return picked as Picked;
};

const table = sqliteTable('name', {
  id: integer('id').primaryKey(),
  n1: integer('n1').notNull(),
  n2: integer('n2')
});

// { n1: SQLiteColumn<...> }
const picked = pickTableColumns(table, ['n1']);
console.log(Object.keys(picked)); // [ 'n1' ]

// { n1: number }[];
await db.select(picked).from(table);

I adapted your solution for the use case of picking column names with the return type [string, ...string[]] so they can be passed to zod for example:

function pickTableColumnNames<
  T extends Table,
  const ColumnNames extends readonly [keyof T['_']['columns'], ...(keyof T['_']['columns'])[]],
>(_table: T, columns: ColumnNames): ColumnNames {
  return columns;
}

// eg.
const picked = pickTableColumnNames(table, ['id', 'n1']); // type: readonly ["id", "n1"]
console.log(picked) // ["id", "n1"]

// usage in zod
const schema = z.object({
  orderBy: z.enum(picked).catch('n1'),
})


// it's equivalent to the following, in case you prefer to skip the function:
const picked2 = ['id', 'n1'] as const satisfies (keyof (typeof table)['_']['columns'])[];

@patte You could comment this improvement in this GitHub issue.