1 files changed, 247 insertions, 0 deletions
diff --git a/packages/db/src/core/schemas.ts b/packages/db/src/core/schemas.ts
new file mode 100644
index 000000000..c9575a79a
--- /dev/null
+++ b/packages/db/src/core/schemas.ts
@@ -0,0 +1,247 @@
+import { SQL } from 'drizzle-orm';
+import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core';
+import { type ZodTypeDef, z } from 'zod';
+import { SERIALIZED_SQL_KEY, type SerializedSQL } from '../runtime/types.js';
+import { errorMap } from './integration/error-map.js';
+import type { NumberColumn, TextColumn } from './types.js';
+import { mapObject } from './utils.js';
+
+export type MaybeArray<T> = T | T[];
+
+// Transform to serializable object for migration files
+const sqlite = new SQLiteAsyncDialect();
+
+const sqlSchema = z.instanceof(SQL<any>).transform(
+	(sqlObj): SerializedSQL => ({
+		[SERIALIZED_SQL_KEY]: true,
+		sql: sqlite.sqlToQuery(sqlObj).sql,
+	}),
+);
+
+const baseColumnSchema = z.object({
+	label: z.string().optional(),
+	optional: z.boolean().optional().default(false),
+	unique: z.boolean().optional().default(false),
+	deprecated: z.boolean().optional().default(false),
+
+	// Defined when `defineDb()` is called to resolve `references`
+	name: z.string().optional(),
+	// TODO: Update to `table`. Will need migration file version change
+	collection: z.string().optional(),
+});
+
+export const booleanColumnSchema = z.object({
+	type: z.literal('boolean'),
+	schema: baseColumnSchema.extend({
+		default: z.union([z.boolean(), sqlSchema]).optional(),
+	}),
+});
+
+const numberColumnBaseSchema = baseColumnSchema.omit({ optional: true }).and(
+	z.union([
+		z.object({
+			primaryKey: z.literal(false).optional().default(false),
+			optional: baseColumnSchema.shape.optional,
+			default: z.union([z.number(), sqlSchema]).optional(),
+		}),
+		z.object({
+			// `integer primary key` uses ROWID as the default value.
+			// `optional` and `default` do not have an effect,
+			// so disable these config options for primary keys.
+			primaryKey: z.literal(true),
+			optional: z.literal(false).optional(),
+			default: z.literal(undefined).optional(),
+		}),
+	]),
+);
+
+export const numberColumnOptsSchema: z.ZodType<
+	z.infer<typeof numberColumnBaseSchema> & {
+		// ReferenceableColumn creates a circular type. Define ZodType to resolve.
+		references?: NumberColumn;
+	},
+	ZodTypeDef,
+	z.input<typeof numberColumnBaseSchema> & {
+		references?: () => z.input<typeof numberColumnSchema>;
+	}
+> = numberColumnBaseSchema.and(
+	z.object({
+		references: z
+			.function()
+			.returns(z.lazy(() => numberColumnSchema))
+			.optional()
+			.transform((fn) => fn?.()),
+	}),
+);
+
+export const numberColumnSchema = z.object({
+	type: z.literal('number'),
+	schema: numberColumnOptsSchema,
+});
+
+const textColumnBaseSchema = baseColumnSchema
+	.omit({ optional: true })
+	.extend({
+		default: z.union([z.string(), sqlSchema]).optional(),
+		multiline: z.boolean().optional(),
+	})
+	.and(
+		z.union([
+			z.object({
+				primaryKey: z.literal(false).optional().default(false),
+				optional: baseColumnSchema.shape.optional,
+			}),
+			z.object({
+				// text primary key allows NULL values.
+				// NULL values bypass unique checks, which could
+				// lead to duplicate URLs per record in Astro Studio.
+				// disable `optional` for primary keys.
+				primaryKey: z.literal(true),
+				optional: z.literal(false).optional(),
+			}),
+		]),
+	);
+
+export const textColumnOptsSchema: z.ZodType<
+	z.infer<typeof textColumnBaseSchema> & {
+		// ReferenceableColumn creates a circular type. Define ZodType to resolve.
+		references?: TextColumn;
+	},
+	ZodTypeDef,
+	z.input<typeof textColumnBaseSchema> & {
+		references?: () => z.input<typeof textColumnSchema>;
+	}
+> = textColumnBaseSchema.and(
+	z.object({
+		references: z
+			.function()
+			.returns(z.lazy(() => textColumnSchema))
+			.optional()
+			.transform((fn) => fn?.()),
+	}),
+);
+
+export const textColumnSchema = z.object({
+	type: z.literal('text'),
+	schema: textColumnOptsSchema,
+});
+
+export const dateColumnSchema = z.object({
+	type: z.literal('date'),
+	schema: baseColumnSchema.extend({
+		default: z
+			.union([
+				sqlSchema,
+				// transform to ISO string for serialization
+				z
+					.date()
+					.transform((d) => d.toISOString()),
+			])
+			.optional(),
+	}),
+});
+
+export const jsonColumnSchema = z.object({
+	type: z.literal('json'),
+	schema: baseColumnSchema.extend({
+		default: z.unknown().optional(),
+	}),
+});
+
+export const columnSchema = z.discriminatedUnion('type', [
+	booleanColumnSchema,
+	numberColumnSchema,
+	textColumnSchema,
+	dateColumnSchema,
+	jsonColumnSchema,
+]);
+export const referenceableColumnSchema = z.union([textColumnSchema, numberColumnSchema]);
+
+export const columnsSchema = z.record(columnSchema);
+
+type ForeignKeysInput = {
+	columns: MaybeArray<string>;
+	references: () => MaybeArray<Omit<z.input<typeof referenceableColumnSchema>, 'references'>>;
+};
+
+type ForeignKeysOutput = Omit<ForeignKeysInput, 'references'> & {
+	// reference fn called in `transform`. Ensures output is JSON serializable.
+	references: MaybeArray<Omit<z.output<typeof referenceableColumnSchema>, 'references'>>;
+};
+
+const foreignKeysSchema: z.ZodType<ForeignKeysOutput, ZodTypeDef, ForeignKeysInput> = z.object({
+	columns: z.string().or(z.array(z.string())),
+	references: z
+		.function()
+		.returns(z.lazy(() => referenceableColumnSchema.or(z.array(referenceableColumnSchema))))
+		.transform((fn) => fn()),
+});
+
+export const resolvedIndexSchema = z.object({
+	on: z.string().or(z.array(z.string())),
+	unique: z.boolean().optional(),
+});
+/** @deprecated */
+const legacyIndexesSchema = z.record(resolvedIndexSchema);
+
+export const indexSchema = z.object({
+	on: z.string().or(z.array(z.string())),
+	unique: z.boolean().optional(),
+	name: z.string().optional(),
+});
+const indexesSchema = z.array(indexSchema);
+
+export const tableSchema = z.object({
+	columns: columnsSchema,
+	indexes: indexesSchema.or(legacyIndexesSchema).optional(),
+	foreignKeys: z.array(foreignKeysSchema).optional(),
+	deprecated: z.boolean().optional().default(false),
+});
+
+export const tablesSchema = z.preprocess((rawTables) => {
+	// Use `z.any()` to avoid breaking object references
+	const tables = z.record(z.any()).parse(rawTables, { errorMap });
+	for (const [tableName, table] of Object.entries(tables)) {
+		// Append table and column names to columns.
+		// Used to track table info for references.
+		table.getName = () => tableName;
+		const { columns } = z.object({ columns: z.record(z.any()) }).parse(table, { errorMap });
+		for (const [columnName, column] of Object.entries(columns)) {
+			column.schema.name = columnName;
+			column.schema.collection = tableName;
+		}
+	}
+	return rawTables;
+}, z.record(tableSchema));
+
+export const dbConfigSchema = z
+	.object({
+		tables: tablesSchema.optional(),
+	})
+	.transform(({ tables = {}, ...config }) => {
+		return {
+			...config,
+			tables: mapObject(tables, (tableName, table) => {
+				const { indexes = {} } = table;
+				if (!Array.isArray(indexes)) {
+					return { ...table, indexes };
+				}
+				const resolvedIndexes: Record<string, z.infer<typeof resolvedIndexSchema>> = {};
+				for (const index of indexes) {
+					if (index.name) {
+						const { name, ...rest } = index;
+						resolvedIndexes[index.name] = rest;
+						continue;
+					}
+					// Sort index columns to ensure consistent index names
+					const indexOn = Array.isArray(index.on) ? index.on.sort().join('_') : index.on;
+					const name = tableName + '_' + indexOn + '_idx';
+					resolvedIndexes[name] = index;
+				}
+				return {
+					...table,
+					indexes: resolvedIndexes,
+				};
+			}),
+		};
+	});