feat: FINALLY collection and name attached

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

@@ -6,7 +6,7 @@ import { existsSync } from 'fs';
 import { mkdir, rm, writeFile } from 'fs/promises';
 import { DB_PATH } from '../consts.js';
 import { createLocalDatabaseClient } from '../../runtime/db-client.js';
-import { astroConfigWithDbSchema } from '../types.js';
+import { astroConfigWithDbSchema, attachTableMetaHandler } from '../types.js';
 import { getAstroStudioEnv, type VitePlugin } from '../utils.js';
 import { appTokenError } from '../errors.js';
 import { errorMap } from './error-map.js';
@@ -15,6 +15,7 @@ import { fileURLToPath } from 'url';
 import { bold } from 'kleur/colors';
 import { fileURLIntegration } from './file-url.js';
 import { setupDbTables } from '../queries.js';
+import { collectionToTable } from '../../runtime/index.js';
 
 function astroDBIntegration(): AstroIntegration {
 	return {
@@ -27,6 +28,8 @@ function astroDBIntegration(): AstroIntegration {
 				// @matthewp: may want to load collections by path at runtime
 				const configWithDb = astroConfigWithDbSchema.parse(config, { errorMap });
 				const collections = configWithDb.db?.collections ?? {};
+				setCollectionsMeta(collections);
+
 				const studio = configWithDb.db?.studio ?? false;
 				if (!studio && Object.values(collections).some((c) => c.writable)) {
 					logger.warn(
@@ -105,6 +108,19 @@ function astroDBIntegration(): AstroIntegration {
 	};
 }
 
+/**
+ * We need to attach the Drizzle `table` and collection name at runtime.
+ * These cannot be determined from `defineCollection()`,
+ * since we don't know the collection name until the `db` config is resolved.
+ */
+function setCollectionsMeta(collections: Record<string, any>) {
+	for (const [name, collection] of Object.entries(collections)) {
+		const table = collectionToTable(name, collection);
+		collection[name] = attachTableMetaHandler(collection);
+		collection[name]._setMeta?.({ table });
+	}
+}
+
 export function integration(): AstroIntegration[] {
 	return [astroDBIntegration(), fileURLIntegration()];
 }

packages/db/src/core/queries.ts

@@ -14,7 +14,7 @@ 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 { DBUserConfig, ReferenceableField } from '../core/types.js';
 import { collectionToTable, hasPrimaryKey } from '../runtime/index.js';
 
 const sqlite = new SQLiteAsyncDialect();
@@ -43,10 +43,6 @@ export async function setupDbTables({
 		await db.run(q);
 	}
 	if (data) {
-		for (const [name, collection] of Object.entries(collections)) {
-			const table = collectionToTable(name, collection);
-			collection._setMeta?.({ table });
-		}
 		try {
 			await data({
 				async seed({ table }, values) {
@@ -139,9 +135,26 @@ export function getModifiers(fieldName: string, field: DBField) {
 	if (hasDefault(field)) {
 		modifiers += ` DEFAULT ${getDefaultValueSql(fieldName, field)}`;
 	}
+	const references = getReferencesConfig(field);
+	if (references) {
+		const { collection, name } = references;
+		if (!collection || !name) {
+			throw new Error(
+				`Invalid reference for field ${fieldName}. This is an unexpected error that should be reported to the Astro team.`
+			);
+		}
+
+		modifiers += ` REFERENCES ${sqlite.escapeName(collection)} (${sqlite.escapeName(name)})`;
+	}
 	return modifiers;
 }
 
+function getReferencesConfig(field: DBField) {
+	const canHaveReferences = field.type === 'number' || field.type === 'text';
+	if (!canHaveReferences) return undefined;
+	return field.references as ReferenceableField;
+}
+
 // Using `DBField` will not narrow `default` based on the column `type`
 // Handle each field separately
 type WithDefaultDefined<T extends DBField> = T & Required<Pick<T, 'default'>>;

packages/db/src/core/types.ts

@@ -1,6 +1,7 @@
 import type { SQLiteInsertValue } from 'drizzle-orm/sqlite-core';
 import type { SqliteDB, Table } from '../runtime/index.js';
 import { z } from 'zod';
+import { getTableName } from 'drizzle-orm';
 
 export type MaybePromise<T> = T | Promise<T>;
 export type MaybeArray<T> = T | T[];
@@ -9,6 +10,10 @@ const baseFieldSchema = z.object({
 	label: z.string().optional(),
 	optional: z.boolean().optional(),
 	unique: z.boolean().optional(),
+
+	// Defined when `defineCollection()` is called
+	name: z.string().optional(),
+	collection: z.string().optional(),
 });
 
 const booleanFieldSchema = baseFieldSchema.extend({
@@ -16,36 +21,22 @@ const booleanFieldSchema = baseFieldSchema.extend({
 	default: z.boolean().optional(),
 });
 
-// NOTE (bholmesdev): `references` creates a recursive type. This is not supported by zod.
-// Declare `NumberField` and `TextField` manually and use `z.lazy()` in schema.
-// see https://zod.dev/?id=recursive-types
-export type NumberField = z.infer<typeof baseFieldSchema> & {
-	type: 'number';
-	default?: number;
-	references?: ReferenceableField;
-	primaryKey?: boolean;
-};
-
-const numberFieldSchema: z.ZodType<NumberField> = baseFieldSchema.extend({
+const numberFieldSchema = baseFieldSchema.extend({
 	type: z.literal('number'),
 	default: z.number().optional(),
-	references: z.lazy(() => referenceableFieldSchema).optional(),
+	// Need to avoid `z.object()`. Otherwise, object references are broken,
+	// and we cannot set the `collection` field at runtime.
+	references: z.any().optional(),
 	primaryKey: z.boolean().optional(),
 });
 
-export type TextField = z.infer<typeof baseFieldSchema> & {
-	type: 'text';
-	multiline?: boolean;
-	default?: string;
-	references?: ReferenceableField;
-	primaryKey?: boolean;
-};
-
-const textFieldSchema: z.ZodType<TextField> = baseFieldSchema.extend({
+const textFieldSchema = baseFieldSchema.extend({
 	type: z.literal('text'),
 	multiline: z.boolean().optional(),
 	default: z.string().optional(),
-	references: z.lazy(() => referenceableFieldSchema).optional(),
+	// Need to avoid `z.object()`. Otherwise, object references are broken,
+	// and we cannot set the `collection` field at runtime.
+	references: z.any().optional(),
 	primaryKey: z.boolean().optional(),
 });
 
@@ -109,6 +100,8 @@ export const collectionSchema = z.union([readableCollectionSchema, writableColle
 export const collectionsSchema = z.record(collectionSchema);
 
 export type BooleanField = z.infer<typeof booleanFieldSchema>;
+export type NumberField = z.infer<typeof numberFieldSchema>;
+export type TextField = z.infer<typeof textFieldSchema>;
 export type DateField = z.infer<typeof dateFieldSchema>;
 // Type `Date` is the config input, `string` is the output for D1 storage
 export type DateFieldInput = z.input<typeof dateFieldSchema>;
@@ -209,40 +202,63 @@ export type ResolvedCollectionConfig<
 	table: Table<string, TFields>;
 };
 
-export function defineCollection<TFields extends FieldsConfig>(
-	userConfig: CollectionConfig<TFields>
-): ResolvedCollectionConfig<TFields, false> {
-	const meta: CollectionMeta<TFields> = { table: null! };
-	function _setMeta(values: CollectionMeta<TFields>) {
+/**
+ * Handler to attach the Drizzle `table` and collection name at runtime.
+ * These cannot be determined from `defineCollection()`,
+ * since we don't know the collection name until the `db` config is resolved.
+ */
+export function attachTableMetaHandler<TFields extends FieldsConfig, TWritable extends boolean>(
+	collectionConfig: ResolvedCollectionConfig<TFields, TWritable>
+): ResolvedCollectionConfig<TFields, TWritable> {
+	const meta: CollectionMeta<TFields> = { table: collectionConfig.table };
+	const _setMeta = (values: CollectionMeta<TFields>) => {
+		// `_setMeta` is called twice: once from the user's config,
+		// and once after the config is parsed via Zod.
+		(collectionConfig as any)._setMeta?.(values);
 		Object.assign(meta, values);
+
+		const tableName = getTableName(meta.table);
+		for (const fieldName in collectionConfig.fields) {
+			const field = collectionConfig.fields[fieldName];
+			field.collection = tableName;
+		}
+	};
+	for (const fieldName in collectionConfig.fields) {
+		const field = collectionConfig.fields[fieldName];
+		field.name = fieldName;
 	}
+
 	return {
-		...userConfig,
-		writable: false,
+		...collectionConfig,
 		get table() {
 			return meta.table;
 		},
-		// @ts-expect-error private field
+		// @ts-expect-error private setter
 		_setMeta,
 	};
 }
 
+function baseDefineCollection<TFields extends FieldsConfig, TWritable extends boolean>(
+	userConfig: CollectionConfig<TFields>,
+	writable: TWritable
+): ResolvedCollectionConfig<TFields, TWritable> {
+	return attachTableMetaHandler({
+		...userConfig,
+		writable,
+		table: null!,
+	});
+}
+
+export function defineCollection<TFields extends FieldsConfig>(
+	userConfig: CollectionConfig<TFields>
+): ResolvedCollectionConfig<TFields, false> {
+	return baseDefineCollection(userConfig, false);
+}
+
 export function defineWritableCollection<TFields extends FieldsConfig>(
 	userConfig: CollectionConfig<TFields>
 ): ResolvedCollectionConfig<TFields, true> {
-	const meta: CollectionMeta<TFields> = { table: null! };
-	function _setMeta(values: CollectionMeta<TFields>) {
-		Object.assign(meta, values);
-	}
-	return {
-		...userConfig,
-		writable: true,
-		get table() {
-			return meta.table;
-		},
-		// @ts-expect-error private field
-		_setMeta,
-	};
+	return baseDefineCollection(userConfig, true);
 }
 
 export type AstroConfigWithDB = z.infer<typeof astroConfigWithDbSchema>;