From 5b48cc0fc8383b0659a595afd3a6ee28b28779c3 Mon Sep 17 00:00:00 2001
From: Ben Holmes <hey@bholmes.dev>
Date: Fri, 15 Mar 2024 13:58:45 -0400
Subject: [PATCH] feat(db): Run db type generation on `astro sync` (#10438)

* feat: db typegen on astro sync

* fix: avoid requiring db to be installed

* fix: make typegen optional for backwards compat

* chore: changeset

* fix: required -> optional

* fix: remove flags from sync API signature
---
 .changeset/popular-radios-grow.md           |  6 +++++
 packages/astro/src/cli/install-package.ts   |  2 ++
 packages/astro/src/core/build/index.ts      |  4 +--
 packages/astro/src/core/sync/index.ts       | 30 +++++++++++++++++----
 packages/db/src/core/integration/index.ts   |  6 ++---
 packages/db/src/core/integration/typegen.ts | 13 +++++++--
 packages/db/src/core/load-file.ts           |  5 +++-
 packages/db/src/index.ts                    |  1 +
 8 files changed, 54 insertions(+), 13 deletions(-)
 create mode 100644 .changeset/popular-radios-grow.md

diff --git a/.changeset/popular-radios-grow.md b/.changeset/popular-radios-grow.md
new file mode 100644
index 000000000000..311dbde58f9f
--- /dev/null
+++ b/.changeset/popular-radios-grow.md
@@ -0,0 +1,6 @@
+---
+"astro": patch
+"@astrojs/db": patch
+---
+
+Generate Astro DB types when running `astro sync`.
diff --git a/packages/astro/src/cli/install-package.ts b/packages/astro/src/cli/install-package.ts
index e1db88d64d88..2c5af58c244a 100644
--- a/packages/astro/src/cli/install-package.ts
+++ b/packages/astro/src/cli/install-package.ts
@@ -13,6 +13,7 @@ const require = createRequire(import.meta.url);
 
 type GetPackageOptions = {
 	skipAsk?: boolean;
+	optional?: boolean;
 	cwd?: string;
 };
 
@@ -37,6 +38,7 @@ export async function getPackage<T>(
 		const packageImport = await import(packageName);
 		return packageImport as T;
 	} catch (e) {
+		if (options.optional) return undefined;
 		logger.info(
 			'SKIP_FORMAT',
 			`To continue, Astro requires the following dependency to be installed: ${bold(packageName)}.`
diff --git a/packages/astro/src/core/build/index.ts b/packages/astro/src/core/build/index.ts
index 0ebf98edd46e..d77e69fd2726 100644
--- a/packages/astro/src/core/build/index.ts
+++ b/packages/astro/src/core/build/index.ts
@@ -144,8 +144,8 @@ class AstroBuilder {
 		);
 		await runHookConfigDone({ settings: this.settings, logger: logger });
 
-		const { syncInternal } = await import('../sync/index.js');
-		const syncRet = await syncInternal(this.settings, { logger: logger, fs });
+		const { syncContentCollections } = await import('../sync/index.js');
+		const syncRet = await syncContentCollections(this.settings, { logger: logger, fs });
 		if (syncRet !== 0) {
 			return process.exit(syncRet);
 		}
diff --git a/packages/astro/src/core/sync/index.ts b/packages/astro/src/core/sync/index.ts
index 18f854ea4faa..e693ad3c4e15 100644
--- a/packages/astro/src/core/sync/index.ts
+++ b/packages/astro/src/core/sync/index.ts
@@ -3,7 +3,7 @@ import { performance } from 'node:perf_hooks';
 import { fileURLToPath } from 'node:url';
 import { dim } from 'kleur/colors';
 import { type HMRPayload, createServer } from 'vite';
-import type { AstroInlineConfig, AstroSettings } from '../../@types/astro.js';
+import type { AstroConfig, AstroInlineConfig, AstroSettings } from '../../@types/astro.js';
 import { createContentTypesGenerator } from '../../content/index.js';
 import { globalContentConfigObserver } from '../../content/utils.js';
 import { telemetry } from '../../events/index.js';
@@ -20,6 +20,8 @@ import { AstroError, AstroErrorData, createSafeError, isAstroError } from '../er
 import type { Logger } from '../logger/core.js';
 import { formatErrorMessage } from '../messages.js';
 import { ensureProcessNodeEnv } from '../util.js';
+import { getPackage } from '../../cli/install-package.js';
+import type { Arguments } from 'yargs-parser';
 
 export type ProcessExit = 0 | 1;
 
@@ -34,6 +36,10 @@ export type SyncInternalOptions = SyncOptions & {
 	logger: Logger;
 };
 
+type DBPackage = {
+	typegen?: (args: Pick<AstroConfig, 'root' | 'integrations'>) => Promise<void>;
+};
+
 /**
  * Generates TypeScript types for all Astro modules. This sets up a `src/env.d.ts` file for type inferencing,
  * and defines the `astro:content` module for the Content Collections API.
@@ -57,8 +63,24 @@ export default async function sync(
 		command: 'build',
 	});
 
+	const timerStart = performance.now();
+	const dbPackage = await getPackage<DBPackage>(
+		'@astrojs/db',
+		logger,
+		{
+			optional: true,
+			cwd: inlineConfig.root,
+		},
+		[]
+	);
+
 	try {
-		return await syncInternal(settings, { ...options, logger });
+		await dbPackage?.typegen?.(astroConfig);
+		const exitCode = await syncContentCollections(settings, { ...options, logger });
+		if (exitCode !== 0) return exitCode;
+
+		logger.info(null, `Types generated ${dim(getTimeStat(timerStart, performance.now()))}`);
+		return 0;
 	} catch (err) {
 		const error = createSafeError(err);
 		logger.error(
@@ -83,11 +105,10 @@ export default async function sync(
  * @param {LogOptions} options.logging Logging options
  * @return {Promise<ProcessExit>}
  */
-export async function syncInternal(
+export async function syncContentCollections(
 	settings: AstroSettings,
 	{ logger, fs }: SyncInternalOptions
 ): Promise<ProcessExit> {
-	const timerStart = performance.now();
 	// Needed to load content config
 	const tempViteServer = await createServer(
 		await createVite(
@@ -150,7 +171,6 @@ export async function syncInternal(
 		await tempViteServer.close();
 	}
 
-	logger.info(null, `Types generated ${dim(getTimeStat(timerStart, performance.now()))}`);
 	await setUpEnvTs({ settings, logger, fs: fs ?? fsMod });
 
 	return 0;
diff --git a/packages/db/src/core/integration/index.ts b/packages/db/src/core/integration/index.ts
index 78a1fab6d498..2dda3b7a9515 100644
--- a/packages/db/src/core/integration/index.ts
+++ b/packages/db/src/core/integration/index.ts
@@ -2,7 +2,7 @@ import { existsSync } from 'fs';
 import { dirname } from 'path';
 import { fileURLToPath } from 'url';
 import type { AstroIntegration } from 'astro';
-import { mkdir, rm, writeFile } from 'fs/promises';
+import { mkdir, writeFile } from 'fs/promises';
 import { blue, yellow } from 'kleur/colors';
 import parseArgs from 'yargs-parser';
 import { CONFIG_FILE_NAMES, DB_PATH } from '../consts.js';
@@ -10,7 +10,7 @@ import { resolveDbConfig } from '../load-file.js';
 import { type ManagedAppToken, getManagedAppTokenOrExit } from '../tokens.js';
 import { type VitePlugin, getDbDirectoryUrl } from '../utils.js';
 import { fileURLIntegration } from './file-url.js';
-import { typegen } from './typegen.js';
+import { typegenInternal } from './typegen.js';
 import { type LateSeedFiles, type LateTables, vitePluginDb } from './vite-plugin-db.js';
 import { vitePluginInjectEnvTs } from './vite-plugin-inject-env-ts.js';
 
@@ -88,7 +88,7 @@ function astroDBIntegration(): AstroIntegration {
 					await writeFile(localDbUrl, '');
 				}
 
-				await typegen({ tables: tables.get() ?? {}, root: config.root });
+				await typegenInternal({ tables: tables.get() ?? {}, root: config.root });
 			},
 			'astro:server:start': async ({ logger }) => {
 				// Wait for the server startup to log, so that this can come afterwards.
diff --git a/packages/db/src/core/integration/typegen.ts b/packages/db/src/core/integration/typegen.ts
index 9133c5dd4054..817cd79f899d 100644
--- a/packages/db/src/core/integration/typegen.ts
+++ b/packages/db/src/core/integration/typegen.ts
@@ -2,9 +2,18 @@ import { existsSync } from 'node:fs';
 import { mkdir, writeFile } from 'node:fs/promises';
 import { DB_TYPES_FILE, RUNTIME_IMPORT } from '../consts.js';
 import type { DBTable, DBTables } from '../types.js';
+import type { AstroConfig } from 'astro';
+import { resolveDbConfig } from '../load-file.js';
 
-export async function typegen({ tables, root }: { tables: DBTables; root: URL }) {
-	const content = `// This file is generated by \`studio sync\`
+// Exported for use in Astro core CLI
+export async function typegen(astroConfig: Pick<AstroConfig, 'root' | 'integrations'>) {
+	const { dbConfig } = await resolveDbConfig(astroConfig);
+
+	await typegenInternal({ tables: dbConfig.tables, root: astroConfig.root });
+}
+
+export async function typegenInternal({ tables, root }: { tables: DBTables; root: URL }) {
+	const content = `// This file is generated by Astro DB
 declare module 'astro:db' {
 	export const db: import(${RUNTIME_IMPORT}).SqliteDB;
 	export const dbUrl: string;
diff --git a/packages/db/src/core/load-file.ts b/packages/db/src/core/load-file.ts
index e4da7688ec8a..7bc7387c824c 100644
--- a/packages/db/src/core/load-file.ts
+++ b/packages/db/src/core/load-file.ts
@@ -18,7 +18,10 @@ const isDbIntegration = (integration: AstroIntegration): integration is AstroDbI
 /**
  * Load a user’s `astro:db` configuration file and additional configuration files provided by integrations.
  */
-export async function resolveDbConfig({ root, integrations }: AstroConfig) {
+export async function resolveDbConfig({
+	root,
+	integrations,
+}: Pick<AstroConfig, 'root' | 'integrations'>) {
 	const { mod, dependencies } = await loadUserConfigFile(root);
 	const userDbConfig = dbConfigSchema.parse(mod?.default ?? {}, { errorMap });
 	/** Resolved `astro:db` config including tables provided by integrations. */
diff --git a/packages/db/src/index.ts b/packages/db/src/index.ts
index a275433764d3..3e694354ebe0 100644
--- a/packages/db/src/index.ts
+++ b/packages/db/src/index.ts
@@ -1,3 +1,4 @@
 export type { ResolvedCollectionConfig, TableConfig } from './core/types.js';
 export { cli } from './core/cli/index.js';
 export { integration as default } from './core/integration/index.js';
+export { typegen } from './core/integration/typegen.js';