refactor: rewrite to injected set() function

packages/db/src/config.ts

@@ -13,7 +13,6 @@ import {
 	type DBField,
 } from './types.js';
 import { z } from 'zod';
-import { type SqliteDB } from './internal.js';
 import type { SQLiteInsertValue } from 'drizzle-orm/sqlite-core';
 
 export const dbConfigSchema = z.object({
@@ -27,18 +26,33 @@ export const dbConfigSchema = z.object({
 		.optional(),
 });
 
-export type DBUserConfig = Omit<z.input<typeof dbConfigSchema>, 'data' | 'collections'> & {
-	collections: Record<
-		string,
-		ResolvedCollectionConfig<z.input<typeof collectionSchema>['fields'], boolean>
-	>;
-	data(): MaybePromise<void>;
+export type DBUserConfig = Omit<z.input<typeof dbConfigSchema>, 'data'> & {
+	data(params: {
+		set<TFields extends z.input<typeof collectionSchema>['fields']>(
+			collection: ResolvedCollectionConfig<TFields, boolean>,
+			data: MaybeArray<
+				SQLiteInsertValue<
+					Table<
+						string,
+						/** TODO: true type inference */ Record<Extract<keyof TFields, string>, DBField>
+					>
+				>
+			>
+		): Promise<any> /** TODO: type output */;
+	}): MaybePromise<void>;
 };
 
 export const astroConfigWithDbSchema = z.object({
 	db: dbConfigSchema.optional(),
 });
 
+type CollectionMeta = {
+	// Collection name is set later when running the data() function.
+	// Collection config is assigned to an object key,
+	// so the collection itself does not know the table name.
+	name?: string;
+};
+
 type CollectionConfig<
 	TFields extends z.input<typeof collectionSchema>['fields'],
 	Writable extends boolean,
@@ -49,6 +63,7 @@ type CollectionConfig<
 			seed?: Writable extends false
 				? never
 				: () => MaybePromise<Array<Record<keyof TFields, any> & { id?: string }>>;
+			_: CollectionMeta;
 		}
 	: {
 			fields: TFields;
@@ -56,6 +71,7 @@ type CollectionConfig<
 			data?: Writable extends true
 				? never
 				: () => MaybePromise<Array<Record<keyof TFields, any> & { id?: string }>>;
+			_: CollectionMeta;
 		};
 
 type ResolvedCollectionConfig<
@@ -63,59 +79,37 @@ type ResolvedCollectionConfig<
 	Writable extends boolean,
 > = CollectionConfig<TFields, Writable> & {
 	writable: Writable;
-	set(
-		data: MaybeArray<
-			SQLiteInsertValue<
-				Table<
-					string,
-					/** TODO: true type inference */ Record<Extract<keyof TFields, string>, DBField>
-				>
-			>
-		>
-	): Promise<any> /** TODO: type output */;
 };
 
 export function defineCollection<TFields extends z.input<typeof collectionSchema>['fields']>(
 	userConfig: CollectionConfig<TFields, false>
 ): ResolvedCollectionConfig<TFields, false> {
-	let db: SqliteDB | undefined;
-	let table:
-		| Table<
-				string,
-				/** TODO: true type inference */ Record<Extract<keyof TFields, string>, DBField>
-		  >
-		| undefined;
-	function _setEnv(env: { db: SqliteDB; table: Table<string, TFields> }) {
-		db = env.db;
-		table = env.table;
+	const _ = {};
+	function _setMeta(values: CollectionMeta) {
+		Object.assign(_, values);
 	}
 	return {
 		...userConfig,
 		writable: false,
-		// @ts-expect-error keep private
-		_setEnv,
-		set: async (values) => {
-			if (!db || !table) {
-				throw new Error('Collection `.set()` can only be called during `data()` seeding.');
-			}
-
-			const result = Array.isArray(values)
-				? await db.insert(table).values(values).returning()
-				: await db.insert(table).values(values).returning().get();
-			return result;
-		},
+		_,
+		// @ts-expect-error private field
+		_setMeta,
 	};
 }
 
 export function defineWritableCollection<
 	TFields extends z.input<typeof collectionSchema>['fields'],
 >(userConfig: CollectionConfig<TFields, true>): ResolvedCollectionConfig<TFields, true> {
+	const _ = {};
+	function _setMeta(values: CollectionMeta) {
+		Object.assign(_, values);
+	}
 	return {
 		...userConfig,
 		writable: true,
-		set: () => {
-			throw new Error('TODO: implement for writable');
-		},
+		_,
+		// @ts-expect-error private field
+		_setMeta,
 	};
 }
 

packages/db/src/internal.ts

@@ -1,7 +1,6 @@
 import type { SqliteRemoteDatabase } from 'drizzle-orm/sqlite-proxy';
 import { createClient } from '@libsql/client';
 import {
-	type ReadableDBCollection,
 	type BooleanField,
 	type DBCollection,
 	type DBCollections,
@@ -11,7 +10,7 @@ import {
 	type JsonField,
 	type NumberField,
 	type TextField,
-	type MaybePromise,
+	collectionSchema,
 } from './types.js';
 import { type LibSQLDatabase, drizzle } from 'drizzle-orm/libsql';
 import { bold } from 'kleur/colors';
@@ -33,10 +32,11 @@ import {
 } from 'drizzle-orm/sqlite-core';
 import { z } from 'zod';
 import type { AstroIntegrationLogger } from 'astro';
-export { createRemoteDatabaseClient, findLocalDatabase } from './utils-runtime.js';
+import type { DBUserConfig } from './config.js';
 
 export type SqliteDB = SqliteRemoteDatabase;
 export type { Table } from './types.js';
+export { createRemoteDatabaseClient, findLocalDatabase } from './utils-runtime.js';
 
 const sqlite = new SQLiteAsyncDialect();
 
@@ -92,7 +92,7 @@ export async function setupDbTables({
 	mode,
 }: {
 	db: LibSQLDatabase;
-	data?: () => MaybePromise<void>;
+	data?: DBUserConfig['data'];
 	collections: DBCollections;
 	logger: AstroIntegrationLogger;
 	mode: 'dev' | 'build';
@@ -108,10 +108,24 @@ export async function setupDbTables({
 	}
 	if (data) {
 		for (const [name, collection] of Object.entries(collections)) {
-			(collection as any)._setEnv({ db, table: collectionToTable(name, collection) });
+			(collection as any)._setMeta?.({ name });
 		}
 		try {
-			await data();
+			await data({
+				async set(collection, values) {
+					const collectionName = collection._.name;
+					if (!collectionName) {
+						throw new Error(
+							`Failed to set collection data. Did you call set() outside of the data() function?`
+						);
+					}
+					const table = collectionToTable(collectionName, collectionSchema.parse(collection));
+					const result = Array.isArray(values)
+						? await db.insert(table).values(values).returning()
+						: await db.insert(table).values(values).returning().get();
+					return result;
+				},
+			});
 		} catch (e) {
 			logger.error(
 				`Failed to seed data. Did you update to match recent schema changes? Full error:\n\n${e}`

packages/db/src/types.ts

@@ -54,16 +54,16 @@ const fieldsSchema = z.record(fieldSchema);
 
 export const readableCollectionSchema = z.object({
 	fields: fieldsSchema,
-	set: z.function(),
-	_setEnv: z.function(),
 	writable: z.literal(false),
+	_: z.object({ name: z.string().optional() }).optional(),
+	_setMeta: z.function().optional(),
 });
 
 export const writableCollectionSchema = z.object({
 	fields: fieldsSchema,
-	set: z.function(),
-	_setEnv: z.function(),
 	writable: z.literal(true),
+	_: z.object({ name: z.string().optional() }).optional(),
+	_setMeta: z.function().optional(),
 });
 
 export const collectionSchema = z.union([readableCollectionSchema, writableCollectionSchema]);

packages/db/test/fixtures/recipes/astro.config.ts

@@ -22,13 +22,13 @@ export default defineConfig({
 	integrations: [db()],
 	db: {
 		collections: { Recipe, Ingredient },
-		async data() {
-			const pancakes = await Recipe.set({
+		async data({ set }) {
+			const pancakes = await set(Recipe, {
 				title: 'Pancakes',
 				description: 'A delicious breakfast',
 			});
 
-			Ingredient.set([
+			set(Ingredient, [
 				{
 					name: 'Flour',
 					quantity: 1,
@@ -46,12 +46,12 @@ export default defineConfig({
 				},
 			]);
 
-			const pizza = await Recipe.set({
+			const pizza = await set(Recipe, {
 				title: 'Pizza',
 				description: 'A delicious dinner',
 			});
 
-			Ingredient.set([
+			set(Ingredient, [
 				{
 					name: 'Flour',
 					quantity: 1,