Nvidia/BossAI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test:Add database migration service with normalization and masking functions

Browse Source 
- Implemented `normalizeMigrationInput` to standardize database migration input for sqlite, mysql, and postgres.
- Added `maskConnectionString` to obscure sensitive credentials in connection strings.
- Created comprehensive tests for migration input normalization and connection string masking.
- Established database schema creation logic for multiple dialects (sqlite, mysql, postgres).
- Developed functions for managing database transactions, including begin, commit, and rollback.
- Implemented data insertion logic from backup snapshots into the target database.
This commit is contained in:
cita
2026-03-05 17:38:52 +08:00
parent 0168592c95
commit 283c052fa5
79 changed files with 3410 additions and 1111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/server/routes/proxy/chat.ts
+6 -6
View File
@@ -139,7 +139,7 @@ async function handleChatProxyRequest(
const { requestedModel, isStream, upstreamBody, claudeOriginalBody } = parsedRequest.value!;
const downstreamPath = downstreamFormat === 'claude' ? '/v1/messages' : '/v1/chat/completions';
if (!ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
if (!await ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
const downstreamPolicy = getDownstreamRoutingPolicy(request);
const excludeChannelIds: number[] = [];
@@ -147,12 +147,12 @@ async function handleChatProxyRequest(
while (retryCount <= MAX_RETRIES) {
let selected = retryCount === 0
? tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
? await tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: await tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
if (!selected && retryCount === 0) {
await refreshModelsAndRebuildRoutes();
selected = tokenRouter.selectChannel(requestedModel, downstreamPolicy);
selected = await tokenRouter.selectChannel(requestedModel, downstreamPolicy);
}
if (!selected) {
@@ -703,7 +703,7 @@ async function handleChatProxyRequest(
}
}
function logProxy(
async function logProxy(
selected: any,
modelRequested: string,
status: string,
@@ -724,7 +724,7 @@ function logProxy(
upstreamPath,
errorMessage,
});
db.insert(schema.proxyLogs).values({
await db.insert(schema.proxyLogs).values({
routeId: selected.channel.routeId,
channelId: selected.channel.id,
accountId: selected.account.id,
src/server/routes/proxy/completions.ts
+6 -6
View File
@@ -22,7 +22,7 @@ export async function completionsProxyRoute(app: FastifyInstance) {
if (!requestedModel) {
return reply.code(400).send({ error: { message: 'model is required', type: 'invalid_request_error' } });
}
if (!ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
if (!await ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
const downstreamPolicy = getDownstreamRoutingPolicy(request);
const isStream = body.stream === true;
@@ -31,12 +31,12 @@ export async function completionsProxyRoute(app: FastifyInstance) {
while (retryCount <= MAX_RETRIES) {
let selected = retryCount === 0
? tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
? await tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: await tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
if (!selected && retryCount === 0) {
await refreshModelsAndRebuildRoutes();
selected = tokenRouter.selectChannel(requestedModel, downstreamPolicy);
selected = await tokenRouter.selectChannel(requestedModel, downstreamPolicy);
}
if (!selected) {
@@ -228,7 +228,7 @@ export async function completionsProxyRoute(app: FastifyInstance) {
});
}
function logProxy(
async function logProxy(
selected: any,
modelRequested: string,
status: string,
@@ -246,7 +246,7 @@ function logProxy(
downstreamPath: '/v1/completions',
errorMessage,
});
db.insert(schema.proxyLogs).values({
await db.insert(schema.proxyLogs).values({
routeId: selected.channel.routeId,
channelId: selected.channel.id,
accountId: selected.account.id,
src/server/routes/proxy/downstreamPolicy.ts
+4 -4
View File
@@ -7,15 +7,15 @@ export function getDownstreamRoutingPolicy(request: FastifyRequest): DownstreamR
return getProxyAuthContext(request)?.policy || EMPTY_DOWNSTREAM_ROUTING_POLICY;
}
export function ensureModelAllowedForDownstreamKey(
export async function ensureModelAllowedForDownstreamKey(
request: FastifyRequest,
reply: FastifyReply,
requestedModel: string,
): boolean {
): Promise<boolean> {
const authContext = getProxyAuthContext(request);
if (!authContext) return true;
if (isModelAllowedByPolicyOrAllowedRoutes(requestedModel, authContext.policy)) {
if (await isModelAllowedByPolicyOrAllowedRoutes(requestedModel, authContext.policy)) {
return true;
}
@@ -31,5 +31,5 @@ export function ensureModelAllowedForDownstreamKey(
export function recordDownstreamCostUsage(request: FastifyRequest, estimatedCost: number): void {
const authContext = getProxyAuthContext(request);
if (!authContext || authContext.keyId === null) return;
recordManagedKeyCostUsage(authContext.keyId, estimatedCost);
void recordManagedKeyCostUsage(authContext.keyId, estimatedCost);
}
src/server/routes/proxy/embeddings.ts
+6 -6
View File
@@ -22,7 +22,7 @@ export async function embeddingsProxyRoute(app: FastifyInstance) {
if (!requestedModel) {
return reply.code(400).send({ error: { message: 'model is required', type: 'invalid_request_error' } });
}
if (!ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
if (!await ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
const downstreamPolicy = getDownstreamRoutingPolicy(request);
const excludeChannelIds: number[] = [];
@@ -30,12 +30,12 @@ export async function embeddingsProxyRoute(app: FastifyInstance) {
while (retryCount <= MAX_RETRIES) {
let selected = retryCount === 0
? tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
? await tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: await tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
if (!selected && retryCount === 0) {
await refreshModelsAndRebuildRoutes();
selected = tokenRouter.selectChannel(requestedModel, downstreamPolicy);
selected = await tokenRouter.selectChannel(requestedModel, downstreamPolicy);
}
if (!selected) {
@@ -143,7 +143,7 @@ export async function embeddingsProxyRoute(app: FastifyInstance) {
});
}
function logProxy(
async function logProxy(
selected: any,
modelRequested: string,
status: string,
@@ -161,7 +161,7 @@ function logProxy(
downstreamPath: '/v1/embeddings',
errorMessage,
});
db.insert(schema.proxyLogs).values({
await db.insert(schema.proxyLogs).values({
routeId: selected.channel.routeId,
channelId: selected.channel.id,
accountId: selected.account.id,
src/server/routes/proxy/images.ts
+6 -6
View File
@@ -17,19 +17,19 @@ export async function imagesProxyRoute(app: FastifyInstance) {
app.post('/v1/images/generations', async (request: FastifyRequest, reply: FastifyReply) => {
const body = request.body as any;
const requestedModel = body?.model || 'gpt-image-1';
if (!ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
if (!await ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
const downstreamPolicy = getDownstreamRoutingPolicy(request);
const excludeChannelIds: number[] = [];
let retryCount = 0;
while (retryCount <= MAX_RETRIES) {
let selected = retryCount === 0
? tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
? await tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: await tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
if (!selected && retryCount === 0) {
await refreshModelsAndRebuildRoutes();
selected = tokenRouter.selectChannel(requestedModel, downstreamPolicy);
selected = await tokenRouter.selectChannel(requestedModel, downstreamPolicy);
}
if (!selected) {
@@ -116,7 +116,7 @@ export async function imagesProxyRoute(app: FastifyInstance) {
});
}
function logProxy(
async function logProxy(
selected: any,
modelRequested: string,
status: string,
@@ -131,7 +131,7 @@ function logProxy(
downstreamPath: '/v1/images/generations',
errorMessage,
});
db.insert(schema.proxyLogs).values({
await db.insert(schema.proxyLogs).values({
routeId: selected.channel.routeId,
channelId: selected.channel.id,
accountId: selected.account.id,
src/server/routes/proxy/models.ts
+13 -7
View File
@@ -9,8 +9,8 @@ export async function modelsProxyRoute(app: FastifyInstance) {
app.get('/v1/models', async (request) => {
const downstreamPolicy = getDownstreamRoutingPolicy(request);
const readModels = () => {
const rows = db.select({ modelName: schema.modelAvailability.modelName })
const readModels = async () => {
const rows = await db.select({ modelName: schema.modelAvailability.modelName })
.from(schema.modelAvailability)
.innerJoin(schema.accounts, eq(schema.modelAvailability.accountId, schema.accounts.id))
.innerJoin(schema.sites, eq(schema.accounts.siteId, schema.sites.id))
@@ -22,23 +22,29 @@ export async function modelsProxyRoute(app: FastifyInstance) {
),
)
.all();
const routeAliases = db.select({ displayName: schema.tokenRoutes.displayName })
const routeAliases = (await db.select({ displayName: schema.tokenRoutes.displayName })
.from(schema.tokenRoutes)
.where(eq(schema.tokenRoutes.enabled, true))
.all()
.all())
.map((row) => (row.displayName || '').trim())
.filter((name) => name.length > 0);
const deduped = Array.from(new Set([
...rows.map((r) => r.modelName),
...routeAliases,
])).sort();
return deduped.filter((modelName) => isModelAllowedByPolicyOrAllowedRoutes(modelName, downstreamPolicy));
const allowed: string[] = [];
for (const modelName of deduped) {
if (await isModelAllowedByPolicyOrAllowedRoutes(modelName, downstreamPolicy)) {
allowed.push(modelName);
}
}
return allowed;
};
let models = readModels();
let models = await readModels();
if (models.length === 0) {
await refreshModelsAndRebuildRoutes();
models = readModels();
models = await readModels();
}
const wantsClaudeFormat = typeof request.headers['anthropic-version'] === 'string'
src/server/routes/proxy/responses.ts
+6 -6
View File
@@ -1147,7 +1147,7 @@ export async function responsesProxyRoute(app: FastifyInstance) {
if (!requestedModel) {
return reply.code(400).send({ error: { message: 'model is required', type: 'invalid_request_error' } });
}
if (!ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
if (!await ensureModelAllowedForDownstreamKey(request, reply, requestedModel)) return;
const downstreamPolicy = getDownstreamRoutingPolicy(request);
const isStream = body.stream === true;
@@ -1156,12 +1156,12 @@ export async function responsesProxyRoute(app: FastifyInstance) {
while (retryCount <= MAX_RETRIES) {
let selected = retryCount === 0
? tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
? await tokenRouter.selectChannel(requestedModel, downstreamPolicy)
: await tokenRouter.selectNextChannel(requestedModel, excludeChannelIds, downstreamPolicy);
if (!selected && retryCount === 0) {
await refreshModelsAndRebuildRoutes();
selected = tokenRouter.selectChannel(requestedModel, downstreamPolicy);
selected = await tokenRouter.selectChannel(requestedModel, downstreamPolicy);
}
if (!selected) {
@@ -1561,7 +1561,7 @@ export async function responsesProxyRoute(app: FastifyInstance) {
});
}
function logProxy(
async function logProxy(
selected: any,
modelRequested: string,
status: string,
@@ -1582,7 +1582,7 @@ function logProxy(
upstreamPath,
errorMessage,
});
db.insert(schema.proxyLogs).values({
await db.insert(schema.proxyLogs).values({
routeId: selected.channel.routeId,
channelId: selected.channel.id,
accountId: selected.account.id,