From 911975333668269c90048d905a36832d236d2e97 Mon Sep 17 00:00:00 2001
From: Matthew Phillips <matthew@skypack.dev>
Date: Tue, 16 Jan 2024 17:00:29 -0500
Subject: [PATCH] Implement defineWritableCollection

---
 packages/db/src/config.ts                     | 42 ++++++++++++++-----
 packages/db/src/internal.ts                   | 11 +++--
 packages/db/src/types.ts                      | 18 ++++++--
 .../db/test/fixtures/basics/astro.config.ts   |  6 +--
 4 files changed, 57 insertions(+), 20 deletions(-)

diff --git a/packages/db/src/config.ts b/packages/db/src/config.ts
index f3c474f438..94a8c769f7 100644
--- a/packages/db/src/config.ts
+++ b/packages/db/src/config.ts
@@ -20,19 +20,41 @@ export const astroConfigWithDBValidator = z.object({
 	db: adjustedConfigSchema.optional(),
 });
 
-type CollectionConfig<TFields extends z.input<typeof collectionSchema>['fields']> = {
-	fields: TFields;
-	// TODO: type inference based on field type. Just `any` for now.
-	data?: () => Array<Record<keyof TFields, any> & { id?: string }>;
-	
-	// I don't love this name mp
-	source: 'readable' | 'writable';
+type CollectionConfig<TFields extends z.input<typeof collectionSchema>['fields'], Writable extends boolean> =
+	Writable extends true ? (
+		{
+			fields: TFields;
+			// TODO: type inference based on field type. Just `any` for now.
+			seed?: Writable extends false ? never : () => Array<Record<keyof TFields, any> & { id?: string }>;
+		}
+	) : (
+		{
+			fields: TFields;
+			// TODO: type inference based on field type. Just `any` for now.
+			data?: Writable extends true ? never : () => Array<Record<keyof TFields, any> & { id?: string }>;
+		}
+	)
+
+type ResolvedCollectionConfig<TFields extends z.input<typeof collectionSchema>['fields'], Writable extends boolean> = CollectionConfig<TFields, Writable> & {
+	writable: Writable;
 };
 
 export function defineCollection<TFields extends z.input<typeof collectionSchema>['fields']>(
-	userConfig: CollectionConfig<TFields>
-): CollectionConfig<TFields> {
-	return userConfig;
+	userConfig: CollectionConfig<TFields, false>
+): ResolvedCollectionConfig<TFields, false> {
+	return {
+		...userConfig,
+		writable: false
+	};
+}
+
+export function defineWritableCollection<TFields extends z.input<typeof collectionSchema>['fields']>(
+	userConfig: CollectionConfig<TFields, true>
+): ResolvedCollectionConfig<TFields, true> {
+	return {
+		...userConfig,
+		writable: true
+	};
 }
 
 export type AstroConfigWithDB = z.infer<typeof astroConfigWithDBValidator>;
diff --git a/packages/db/src/internal.ts b/packages/db/src/internal.ts
index 9e80df52a6..b8cef7c6bb 100644
--- a/packages/db/src/internal.ts
+++ b/packages/db/src/internal.ts
@@ -1,7 +1,8 @@
 import type { SqliteRemoteDatabase } from 'drizzle-orm/sqlite-proxy';
 import { createClient } from '@libsql/client';
 import {
-	collectionSchema,
+	type ReadableDBCollection,
+	type WritableDBCollection,
 	type BooleanField,
 	type DBCollection,
 	type DBCollections,
@@ -40,11 +41,15 @@ export type {
 
 const sqlite = new SQLiteAsyncDialect();
 
+function isReadableCollection(collection: DBCollection): collection is ReadableDBCollection {
+	return !collection.writable;
+}
+
 function checkIfModificationIsAllowed(collections: DBCollections, Table: SQLiteTable) {
 	// This totally works, don't worry about it
 	const tableName = (Table as any)[(SQLiteTable as any).Symbol.Name];
 	const collection = collections[tableName];
-	if(collection.source === 'readable') {
+	if(collection.writable) {
 		throw new Error(`The [${tableName}] collection is read-only.`);
 	}
 }
@@ -100,7 +105,7 @@ export async function setupDbTables({
 		await db.run(q);
 	}
 	for (const [name, collection] of Object.entries(collections)) {
-		if (!collection.data) continue;
+		if (!isReadableCollection(collection) || !collection.data) continue;
 
 		const table = collectionToTable(name, collection);
 		try {
diff --git a/packages/db/src/types.ts b/packages/db/src/types.ts
index 1cef4049a8..2e03b9781e 100644
--- a/packages/db/src/types.ts
+++ b/packages/db/src/types.ts
@@ -50,15 +50,25 @@ const fieldSchema = z.union([
 ]);
 const fieldsSchema = z.record(fieldSchema);
 
-export const collectionSchema = z.object({
+export const readableCollectionSchema = z.object({
 	fields: fieldsSchema,
 	data: z
 		.function()
 		.returns(z.array(z.record(z.unknown())))
 		.optional(),
-	source: z.enum(['readable', 'writable'])
+	writable: z.literal(false)
 });
 
+export const writableCollectionSchema = z.object({
+	fields: fieldsSchema,
+	seed: z
+		.function()
+		.returns(z.array(z.record(z.unknown())))
+		.optional(),
+	writable: z.literal(true)
+});
+
+export const collectionSchema = z.union([readableCollectionSchema, writableCollectionSchema]);
 export const collectionsSchema = z.record(collectionSchema);
 
 export type BooleanField = z.infer<typeof booleanFieldSchema>;
@@ -79,8 +89,10 @@ export type FieldType =
 export type DBField = z.infer<typeof fieldSchema>;
 export type DBFieldInput = DateFieldInput | BooleanField | NumberField | TextField | JsonField;
 export type DBFields = z.infer<typeof fieldsSchema>;
-export type DBCollection = z.infer<typeof collectionSchema>;
+export type DBCollection = z.infer<typeof readableCollectionSchema | typeof writableCollectionSchema>;
 export type DBCollections = Record<string, DBCollection>;
+export type ReadableDBCollection = z.infer<typeof readableCollectionSchema>;
+export type WritableDBCollection = z.infer<typeof writableCollectionSchema>;
 
 export type AstroTable<T extends Pick<TableConfig, 'name' | 'columns'>> = SQLiteTableWithColumns<
 	T & {
diff --git a/packages/db/test/fixtures/basics/astro.config.ts b/packages/db/test/fixtures/basics/astro.config.ts
index 7b54863bf9..28630fca73 100644
--- a/packages/db/test/fixtures/basics/astro.config.ts
+++ b/packages/db/test/fixtures/basics/astro.config.ts
@@ -1,11 +1,10 @@
 import { defineConfig } from 'astro/config';
-import db, { defineCollection, field } from '@astrojs/db';
+import db, { defineCollection, defineWritableCollection, field } from '@astrojs/db';
 
 const Author = defineCollection({
 	fields: {
 		name: field.text(),
 	},
-	source: 'readable',
 	data() {
 		return [
 			{
@@ -27,11 +26,10 @@ const Author = defineCollection({
 	}
 });
 
-const Themes = defineCollection({
+const Themes = defineWritableCollection({
 	fields: {
 		name: field.text(),
 	},
-	source: 'writable',
 	data() {
 		return [
 			{