path: root/packages/db/src/core/integration/index.ts

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;
}