fix: pass schema to drizzle client for better type generation (#459)

* fix: pass `schema` to `drizzle` client for better type generation

* see if tests work

* Revert "see if tests work"

This reverts commit 3910a48.

* fix everything

Author: manuel3108

.changeset/swift-poets-wonder.md

@@ -0,0 +1,5 @@
+---
+'sv': patch
+---
+
+fix: pass `schema` to `drizzle` client for better type generation

packages/addons/drizzle/index.ts

@@ -261,6 +261,7 @@ export default defineAddon({
 			const { ast, generateCode } = parseScript(content);
 
 			imports.addNamed(ast, '$env/dynamic/private', { env: 'env' });
+			imports.addNamespace(ast, './schema', 'schema');
 
 			// env var checks
 			const dbURLCheck = common.statementFromString(
@@ -296,20 +297,14 @@ export default defineAddon({
 				}
 			}
 			// MySQL
-			if (options.mysql === 'mysql2') {
+			if (options.mysql === 'mysql2' || options.mysql === 'planetscale') {
 				imports.addDefault(ast, 'mysql2/promise', 'mysql');
 				imports.addNamed(ast, 'drizzle-orm/mysql2', { drizzle: 'drizzle' });
 
 				clientExpression = common.expressionFromString(
 					'await mysql.createConnection(env.DATABASE_URL)'
 				);
 			}
-			if (options.mysql === 'planetscale') {
-				imports.addNamed(ast, '@planetscale/database', { Client: 'Client' });
-				imports.addNamed(ast, 'drizzle-orm/planetscale-serverless', { drizzle: 'drizzle' });
-
-				clientExpression = common.expressionFromString('new Client({ url: env.DATABASE_URL })');
-			}
 			// PostgreSQL
 			if (options.postgresql === 'neon') {
 				imports.addNamed(ast, '@neondatabase/serverless', { neon: 'neon' });
@@ -328,7 +323,20 @@ export default defineAddon({
 			const clientIdentifier = variables.declaration(ast, 'const', 'client', clientExpression);
 			common.addStatement(ast, clientIdentifier);
 
+			// create drizzle function call
 			const drizzleCall = functions.callByIdentifier('drizzle', ['client']);
+
+			// add schema to support `db.query`
+			const paramObject = object.create({
+				schema: variables.identifier('schema')
+			});
+			if (options.database == 'mysql') {
+				const mode = options.mysql == 'planetscale' ? 'planetscale' : 'default';
+				object.property(paramObject, 'mode', common.createLiteral(mode));
+			}
+			drizzleCall.arguments.push(paramObject);
+
+			// create `db` export
 			const db = variables.declaration(ast, 'const', 'db', drizzleCall);
 			exports.namedExport(ast, 'db', db);