1 files changed, 239 insertions, 0 deletions
diff --git a/packages/db/src/core/integration/index.ts b/packages/db/src/core/integration/index.ts
new file mode 100644
index 000000000..c6f58d2fd
--- /dev/null
+++ b/packages/db/src/core/integration/index.ts
@@ -0,0 +1,239 @@
+import { existsSync } from 'node:fs';
+import { mkdir, writeFile } from 'node:fs/promises';
+import { dirname } from 'node:path';
+import { fileURLToPath } from 'node:url';
+import type { ManagedAppToken } from '@astrojs/studio';
+import type { AstroIntegration } from 'astro';
+import { blue, yellow } from 'kleur/colors';
+import {
+	type HMRPayload,
+	type UserConfig,
+	type ViteDevServer,
+	createServer,
+	loadEnv,
+	mergeConfig,
+} from 'vite';
+import parseArgs from 'yargs-parser';
+import { AstroDbError, isDbError } from '../../runtime/utils.js';
+import { CONFIG_FILE_NAMES, DB_PATH, VIRTUAL_MODULE_ID } from '../consts.js';
+import { EXEC_DEFAULT_EXPORT_ERROR, EXEC_ERROR } from '../errors.js';
+import { resolveDbConfig } from '../load-file.js';
+import { SEED_DEV_FILE_NAME } from '../queries.js';
+import { type VitePlugin, getDbDirectoryUrl, getManagedRemoteToken } from '../utils.js';
+import { fileURLIntegration } from './file-url.js';
+import { getDtsContent } from './typegen.js';
+import {
+	type LateSeedFiles,
+	type LateTables,
+	type SeedHandler,
+	vitePluginDb,
+} from './vite-plugin-db.js';
+
+function astroDBIntegration(): AstroIntegration {
+	let connectToRemote = false;
+	let configFileDependencies: string[] = [];
+	let root: URL;
+	let appToken: ManagedAppToken | undefined;
+	// Used during production builds to load seed files.
+	let tempViteServer: ViteDevServer | undefined;
+
+	// Make table loading "late" to pass to plugins from `config:setup`,
+	// but load during `config:done` to wait for integrations to settle.
+	let tables: LateTables = {
+		get() {
+			throw new Error('[astro:db] INTERNAL Tables not loaded yet');
+		},
+	};
+	let seedFiles: LateSeedFiles = {
+		get() {
+			throw new Error('[astro:db] INTERNAL Seed files not loaded yet');
+		},
+	};
+	let seedHandler: SeedHandler = {
+		execute: () => {
+			throw new Error('[astro:db] INTERNAL Seed handler not loaded yet');
+		},
+		inProgress: false,
+	};
+
+	let command: 'dev' | 'build' | 'preview' | 'sync';
+	let finalBuildOutput: string;
+	return {
+		name: 'astro:db',
+		hooks: {
+			'astro:config:setup': async ({ updateConfig, config, command: _command, logger }) => {
+				command = _command;
+				root = config.root;
+
+				if (command === 'preview') return;
+
+				let dbPlugin: VitePlugin | undefined = undefined;
+				const args = parseArgs(process.argv.slice(3));
+				connectToRemote = process.env.ASTRO_INTERNAL_TEST_REMOTE || args['remote'];
+
+				if (connectToRemote) {
+					appToken = await getManagedRemoteToken();
+					dbPlugin = vitePluginDb({
+						connectToStudio: connectToRemote,
+						appToken: appToken.token,
+						tables,
+						root: config.root,
+						srcDir: config.srcDir,
+						output: config.output,
+						seedHandler,
+					});
+				} else {
+					dbPlugin = vitePluginDb({
+						connectToStudio: false,
+						tables,
+						seedFiles,
+						root: config.root,
+						srcDir: config.srcDir,
+						output: config.output,
+						logger,
+						seedHandler,
+					});
+				}
+
+				updateConfig({
+					vite: {
+						assetsInclude: [DB_PATH],
+						plugins: [dbPlugin],
+					},
+				});
+			},
+			'astro:config:done': async ({ config, injectTypes, buildOutput }) => {
+				if (command === 'preview') return;
+
+				finalBuildOutput = buildOutput;
+
+				// TODO: refine where we load tables
+				// @matthewp: may want to load tables by path at runtime
+				const { dbConfig, dependencies, integrationSeedPaths } = await resolveDbConfig(config);
+				tables.get = () => dbConfig.tables;
+				seedFiles.get = () => integrationSeedPaths;
+				configFileDependencies = dependencies;
+
+				const localDbUrl = new URL(DB_PATH, config.root);
+				if (!connectToRemote && !existsSync(localDbUrl)) {
+					await mkdir(dirname(fileURLToPath(localDbUrl)), { recursive: true });
+					await writeFile(localDbUrl, '');
+				}
+
+				injectTypes({
+					filename: 'db.d.ts',
+					content: getDtsContent(tables.get() ?? {}),
+				});
+			},
+			'astro:server:setup': async ({ server, logger }) => {
+				seedHandler.execute = async (fileUrl) => {
+					await executeSeedFile({ fileUrl, viteServer: server });
+				};
+				const filesToWatch = [
+					...CONFIG_FILE_NAMES.map((c) => new URL(c, getDbDirectoryUrl(root))),
+					...configFileDependencies.map((c) => new URL(c, root)),
+				];
+				server.watcher.on('all', (_event, relativeEntry) => {
+					const entry = new URL(relativeEntry, root);
+					if (filesToWatch.some((f) => entry.href === f.href)) {
+						server.restart();
+					}
+				});
+				// Wait for dev server log before showing "connected".
+				setTimeout(() => {
+					logger.info(
+						connectToRemote ? 'Connected to remote database.' : 'New local database created.',
+					);
+					if (connectToRemote) return;
+
+					const localSeedPaths = SEED_DEV_FILE_NAME.map(
+						(name) => new URL(name, getDbDirectoryUrl(root)),
+					);
+					// Eager load astro:db module on startup
+					if (seedFiles.get().length || localSeedPaths.find((path) => existsSync(path))) {
+						server.ssrLoadModule(VIRTUAL_MODULE_ID).catch((e) => {
+							logger.error(e instanceof Error ? e.message : String(e));
+						});
+					}
+				}, 100);
+			},
+			'astro:build:start': async ({ logger }) => {
+				if (!connectToRemote && !databaseFileEnvDefined() && finalBuildOutput === 'server') {
+					const message = `Attempting to build without the --remote flag or the ASTRO_DATABASE_FILE environment variable defined. You probably want to pass --remote to astro build.`;
+					const hint =
+						'Learn more connecting to Studio: https://docs.astro.build/en/guides/astro-db/#connect-to-astro-studio';
+					throw new AstroDbError(message, hint);
+				}
+
+				logger.info('database: ' + (connectToRemote ? yellow('remote') : blue('local database.')));
+			},
+			'astro:build:setup': async ({ vite }) => {
+				tempViteServer = await getTempViteServer({ viteConfig: vite });
+				seedHandler.execute = async (fileUrl) => {
+					await executeSeedFile({ fileUrl, viteServer: tempViteServer! });
+				};
+			},
+			'astro:build:done': async ({}) => {
+				await appToken?.destroy();
+				await tempViteServer?.close();
+			},
+		},
+	};
+}
+
+function databaseFileEnvDefined() {
+	const env = loadEnv('', process.cwd());
+	return env.ASTRO_DATABASE_FILE != null || process.env.ASTRO_DATABASE_FILE != null;
+}
+
+export function integration(): AstroIntegration[] {
+	return [astroDBIntegration(), fileURLIntegration()];
+}
+
+async function executeSeedFile({
+	fileUrl,
+	viteServer,
+}: {
+	fileUrl: URL;
+	viteServer: ViteDevServer;
+}) {
+	// Use decodeURIComponent to handle paths with spaces correctly
+	// This ensures that %20 in the pathname is properly handled
+	const pathname = decodeURIComponent(fileUrl.pathname);
+	const mod = await viteServer.ssrLoadModule(pathname);
+	if (typeof mod.default !== 'function') {
+		throw new AstroDbError(EXEC_DEFAULT_EXPORT_ERROR(fileURLToPath(fileUrl)));
+	}
+	try {
+		await mod.default();
+	} catch (e) {
+		if (isDbError(e)) {
+			throw new AstroDbError(EXEC_ERROR(e.message));
+		}
+		throw e;
+	}
+}
+
+/**
+ * Inspired by Astro content collection config loader.
+ */
+async function getTempViteServer({ viteConfig }: { viteConfig: UserConfig }) {
+	const tempViteServer = await createServer(
+		mergeConfig(viteConfig, {
+			server: { middlewareMode: true, hmr: false, watch: null, ws: false },
+			optimizeDeps: { noDiscovery: true },
+			ssr: { external: [] },
+			logLevel: 'silent',
+		}),
+	);
+
+	const hotSend = tempViteServer.hot.send;
+	tempViteServer.hot.send = (payload: HMRPayload) => {
+		if (payload.type === 'error') {
+			throw payload.err;
+		}
+		return hotSend(payload);
+	};
+
+	return tempViteServer;
+}