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);