Skip to content

Commit 5604a6f

Browse files
Refactor localExecution handling in vitePluginDb and related functions
1 parent 9fde9ec commit 5604a6f

File tree

1 file changed

+47
-13
lines changed

1 file changed

+47
-13
lines changed

packages/db/src/core/integration/vite-plugin-db.ts

Lines changed: 47 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,14 @@ import type { AstroConfig, AstroIntegrationLogger } from 'astro';
44
import { type SQL, sql } from 'drizzle-orm';
55
import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core';
66
import { normalizeDatabaseUrl } from '../../runtime/index.js';
7-
import { DB_CLIENTS, DB_PATH, RUNTIME_IMPORT, RUNTIME_VIRTUAL_IMPORT, VIRTUAL_CLIENT_MODULE_ID, VIRTUAL_MODULE_ID } from '../consts.js';
7+
import {
8+
DB_CLIENTS,
9+
DB_PATH,
10+
RUNTIME_IMPORT,
11+
RUNTIME_VIRTUAL_IMPORT,
12+
VIRTUAL_CLIENT_MODULE_ID,
13+
VIRTUAL_MODULE_ID,
14+
} from '../consts.js';
815
import { createClient } from '../db-client/libsql-local.js';
916
import { getResolvedFileUrl } from '../load-file.js';
1017
import { getCreateIndexQueries, getCreateTableQuery, SEED_DEV_FILE_NAME } from '../queries.js';
@@ -77,6 +84,7 @@ export function vitePluginDb(params: VitePluginDBParams): VitePlugin {
7784
tables: params.tables.get(),
7885
isBuild: command === 'build',
7986
output: params.output,
87+
localExecution: false,
8088
});
8189
}
8290

@@ -87,6 +95,7 @@ export function vitePluginDb(params: VitePluginDBParams): VitePlugin {
8795
return getLocalVirtualModContents({
8896
root: params.root,
8997
tables: params.tables.get(),
98+
localExecution: false,
9099
});
91100
}
92101

@@ -108,6 +117,7 @@ export function vitePluginDb(params: VitePluginDBParams): VitePlugin {
108117
return getLocalVirtualModContents({
109118
root: params.root,
110119
tables: params.tables.get(),
120+
localExecution: false,
111121
});
112122
},
113123
};
@@ -117,23 +127,43 @@ export function getConfigVirtualModContents() {
117127
return `export * from ${RUNTIME_VIRTUAL_IMPORT}`;
118128
}
119129

130+
/**
131+
* Get the module import for the DB client.
132+
* This is used to pick which module to import based on whether
133+
* the DB client is being used by the CLI, or in the Astro runtime.
134+
*
135+
* This is important for the `astro db execute` command to work correctly.
136+
*
137+
* @param localExecution - Whether the DB client is being used in a local execution context (e.g. CLI commands).
138+
* @returns The module import string for the DB client.
139+
*/
140+
function getDBModule(localExecution: boolean) {
141+
return localExecution
142+
? `import { createClient } from '${DB_CLIENTS.node}';`
143+
: `import { createClient } from '${VIRTUAL_CLIENT_MODULE_ID}';`;
144+
}
145+
120146
export function getLocalVirtualModContents({
121147
tables,
122148
root,
123-
localExecution = false,
149+
localExecution,
124150
}: {
125151
tables: DBTables;
126152
root: URL;
127-
// Request module to be loaded immediately in process
128-
localExecution?: boolean;
153+
/**
154+
* Used for the execute command to import the client directly.
155+
* In other cases, we use the runtime only vite virtual module.
156+
*
157+
* This is used to ensure that the client is imported correctly
158+
* when executing commands like `astro db execute`.
159+
*/
160+
localExecution: boolean;
129161
}) {
130162
const { ASTRO_DATABASE_FILE } = getAstroEnv();
131163
const dbUrl = new URL(DB_PATH, root);
132164

133165
// If this is for the execute command, we need to import the client directly instead of using the runtime only virtual module.
134-
const clientImport = localExecution
135-
? `import { createClient } from '${DB_CLIENTS.node}';`
136-
: `import { createClient } from '${VIRTUAL_CLIENT_MODULE_ID}';`;
166+
const clientImport = getDBModule(localExecution);
137167

138168
return `
139169
import { asDrizzleTable, normalizeDatabaseUrl } from ${RUNTIME_IMPORT};
@@ -153,14 +183,20 @@ export function getRemoteVirtualModContents({
153183
appToken,
154184
isBuild,
155185
output,
156-
localExecution = false, // Used for execute command
186+
localExecution,
157187
}: {
158188
tables: DBTables;
159189
appToken: string;
160190
isBuild: boolean;
161191
output: AstroConfig['output'];
162-
// Request module to be loaded immediately in process
163-
localExecution?: boolean;
192+
/**
193+
* Used for the execute command to import the client directly.
194+
* In other cases, we use the runtime only vite virtual module.
195+
*
196+
* This is used to ensure that the client is imported correctly
197+
* when executing commands like `astro db execute`.
198+
*/
199+
localExecution: boolean;
164200
}) {
165201
const dbInfo = getRemoteDatabaseInfo();
166202

@@ -190,9 +226,7 @@ export function getRemoteVirtualModContents({
190226
}
191227

192228
// If this is for the execute command, we need to import the client directly instead of using the runtime only virtual module.
193-
const clientImport = localExecution
194-
? `import { createClient } from '${DB_CLIENTS.node}';`
195-
: `import { createClient } from '${VIRTUAL_CLIENT_MODULE_ID}';`;
229+
const clientImport = getDBModule(localExecution);
196230

197231
return `
198232
import {asDrizzleTable} from ${RUNTIME_IMPORT};

0 commit comments

Comments
 (0)