Improve asDrizzleTable() types

withastro · Apr 25, 2024 · e93cfe6 · e93cfe6
1 parent 25edcb7
commit e93cfe6
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/packages/db/src/runtime/index.ts b/packages/db/src/runtime/index.ts
@@ -8,9 +8,10 @@ import {
 	sqliteTable,
 	text,
 } from 'drizzle-orm/sqlite-core';
-import { type DBColumn, type DBTable } from '../core/types.js';
-import { type SerializedSQL, isSerializedSQL } from './types.js';
+import { type ColumnsConfig, type DBColumn, type DBTable, type TableConfig } from '../core/types.js';
+import { type SerializedSQL, isSerializedSQL, type Table } from './types.js';
 import { pathToFileURL } from './utils.js';
+import { tableSchema } from '../core/schemas.js';
 
 export type { Table } from './types.js';
 export { createRemoteDatabaseClient, createLocalDatabaseClient } from './db-client.js';
@@ -48,7 +49,8 @@ type D1ColumnBuilder = SQLiteColumnBuilderBase<
 	ColumnBuilderBaseConfig<ColumnDataType, string> & { data: unknown }
 >;
 
-export function asDrizzleTable(name: string, table: DBTable) {
+export function asDrizzleTable<TableName extends string = string, TColumns extends ColumnsConfig = ColumnsConfig>(name: TableName, tbl: TableConfig<TColumns>): Table<TableName, TColumns> {
+	const table = tableSchema.parse(tbl);
 	const columns: Record<string, D1ColumnBuilder> = {};
 	if (!Object.entries(table.columns).some(([, column]) => hasPrimaryKey(column))) {
 		columns['_id'] = integer('_id').primaryKey();
@@ -67,7 +69,7 @@ export function asDrizzleTable(name: string, table: DBTable) {
 		}
 		return indexes;
 	});
-	return drizzleTable;
+	return drizzleTable as Table<TableName, TColumns>;
 }
 
 function atLeastOne<T>(arr: T[]): arr is [T, ...T[]] {