From e3ba7de572cd0419d9b5c6e63ecb5b3d95c2e372 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               | 23 ++++++++++---------
 packages/db/src/core/types.ts                 |  6 ++++-
 .../db/test/fixtures/recipes/astro.config.ts  | 10 ++++----
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/packages/db/src/core/queries.ts b/packages/db/src/core/queries.ts
index b35f2d09e0..e05031fe19 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';
@@ -53,16 +53,17 @@ export async function seedData({
 	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,
diff --git a/packages/db/src/core/types.ts b/packages/db/src/core/types.ts
index 5a9b7d89fc..e371bb246d 100644
--- a/packages/db/src/core/types.ts
+++ b/packages/db/src/core/types.ts
@@ -260,7 +260,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>>>,
 	>(
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,