Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Validate cache.keyPath on registration #273

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 8 additions & 5 deletions control-plane/src/modules/jobs/create-job.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import {
getServiceDefinition,
parseJobArgs,
} from "../service-definitions";
import { extractWithPath } from "../util";
import { extractWithJsonPath } from "../util";
import { externalServices } from "./external";
import { env } from "../../utilities/env";
import { injectTraceContext } from "../observability/tracer";
Expand All @@ -34,12 +34,15 @@ type CreateJobParams = {

const DEFAULT_RETRY_COUNT_ON_STALL = 0;

const extractKeyFromPath = (path: string, args: unknown) => {
const extractCacheKeyFromJsonPath = (path: string, args: unknown) => {
try {
return extractWithPath(path, args)[0];
return extractWithJsonPath(path, args)[0];
} catch (error) {
if (error instanceof NotFoundError) {
throw new InvalidJobArgumentsError(error.message);
throw new InvalidJobArgumentsError(
`Failed to extract cache key from arguments: ${error.message}`,
"https://docs.inferable.ai/pages/functions#config-cache"
);
}
throw error;
}
Expand Down Expand Up @@ -94,7 +97,7 @@ export const createJob = async (params: {
};

if (config?.cache?.keyPath && config?.cache?.ttlSeconds) {
const cacheKey = extractKeyFromPath(config.cache.keyPath, args);
const cacheKey = extractCacheKeyFromJsonPath(config.cache.keyPath, args);

const { id, created } = await createJobStrategies.cached({
...jobConfig,
Expand Down
45 changes: 44 additions & 1 deletion control-plane/src/modules/service-definitions.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -340,7 +340,7 @@ describe("validateServiceRegistration", () => {
z.object({
test: z.string(),
})
))
)),
},
],
},
Expand Down Expand Up @@ -370,4 +370,47 @@ describe("validateServiceRegistration", () => {
}).not.toThrow();
})

it("should reject invalid cache.keyPath jsonpath", () => {
expect(() => {
validateServiceRegistration({
service: "default",
definition: {
name: "default",
functions: [
{
name: "myFn",
config: {
cache: {
keyPath: "$invalid",
ttlSeconds: 10
}
}
},
],
},
});
}).toThrow(InvalidServiceRegistrationError);
})

it("should accept valid cache.keyPath jsonpath", () => {
expect(() => {
validateServiceRegistration({
service: "default",
definition: {
name: "default",
functions: [
{
name: "myFn",
config: {
cache: {
keyPath: "$.someKey",
ttlSeconds: 10
}
}
},
],
},
});
}).not.toThrow();
})
})
16 changes: 13 additions & 3 deletions control-plane/src/modules/service-definitions.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import { and, eq, lte } from "drizzle-orm";
import {
handleCustomerAuthSchema,
validateDescription,
validateFunctionName,
validateFunctionSchema,
Expand All @@ -20,6 +19,7 @@ import { embeddableEntitiy } from "./embeddings/embeddings";
import { logger } from "./observability/logger";
import { packer } from "./packer";
import { withThrottle } from "./util";
import jsonpath from "jsonpath";

// The time without a ping before a service is considered expired
const SERVICE_LIVE_THRESHOLD_MS = 30 * 60 * 1000; // 30 minutes
Expand Down Expand Up @@ -400,10 +400,20 @@ export const validateServiceRegistration = ({
}
}

const VERIFY_FUNCTION_NAME = "handleCustomerAuth";
const VERIFY_FUNCTION_SERVICE = "default";
if (fn.config?.cache) {
try {
jsonpath.parse(fn.config.cache.keyPath);
} catch {
throw new InvalidServiceRegistrationError(
`${fn.name} cache.keyPath is invalid`,
"https://docs.inferable.ai/pages/functions#config-cache"
)
}
}

// Checks for customer auth handler
const VERIFY_FUNCTION_NAME = "handleCustomerAuth";
const VERIFY_FUNCTION_SERVICE = "default";
if (service === VERIFY_FUNCTION_SERVICE && fn.name === VERIFY_FUNCTION_NAME) {
if (!fn.schema) {
throw new InvalidServiceRegistrationError(
Expand Down
2 changes: 1 addition & 1 deletion control-plane/src/modules/util.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import jsonpath from "jsonpath";
import { NotFoundError } from "../utilities/errors";
import { redisClient } from "./redis";

export const extractWithPath = (path: string, args: unknown) => {
export const extractWithJsonPath = (path: string, args: unknown) => {
const result = jsonpath.query(args, path);
if (!result || result.length === 0) {
throw new NotFoundError(`Path ${path} not found within input`);
Expand Down
Loading