feat: basic defineData implementation

packages/db/config-augment.d.ts

@@ -6,7 +6,7 @@ declare namespace Config {
 declare module 'astro:db' {
 	export const db: import('./dist/runtime/index.js').SqliteDB;
 	export const dbUrl: string;
-	export const defineData: typeof import('./dist/runtime/types.js').defineData;
+	export const defineData: typeof import('./dist/runtime/index.js').defineData;
 
 	export {
 		sql,

packages/db/src/core/consts.ts

@@ -12,3 +12,10 @@ export const DB_TYPES_FILE = 'db-types.d.ts';
 export const VIRTUAL_MODULE_ID = 'astro:db';
 
 export const DB_PATH = '.astro/content.db';
+
+export const SUPPORTED_DATA_FILES = [
+	'astro.data.ts',
+	'astro.data.mts',
+	'astro.data.js',
+	'astro.data.mjs',
+];

packages/db/src/core/integration/index.ts

@@ -16,6 +16,7 @@ import { bold } from 'kleur/colors';
 import { fileURLIntegration } from './file-url.js';
 import { setupDbTables } from '../queries.js';
 import { collectionToTable } from '../../runtime/index.js';
+import { loadDataFile } from './load-astro-data.js';
 
 function astroDBIntegration(): AstroIntegration {
 	return {
@@ -65,10 +66,11 @@ function astroDBIntegration(): AstroIntegration {
 						dbUrl: dbUrl.toString(),
 						seeding: true,
 					});
+					const dataFile = await loadDataFile({ root: config.root, collections });
 					await setupDbTables({
 						db,
 						collections,
-						data: configWithDb.db?.data,
+						data: dataFile?.default,
 						logger,
 						mode: command === 'dev' ? 'dev' : 'build',
 						useForeignKeys: true,

packages/db/src/core/integration/load-astro-data.ts

@@ -0,0 +1,119 @@
+import { build as esbuild } from 'esbuild';
+import { SUPPORTED_DATA_FILES, VIRTUAL_MODULE_ID } from '../consts.js';
+import { fileURLToPath } from 'node:url';
+import { existsSync, unlinkSync, writeFileSync } from 'node:fs';
+import { getVirtualModContents } from './vite-plugin-db.js';
+import { z } from 'zod';
+import type { DBCollections } from '../types.js';
+
+const dataFileSchema = z.object({
+	// TODO: robust type checking
+	default: z.function().returns(z.void().or(z.promise(z.void()))),
+});
+
+export async function loadDataFile({
+	root,
+	collections,
+}: {
+	root: URL;
+	collections: DBCollections;
+}) {
+	let fileUrl: URL | undefined;
+	for (const filename of SUPPORTED_DATA_FILES) {
+		const url = new URL(filename, root);
+		if (!existsSync(url)) continue;
+
+		fileUrl = url;
+		break;
+	}
+
+	if (!fileUrl) {
+		return undefined;
+	}
+
+	const { code } = await bundleDataFile({ root, fileUrl, collections });
+	return dataFileSchema.parse(await loadBundledFile({ code, root }));
+}
+
+/**
+ * Bundle config file to support `.ts` files. Simplified fork from Vite's `bundleConfigFile`
+ * function:
+ *
+ * @see https://github.com/vitejs/vite/blob/main/packages/vite/src/node/config.ts#L961
+ */
+async function bundleDataFile({
+	root,
+	fileUrl,
+	collections,
+}: {
+	root: URL;
+	fileUrl: URL;
+	collections: DBCollections;
+}): Promise<{ code: string; dependencies: string[] }> {
+	const result = await esbuild({
+		absWorkingDir: process.cwd(),
+		entryPoints: [fileURLToPath(fileUrl)],
+		outfile: 'out.js',
+		packages: 'external',
+		write: false,
+		target: ['node16'],
+		platform: 'node',
+		bundle: true,
+		format: 'esm',
+		sourcemap: 'inline',
+		metafile: true,
+		plugins: [
+			{
+				name: 'resolve-astro-db',
+				setup(build) {
+					build.onResolve({ filter: /^astro:db$/ }, ({ path }) => {
+						return { path, namespace: VIRTUAL_MODULE_ID };
+					});
+					build.onLoad({ namespace: VIRTUAL_MODULE_ID, filter: /.*/ }, () => {
+						return {
+							contents: getVirtualModContents({ root, collections, isSeed: true }),
+							// Needed to resolve `@packages/studio` internals
+							resolveDir: process.cwd(),
+						};
+					});
+				},
+			},
+		],
+	});
+
+	const file = result.outputFiles[0];
+	if (!file) {
+		throw new Error(`Unexpected: no output file`);
+	}
+
+	return {
+		code: file.text,
+		dependencies: Object.keys(result.metafile.inputs),
+	};
+}
+
+/**
+ * Forked from Vite config loader, replacing CJS-based path concat with ESM only
+ *
+ * @see https://github.com/vitejs/vite/blob/main/packages/vite/src/node/config.ts#L1074
+ */
+async function loadBundledFile({
+	code,
+	root,
+}: {
+	code: string;
+	root: URL;
+}): Promise<{ default?: unknown }> {
+	// Write it to disk, load it with native Node ESM, then delete the file.
+	const tmpFileUrl = new URL(`astro.data.timestamp-${Date.now()}.mjs`, root);
+	writeFileSync(tmpFileUrl, code);
+	try {
+		return await import(/* @vite-ignore */ tmpFileUrl.pathname);
+	} finally {
+		try {
+			unlinkSync(tmpFileUrl);
+		} catch {
+			// already removed if this function is called twice simultaneously
+		}
+	}
+}

packages/db/src/core/integration/vite-plugin-db.ts

@@ -38,18 +38,35 @@ export function vitePluginDb(
 	};
 }
 
-export function getVirtualModContents({ collections, root }: { collections: DBCollections; root: URL }) {
+export function getVirtualModContents({
+	collections,
+	root,
+	isSeed,
+}: {
+	collections: DBCollections;
+	root: URL;
+	isSeed?: boolean;
+}) {
 	const dbUrl = new URL(DB_PATH, root);
 	return `
-import { collectionToTable, createLocalDatabaseClient } from ${RUNTIME_IMPORT};
-import dbUrl from '${fileURLToPath(dbUrl)}?fileurl';
+import { collectionToTable, createLocalDatabaseClient, defineData as _defineData } from ${RUNTIME_IMPORT};
+${
+	isSeed
+		? `const dbUrl = ${JSON.stringify(dbUrl)};`
+		: `import dbUrl from '${fileURLToPath(dbUrl)}?fileurl';`
+}
 
 const params = ${JSON.stringify({
 		collections,
-		seeding: false,
+		seeding: isSeed,
 	})};
 params.dbUrl = dbUrl;
 
+export const defineData = ${
+		isSeed
+			? '_defineData;'
+			: '() => { throw new Error("defineData() should not be called at runtime.") }'
+	}; 
 export const db = await createLocalDatabaseClient(params);
 
 export * from ${RUNTIME_DRIZZLE_IMPORT};

packages/db/src/core/queries.ts

@@ -14,8 +14,9 @@ import { bold } from 'kleur/colors';
 import { type SQL, sql } from 'drizzle-orm';
 import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core';
 import type { AstroIntegrationLogger } from 'astro';
-import type { DBUserConfig } from '../core/types.js';
+import type { MaybePromise } from '../core/types.js';
 import { hasPrimaryKey } from '../runtime/index.js';
+import type { DBDataContext } from '../runtime/types.js';
 
 const sqlite = new SQLiteAsyncDialect();
 
@@ -29,7 +30,7 @@ export async function setupDbTables({
 	useForeignKeys = false,
 }: {
 	db: SqliteRemoteDatabase;
-	data?: DBUserConfig['data'];
+	data?: (ctx: DBDataContext) => MaybePromise<void>;
 	collections: DBCollections;
 	logger?: AstroIntegrationLogger;
 	mode: 'dev' | 'build';
@@ -48,7 +49,7 @@ export async function setupDbTables({
 	if (data) {
 		try {
 			await data({
-				async seed({ table }, values) {
+				async seed(table, values) {
 					const result = Array.isArray(values)
 						? // TODO: fix values typing once we can infer fields type correctly
 							await db

packages/db/src/core/types.ts

@@ -184,9 +184,7 @@ export const dbConfigSchema = z.object({
 		.optional(),
 });
 
-export type DBUserConfig = Omit<z.input<typeof dbConfigSchema>, 'data'> & {
-	data(params: DBDataContext): MaybePromise<void>;
-};
+export type DBUserConfig = z.input<typeof dbConfigSchema>;
 
 export const astroConfigWithDbSchema = z.object({
 	db: dbConfigSchema.optional(),

packages/db/src/runtime/index.ts

@@ -1,5 +1,5 @@
 import type { SqliteRemoteDatabase } from 'drizzle-orm/sqlite-proxy';
-import { type DBCollection, type DBField } from '../core/types.js';
+import { type DBCollection, type DBField, type MaybePromise } from '../core/types.js';
 import { type ColumnBuilderBaseConfig, type ColumnDataType, sql } from 'drizzle-orm';
 import {
 	customType,
@@ -11,6 +11,7 @@ import {
 	type IndexBuilder,
 } from 'drizzle-orm/sqlite-core';
 import { z } from 'zod';
+import type { DBDataContext } from './types.js';
 
 export type SqliteDB = SqliteRemoteDatabase;
 export type { Table } from './types.js';
@@ -138,3 +139,7 @@ function columnMapper(fieldName: string, field: DBField, isJsonSerializable: boo
 	if (field.unique) c = c.unique();
 	return c;
 }
+
+export function defineData(callback: (ctx: DBDataContext) => MaybePromise<void>) {
+	return callback;
+}

packages/db/src/runtime/types.ts

@@ -112,7 +112,3 @@ export type DBDataContext = {
 	): Promise<any> /** TODO: type output */;
 	mode: 'dev' | 'build';
 };
-
-export function defineData(callback: (ctx: DBDataContext) => MaybePromise<void>) {
-	return callback;
-}

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

@@ -26,57 +26,5 @@ export default defineConfig({
 	integrations: [astroDb()],
 	db: {
 		collections: { Recipe, Ingredient },
-		async data({ seed }) {
-			const pancakes = await seed(Recipe, {
-				title: 'Pancakes',
-				description: 'A delicious breakfast',
-			});
-
-			seed(Ingredient, [
-				{
-					name: 'Flour',
-					quantity: 1,
-					recipeId: pancakes.id,
-				},
-				{
-					name: 'Eggs',
-					quantity: 2,
-					recipeId: pancakes.id,
-				},
-				{
-					name: 'Milk',
-					quantity: 1,
-					recipeId: pancakes.id,
-				},
-			]);
-
-			const pizza = await seed(Recipe, {
-				title: 'Pizza',
-				description: 'A delicious dinner',
-			});
-
-			seed(Ingredient, [
-				{
-					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,
-				},
-			]);
-		},
 	},
 });

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

@@ -1,10 +1,56 @@
 /// <reference path="./src/env.d.ts" />
-import { db, defineData, Ingredient, Recipe } from 'astro:db';
+
+import { defineData, Ingredient, Recipe } from 'astro:db';
 
 export default defineData(async ({ seed }) => {
-	await seed(Recipe, {
+	const pancakes = await seed(Recipe, {
 		title: 'Pancakes',
 		description: 'A delicious breakfast',
 	});
-	await db.insert(Recipe).values({});
+
+	seed(Ingredient, [
+		{
+			name: 'Flour',
+			quantity: 1,
+			recipeId: pancakes.id,
+		},
+		{
+			name: 'Eggs',
+			quantity: 2,
+			recipeId: pancakes.id,
+		},
+		{
+			name: 'Milk',
+			quantity: 1,
+			recipeId: pancakes.id,
+		},
+	]);
+
+	const pizza = await seed(Recipe, {
+		title: 'Pizza',
+		description: 'A delicious dinner',
+	});
+
+	await seed(Ingredient, [
+		{
+			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,
+		},
+	]);
 });