test:Add database migration service with normalization and masking functions
- 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:
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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'
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user