feat: new reference() API that fixes type inference

examples/with-data/src/content/config.ts

@@ -1,12 +1,4 @@
-import { defineCollection, defineDataCollection, z, reference } from 'astro:content';
-
-const blog = defineCollection({
-	schema: z.object({
-		title: z.string(),
-		banner: reference('banners'),
-		author: reference('authors'),
-	}),
-});
+import { defineCollection, defineDataCollection, z } from 'astro:content';
 
 const banners = defineDataCollection({
 	schema: ({ image }) =>
@@ -24,4 +16,12 @@ const authors = defineDataCollection({
 		}),
 });
 
+const blog = defineCollection({
+	schema: z.object({
+		title: z.string(),
+		banner: banners.reference(),
+		author: authors.reference(),
+	}),
+});
+
 export const collections = { blog, authors, banners };

examples/with-data/src/data/authors/fred.json

@@ -0,0 +1,3 @@
+{
+  "name": "Fred K. Schott"
+}

packages/astro/src/content/runtime.ts

@@ -247,18 +247,21 @@ async function render({
 }
 
 export function createReference({
-	dataCollectionToEntryMap,
+	getCollectionName,
+	map,
 }: {
-	dataCollectionToEntryMap: CollectionToEntryMap;
+	getCollectionName(): Promise<string>;
+	map: CollectionToEntryMap;
 }) {
-	return function reference(collection: string) {
+	return function reference() {
 		return z.string().transform(async (entryId: string, ctx) => {
+			const collectionName = await getCollectionName();
 			const flattenedErrorPath = ctx.path.join('.');
-			const entries = dataCollectionToEntryMap[collection];
+			const entries = map[collectionName];
 			if (!entries) {
 				ctx.addIssue({
 					code: z.ZodIssueCode.custom,
-					message: `**${flattenedErrorPath}:** Reference to ${collection} invalid. Collection does not exist or is empty.`,
+					message: `**${flattenedErrorPath}:** Reference to ${collectionName} invalid. Collection does not exist or is empty.`,
 				});
 				return;
 			}
@@ -271,7 +274,7 @@ export function createReference({
 				);
 				ctx.addIssue({
 					code: z.ZodIssueCode.custom,
-					message: `**${flattenedErrorPath}**: Reference to ${collection} invalid. Expected ${entryKeys
+					message: `**${flattenedErrorPath}**: Reference to ${collectionName} invalid. Expected ${entryKeys
 						.map((c) => JSON.stringify(c))
 						.join(' | ')}. Received ${JSON.stringify(entryId)}.`,
 				});
@@ -295,7 +298,7 @@ export function createReference({
 				} else {
 					ctx.addIssue({
 						code: z.ZodIssueCode.custom,
-						message: `**${flattenedErrorPath}:** Referenced entry ${collection} → ${entryId} is invalid.`,
+						message: `**${flattenedErrorPath}:** Referenced entry ${collectionName} → ${entryId} is invalid.`,
 					});
 				}
 			}

packages/astro/src/content/template/types.d.ts

@@ -67,20 +67,21 @@ declare module 'astro:content' {
 	type ContentCollectionConfig<S extends BaseSchema> = {
 		type: 'content';
 		schema?: S | ((context: SchemaContext) => S);
+		reference(): import('astro/zod').ZodEffects<S>;
 	};
 
 	type DataCollectionConfig<S extends BaseSchema> = {
 		type: 'data';
-		key?: string;
 		schema?: S | ((context: SchemaContext) => S);
+		reference(): import('astro/zod').ZodEffects<S>;
 	};
 
 	export function defineCollection<S extends BaseSchema>(
-		input: Omit<ContentCollectionConfig<S>, 'type'>
+		input: Omit<ContentCollectionConfig<S>, 'type' | 'reference'>
 	): ContentCollectionConfig<S>;
 
 	export function defineDataCollection<S extends BaseSchema>(
-		input: Omit<DataCollectionConfig<S>, 'type'>
+		input: Omit<DataCollectionConfig<S>, 'type' | 'reference'>
 	): DataCollectionConfig<S>;
 
 	type AllValuesOf<T> = T extends any ? T[keyof T] : never;

packages/astro/src/content/template/virtual-mod.mjs

@@ -10,21 +10,6 @@ import {
 
 export { z } from 'astro/zod';
 
-export function defineCollection(config) {
-	return Object.assign(config, { type: 'content' });
-}
-
-export function defineDataCollection(config) {
-	return Object.assign(config, { type: 'data' });
-}
-
-// TODO: Remove this when having this fallback is no longer relevant. 2.3? 3.0? - erika, 2023-04-04
-export const image = () => {
-	throw new Error(
-		'Importing `image()` from `astro:content` is no longer supported. See https://docs.astro.build/en/guides/assets/#update-content-collections-schemas for our new import instructions.'
-	);
-};
-
 const contentDir = '@@CONTENT_DIR@@';
 const dataDir = '@@DATA_DIR@@';
 
@@ -52,6 +37,45 @@ const collectionToRenderEntryMap = createCollectionToGlobResultMap({
 	dir: contentDir,
 });
 
+let referenceKeyIncr = 0;
+
+/**
+ * @param {any} partialConfig
+ * @param {'content' | 'data'} type
+ */
+function baseDefineCollection(partialConfig, type) {
+	// We don't know the collection name since this is defined on the `collections` export.
+	// Generate a unique key for the collection that we can use for lookups.
+	const referenceKey = String(referenceKeyIncr++);
+	return {
+		...partialConfig,
+		type,
+		reference: createReference({
+			map: type === 'content' ? contentCollectionToEntryMap : dataCollectionToEntryMap,
+			async getCollectionName() {
+				const { default: map } = await import('@@COLLECTION_NAME_BY_REFERENCE_KEY@@');
+				return map[referenceKey];
+			},
+		}),
+		referenceKey,
+	};
+}
+
+export function defineCollection(config) {
+	return baseDefineCollection(config, 'content');
+}
+
+export function defineDataCollection(config) {
+	return baseDefineCollection(config, 'data');
+}
+
+// TODO: Remove this when having this fallback is no longer relevant. 2.3? 3.0? - erika, 2023-04-04
+export const image = () => {
+	throw new Error(
+		'Importing `image()` from `astro:content` is no longer supported. See https://docs.astro.build/en/guides/assets/#update-content-collections-schemas for our new import instructions.'
+	);
+};
+
 export const getCollection = createGetCollection({
 	contentCollectionToEntryMap,
 	collectionToRenderEntryMap,

packages/astro/src/content/utils.ts

@@ -17,10 +17,12 @@ export const collectionConfigParser = z.union([
 	z.object({
 		type: z.literal('content').optional().default('content'),
 		schema: z.any().optional(),
+		referenceKey: z.string(),
 	}),
 	z.object({
 		type: z.literal('data'),
 		schema: z.any().optional(),
+		referenceKey: z.string(),
 	}),
 ]);
 
@@ -337,6 +339,18 @@ export async function loadContentConfig({
 	}
 	const config = contentConfigParser.safeParse(unparsedConfig);
 	if (config.success) {
+		let collectionNameByReferenceKey: Record<string, string> = {};
+		for (const [collectionName, collection] of Object.entries(config.data.collections)) {
+			collectionNameByReferenceKey[collection.referenceKey] = collectionName;
+		}
+
+		// We need a way to map collection config references *back* to their collection name.
+		// This generates a JSON map we can import when querying.
+		await fs.promises.writeFile(
+			new URL('reference-map.json', contentPaths.cacheDir),
+			JSON.stringify(collectionNameByReferenceKey)
+		);
+
 		// Check that data collections are properly configured using `defineDataCollection()`.
 		const dataEntryExts = getDataEntryExts(settings);
 

packages/astro/src/content/vite-plugin-content-virtual-mod.ts

@@ -23,6 +23,10 @@ export function astroContentVirtualModPlugin({
 
 	const virtualModContents = fsMod
 		.readFileSync(contentPaths.virtualModTemplate, 'utf-8')
+		.replace(
+			'@@COLLECTION_NAME_BY_REFERENCE_KEY@@',
+			new URL('reference-map.json', contentPaths.cacheDir).pathname
+		)
 		.replace('@@CONTENT_DIR@@', relContentDir)
 		.replace('@@DATA_DIR@@', relDataDir)
 		.replace('@@CONTENT_ENTRY_GLOB_PATH@@', `${relContentDir}**/*${getExtGlob(contentEntryExts)}`)