wip: data() -> set() concept

packages/db/src/cli/commands/push/index.ts

@@ -136,12 +136,13 @@ async function pushData({
 	});
 	const queries: Query[] = [];
 
-	for (const [name, collection] of Object.entries(config.db!.collections! as DBCollections)) {
-		if (collection.writable || !collection.data) continue;
-		const table = collectionToTable(name, collection);
-		const insert = db.insert(table).values(await collection.data());
-		queries.push(insert.toSQL());
-	}
+	// TODO: update migration seeding
+	// for (const [name, collection] of Object.entries(config.db!.collections! as DBCollections)) {
+	// 	if (collection.writable || !collection.data) continue;
+	// 	const table = collectionToTable(name, collection);
+	// 	const insert = db.insert(table).values(await collection.data());
+	// 	queries.push(insert.toSQL());
+	// }
 	const url = new URL('/db/query', getRemoteDatabaseUrl());
 	const requestBody: InStatement[] = queries.map((q) => ({
 		sql: q.sql,

packages/db/src/config.ts

@@ -8,17 +8,30 @@ import {
 	type collectionSchema,
 	collectionsSchema,
 	type MaybePromise,
+	type Table,
 } 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({
 	studio: z.boolean().optional(),
 	collections: collectionsSchema.optional(),
 	// TODO: strict types
-	data: z.function().args(z.any()).optional(),
+	data: z
+		.function()
+		.args()
+		.returns(z.union([z.void(), z.promise(z.void())]))
+		.optional(),
 });
 
-export type DBUserConfig = z.input<typeof dbConfigSchema>;
+export type DBUserConfig = Omit<z.input<typeof dbConfigSchema>, 'data' | 'collections'> & {
+	collections: Record<
+		string,
+		ResolvedCollectionConfig<z.input<typeof collectionSchema>['fields'], boolean>
+	>;
+	data(): MaybePromise<void>;
+};
 
 export const astroConfigWithDbSchema = z.object({
 	db: dbConfigSchema.optional(),
@@ -48,14 +61,36 @@ type ResolvedCollectionConfig<
 	Writable extends boolean,
 > = CollectionConfig<TFields, Writable> & {
 	writable: Writable;
+	set(data: SQLiteInsertValue<Table<string, TFields>>): Promise<any> /** TODO: type output */;
 };
+type SetData<TFields extends z.input<typeof collectionSchema>['fields']> = SQLiteInsertValue<
+	Table<string, TFields>
+>;
 
 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, TFields> | undefined;
+	function _setEnv(env: { db: SqliteDB; table: Table<string, TFields> }) {
+		db = env.db;
+		table = env.table;
+	}
 	return {
 		...userConfig,
 		writable: false,
+		// @ts-expect-error keep private
+		_setEnv,
+		set: async (values: SetData<TFields>) => {
+			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;
+		},
 	};
 }
 
@@ -65,6 +100,9 @@ export function defineWritableCollection<
 	return {
 		...userConfig,
 		writable: true,
+		set: () => {
+			throw new Error('TODO: implement for writable');
+		},
 	};
 }
 

packages/db/src/internal.ts

@@ -11,6 +11,7 @@ import {
 	type JsonField,
 	type NumberField,
 	type TextField,
+	type MaybePromise,
 } from './types.js';
 import { type LibSQLDatabase, drizzle } from 'drizzle-orm/libsql';
 import { bold } from 'kleur/colors';
@@ -43,10 +44,6 @@ export function hasPrimaryKey(field: DBField) {
 	return 'primaryKey' in field && !!field.primaryKey;
 }
 
-function isReadableCollection(collection: DBCollection): collection is ReadableDBCollection {
-	return !collection.writable;
-}
-
 function checkIfModificationIsAllowed(collections: DBCollections, Table: SQLiteTable) {
 	const tableName = getTableName(Table);
 	const collection = collections[tableName];
@@ -95,7 +92,7 @@ export async function setupDbTables({
 	mode,
 }: {
 	db: LibSQLDatabase;
-	data?: (...params: any) => any;
+	data?: () => MaybePromise<void>;
 	collections: DBCollections;
 	logger: AstroIntegrationLogger;
 	mode: 'dev' | 'build';
@@ -110,26 +107,14 @@ export async function setupDbTables({
 		await db.run(q);
 	}
 	if (data) {
-		const ormObjects = Object.fromEntries(
-			Object.entries(collections).map(([name, collection]) => {
-				const table = collectionToTable(name, collection, false);
-				return [name, table];
-			})
-		);
-		await data({ db, ...ormObjects, mode });
-	}
-	// TODO: decide on removing collection-level data
-	for (const [name, collection] of Object.entries(collections)) {
-		if (!isReadableCollection(collection) || !collection.data) continue;
-
-		const table = collectionToTable(name, collection);
+		for (const [name, collection] of Object.entries(collections)) {
+			(collection as any)._setEnv({ db, table: collectionToTable(name, collection) });
+		}
 		try {
-			await db.insert(table).values(await collection.data());
+			await data();
 		} catch (e) {
 			logger.error(
-				`Failed to seed ${bold(
-					name
-				)} data. Did you update to match recent schema changes? Full error:\n\n${e}`
+				`Failed to seed data. Did you update to match recent schema changes? Full error:\n\n${e}`
 			);
 		}
 	}

packages/db/src/types.ts

@@ -52,23 +52,17 @@ const fieldSchema = z.union([
 ]);
 const fieldsSchema = z.record(fieldSchema);
 
-const dataResponse = z.array(z.record(z.unknown()));
-
 export const readableCollectionSchema = z.object({
 	fields: fieldsSchema,
-	data: z
-		.function()
-		.returns(z.union([dataResponse, z.promise(dataResponse)]))
-		.optional(),
+	set: z.function(),
+	_setEnv: z.function(),
 	writable: z.literal(false),
 });
 
 export const writableCollectionSchema = z.object({
 	fields: fieldsSchema,
-	seed: z
-		.function()
-		.returns(z.union([dataResponse, z.promise(dataResponse)]))
-		.optional(),
+	set: z.function(),
+	_setEnv: z.function(),
 	writable: z.literal(true),
 });
 
@@ -99,7 +93,7 @@ export type DBCollection = z.infer<
 export type DBCollections = Record<string, DBCollection>;
 export type DBSnapshot = {
 	schema: Record<string, DBCollection>;
-	/** 
+	/**
 	 * Snapshot version. Breaking changes to the snapshot format increment this number.
 	 * @todo Rename to "version" once closer to release.
 	 */

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

@@ -3,7 +3,7 @@ import db, { defineCollection, field } from '@astrojs/db';
 
 const Recipe = defineCollection({
 	fields: {
-		id: field.number({ primaryKey: true }),
+		id: field.number({ primaryKey: true, optional: true }),
 		title: field.text(),
 		description: field.text(),
 	},
@@ -11,7 +11,7 @@ const Recipe = defineCollection({
 
 const Ingredient = defineCollection({
 	fields: {
-		id: field.number({ primaryKey: true }),
+		id: field.number({ primaryKey: true, optional: true }),
 		name: field.text(),
 		quantity: field.number(),
 		recipeId: field.text(),
@@ -22,25 +22,57 @@ export default defineConfig({
 	integrations: [db()],
 	db: {
 		collections: { Recipe, Ingredient },
-		async data({ db, Recipe, Ingredient }) {
-			const pancakes = await db
-				.insert(Recipe)
-				.values({ title: 'Pancakes', description: 'A delicious breakfast' })
-				.returning()
-				.get();
-			await db.insert(Ingredient).values({ name: 'Flour', quantity: 1, recipeId: pancakes.id });
-			await db.insert(Ingredient).values({ name: 'Eggs', quantity: 2, recipeId: pancakes.id });
-			await db.insert(Ingredient).values({ name: 'Milk', quantity: 1, recipeId: pancakes.id });
+		async data() {
+			const pancakes = await Recipe.set({
+				title: 'Pancakes',
+				description: 'A delicious breakfast',
+			});
 
-			const pizza = await db
-				.insert(Recipe)
-				.values({ title: 'Pizza', description: 'A delicious dinner' })
-				.returning()
-				.get();
-			await db.insert(Ingredient).values({ name: 'Flour', quantity: 1, recipeId: pizza.id });
-			await db.insert(Ingredient).values({ name: 'Eggs', quantity: 2, recipeId: pizza.id });
-			await db.insert(Ingredient).values({ name: 'Milk', quantity: 1, recipeId: pizza.id });
-			await db.insert(Ingredient).values({ name: 'Tomato Sauce', quantity: 1, recipeId: pizza.id });
+			Ingredient.set([
+				{
+					name: 'Flour',
+					quantity: 1,
+					recipeId: pancakes.id,
+				},
+				{
+					name: 'Eggs',
+					quantity: 2,
+					recipeId: pancakes.id,
+				},
+				{
+					name: 'Milk',
+					quantity: 1,
+					recipeId: pancakes.id,
+				},
+			]);
+
+			const pizza = await Recipe.set({
+				title: 'Pizza',
+				description: 'A delicious dinner',
+			});
+
+			Ingredient.set([
+				{
+					name: 'Flour',
+					quantity: 1,
+					recipeId: pizza.id,
+				},
+				{
+					name: 'Eggs',
+					quantity: 2,
+					recipeId: pizza.id,
+				},
+				{
+					name: 'Milk',
+					quantity: 1,
+					recipeId: pizza.id,
+				},
+				{
+					name: 'Tomato Sauce',
+					quantity: 1,
+					recipeId: pizza.id,
+				},
+			]);
 		},
 	},
 });