From 463796c48b47aeaa0d078ec61ee34974ff24b551 Mon Sep 17 00:00:00 2001
From: bholmesdev <hey@bholmes.dev>
Date: Tue, 13 Feb 2024 11:37:29 -0500
Subject: [PATCH] feat: seed -> seedReturning

---
 packages/db/src/core/queries.ts               | 25 ++++++++++---------
 packages/db/src/core/types.ts                 | 18 +++++++------
 .../db/test/fixtures/recipes/astro.config.ts  | 10 ++++----
 3 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/packages/db/src/core/queries.ts b/packages/db/src/core/queries.ts
index 7ae2eb2b47..8f6c6fbd0f 100644
--- a/packages/db/src/core/queries.ts
+++ b/packages/db/src/core/queries.ts
@@ -12,7 +12,7 @@ import {
 } from '../core/types.js';
 import { bold } from 'kleur/colors';
 import { type SQL, sql } from 'drizzle-orm';
-import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core';
+import { SQLiteAsyncDialect, type SQLiteInsert } from 'drizzle-orm/sqlite-core';
 import type { AstroIntegrationLogger } from 'astro';
 import type { DBUserConfig } from '../core/types.js';
 import { hasPrimaryKey } from '../runtime/index.js';
@@ -48,16 +48,17 @@ export async function setupDbTables({
 		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;
+					await db.insert(table).values(values as any);
+				},
+				seedReturning: async ({ table }, values) => {
+					let result: SQLiteInsert<any, any, any, any> = db
+						.insert(table)
+						.values(values as any)
+						.returning();
+					if (!Array.isArray(values)) {
+						result = result.get();
+					}
+					return result;
 				},
 				db,
 				mode,
@@ -192,7 +193,7 @@ export function getReferencesConfig(field: DBField) {
 // Using `DBField` will not narrow `default` based on the column `type`
 // Handle each field separately
 type WithDefaultDefined<T extends DBField> = T & {
-	schema: Required<Pick<T['schema'], 'default'>>
+	schema: Required<Pick<T['schema'], 'default'>>;
 };
 type DBFieldWithDefault =
 	| WithDefaultDefined<TextField>
diff --git a/packages/db/src/core/types.ts b/packages/db/src/core/types.ts
index a71ec850d7..c95780a96c 100644
--- a/packages/db/src/core/types.ts
+++ b/packages/db/src/core/types.ts
@@ -75,7 +75,7 @@ const numberFieldOptsSchema: z.ZodType<
 
 const numberFieldSchema = z.object({
 	type: z.literal('number'),
-	schema: numberFieldOptsSchema
+	schema: numberFieldOptsSchema,
 });
 
 const textFieldBaseSchema = baseFieldSchema
@@ -122,7 +122,7 @@ const textFieldOptsSchema: z.ZodType<
 
 const textFieldSchema = z.object({
 	type: z.literal('text'),
-	schema: textFieldOptsSchema
+	schema: textFieldOptsSchema,
 });
 
 const dateFieldSchema = z.object({
@@ -136,14 +136,14 @@ const dateFieldSchema = z.object({
 				z.coerce.date().transform((d) => d.toISOString()),
 			])
 			.optional(),
-	})
+	}),
 });
 
 const jsonFieldSchema = z.object({
 	type: z.literal('json'),
 	schema: baseFieldSchema.extend({
 		default: z.unknown().optional(),
-	})
+	}),
 });
 
 const fieldSchema = z.union([
@@ -261,7 +261,11 @@ export type WritableDBCollection = z.infer<typeof writableCollectionSchema>;
 
 export type DBDataContext = {
 	db: SqliteDB;
-	seed: <
+	seed: <TFields extends FieldsConfig>(
+		collection: ResolvedCollectionConfig<TFields>,
+		data: MaybeArray<SQLiteInsertValue<Table<string, TFields>>>
+	) => Promise<void>;
+	seedReturning: <
 		TFields extends FieldsConfig,
 		TData extends MaybeArray<SQLiteInsertValue<Table<string, TFields>>>,
 	>(
@@ -361,7 +365,7 @@ function createField<S extends string, T extends Record<string, unknown>>(type:
 		/**
 		 * @internal
 		 */
-		schema
+		schema,
 	};
 }
 
@@ -370,7 +374,7 @@ export const field = {
 		return createField('number', opts) satisfies { type: 'number' };
 	},
 	boolean: <T extends FieldOpts<BooleanFieldInput>>(opts: T = {} as T) => {
-		return createField('boolean', opts) satisfies { type: 'boolean' }
+		return createField('boolean', opts) satisfies { type: 'boolean' };
 	},
 	text: <T extends TextFieldOpts>(opts: T = {} as T) => {
 		return createField('text', opts) satisfies { type: 'text' };
diff --git a/packages/db/test/fixtures/recipes/astro.config.ts b/packages/db/test/fixtures/recipes/astro.config.ts
index e0d1a1489b..4a61f7c2c6 100644
--- a/packages/db/test/fixtures/recipes/astro.config.ts
+++ b/packages/db/test/fixtures/recipes/astro.config.ts
@@ -26,13 +26,13 @@ export default defineConfig({
 	integrations: [astroDb()],
 	db: {
 		collections: { Recipe, Ingredient },
-		async data({ seed }) {
-			const pancakes = await seed(Recipe, {
+		async data({ seed, seedReturning }) {
+			const pancakes = await seedReturning(Recipe, {
 				title: 'Pancakes',
 				description: 'A delicious breakfast',
 			});
 
-			seed(Ingredient, [
+			await seed(Ingredient, [
 				{
 					name: 'Flour',
 					quantity: 1,
@@ -50,12 +50,12 @@ export default defineConfig({
 				},
 			]);
 
-			const pizza = await seed(Recipe, {
+			const pizza = await seedReturning(Recipe, {
 				title: 'Pizza',
 				description: 'A delicious dinner',
 			});
 
-			seed(Ingredient, [
+			await seed(Ingredient, [
 				{
 					name: 'Flour',
 					quantity: 1,