From 3ed9ae33e559a02fdb8c137e0ea21f9717ff450d Mon Sep 17 00:00:00 2001
From: "Fred K. Schott" <fkschott@gmail.com>
Date: Fri, 9 Feb 2024 10:50:01 -0800
Subject: [PATCH] wip

---
 packages/db/src/core/queries.ts | 36 +++++++++++++++++++++++----------
 packages/db/src/core/types.ts   |  3 ++-
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/packages/db/src/core/queries.ts b/packages/db/src/core/queries.ts
index d64149ee84..47e160e66c 100644
--- a/packages/db/src/core/queries.ts
+++ b/packages/db/src/core/queries.ts
@@ -16,6 +16,7 @@ import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core';
 import type { AstroIntegrationLogger } from 'astro';
 import type { DBUserConfig } from '../core/types.js';
 import { hasPrimaryKey } from '../runtime/index.js';
+import type { SQLiteInsert } from 'drizzle-orm/sqlite-core';
 
 const sqlite = new SQLiteAsyncDialect();
 
@@ -37,6 +38,10 @@ export async function setupDbTables({
 }) {
 	const setupQueries: SQL[] = [];
 	for (const [name, collection] of Object.entries(collections)) {
+		// We don't reset writable collections in production
+		if (mode === 'build' && collection.writable) {
+			continue;
+		}
 		const dropQuery = sql.raw(`DROP TABLE IF EXISTS ${name}`);
 		const createQuery = sql.raw(getCreateTableQuery(name, collection, useForeignKeys));
 		const indexQueries = getCreateIndexQueries(name, collection);
@@ -48,17 +53,26 @@ export async function setupDbTables({
 	if (data) {
 		try {
 			await data({
-				seed: async ({ table }, values) => {
-					const result = Array.isArray(values)
-						? db.insert(table).values(values).returning()
-						: db
-								.insert(table)
-								.values(values as any)
-								.returning()
-								.get();
-					// Drizzle types don't *quite* line up, and it's tough to debug why.
-					// we're casting and calling this close enough :)
-					return result as any;
+				seed: async (collection, values, options = {}) => {
+					const shouldRunInDev = (options.when === 'dev' || options.when === 'always')  ? true : true;
+					const shouldRunInBuild = (options.when === 'build' || options.when === 'always')  ? true : !collection.writable;
+					const shouldRun = mode === 'dev' ? shouldRunInDev : shouldRunInBuild;
+					if (options.returning && (!shouldRunInDev || !shouldRunInBuild)) {
+						throw new Error(`'returning' is set to true, but this data will not always seed. Set 'when' to 'always' to seed this data on every run, including production.`);
+					}
+					if (!shouldRun) {
+						return;
+					}
+					let result: SQLiteInsert<any, any, any, any> = db
+						.insert(collection.table)
+						.values(values as any);
+					if (options.returning) {
+						result = result.returning();
+					}
+					if (!Array.isArray(values)) {
+						result = result.get();
+					}
+					return await result.execute();
 				},
 				db,
 				mode,
diff --git a/packages/db/src/core/types.ts b/packages/db/src/core/types.ts
index ddb792fe62..a2ede0236e 100644
--- a/packages/db/src/core/types.ts
+++ b/packages/db/src/core/types.ts
@@ -207,7 +207,8 @@ export type DBDataContext = {
 		TData extends MaybeArray<SQLiteInsertValue<Table<string, TFields>>>,
 	>(
 		collection: ResolvedCollectionConfig<TFields>,
-		data: TData
+		data: TData,
+		options?: { returning?: boolean, when?: 'build' | 'dev' | 'always' | 'never'}
 	) => Promise<
 		TData extends Array<SQLiteInsertValue<Table<string, TFields>>>
 			? InferSelectModel<Table<string, TFields>>[]