From 5cb720582254e76c3075717c031e28735f18996e Mon Sep 17 00:00:00 2001 From: Kelvin Lu Date: Tue, 15 Oct 2024 16:59:58 -0700 Subject: [PATCH] Add branch argument to chalk-ts Signed-off-by: Kelvin Lu --- src/__test__/client.test.ts | 14 +++++++++---- src/_client.ts | 42 ++++++++++++++++++++++++++++--------- src/_http.ts | 5 +++++ src/_types.ts | 3 ++- 4 files changed, 49 insertions(+), 15 deletions(-) diff --git a/src/__test__/client.test.ts b/src/__test__/client.test.ts index e6eb56a..2148827 100644 --- a/src/__test__/client.test.ts +++ b/src/__test__/client.test.ts @@ -34,16 +34,18 @@ describe("ChalkClient", () => { const client = new ChalkClient({ activeEnvironment: "a", apiServer: "b", - clientId: "c", - clientSecret: "d", + branch: "c", + clientId: "d", + clientSecret: "e", timestampFormat: TimestampFormat.EPOCH_MILLIS, }); expect(getConfig(client)).toEqual({ activeEnvironment: "a", apiServer: "b", - clientId: "c", - clientSecret: "d", + branch: "c", + clientId: "d", + clientSecret: "e", queryServer: "b", timestampFormat: TimestampFormat.EPOCH_MILLIS, }); @@ -52,6 +54,7 @@ describe("ChalkClient", () => { it("reads from environment variables when set", () => { process.env._CHALK_ACTIVE_ENVIRONMENT = "env"; process.env._CHALK_API_SERVER = "http://localhost:8000"; + process.env._CHALK_BRANCH = "not_a_real_branch"; process.env._CHALK_CLIENT_ID = "client_id"; process.env._CHALK_CLIENT_SECRET = "secret"; process.env._CHALK_QUERY_SERVER = "http://localhost:1337"; @@ -61,6 +64,7 @@ describe("ChalkClient", () => { expect(getConfig(client)).toEqual({ activeEnvironment: "env", apiServer: "http://localhost:8000", + branch: "not_a_real_branch", clientId: "client_id", clientSecret: "secret", queryServer: "http://localhost:1337", @@ -84,6 +88,7 @@ describe("ChalkClient", () => { expect(getConfig(client)).toEqual({ activeEnvironment: undefined, apiServer: DEFAULT_API_SERVER, + branch: undefined, clientId: "client_id", clientSecret: "secret", queryServer: DEFAULT_API_SERVER, @@ -99,6 +104,7 @@ describe("ChalkClient", () => { expect(getConfig(client)).toEqual({ activeEnvironment: undefined, apiServer: DEFAULT_API_SERVER, + branch: undefined, clientId: "client_id", clientSecret: "secret", queryServer: "query server", diff --git a/src/_client.ts b/src/_client.ts index 3cc008d..99a2926 100644 --- a/src/_client.ts +++ b/src/_client.ts @@ -56,12 +56,21 @@ export interface ChalkClientOpts { queryServer?: string; /** - * The environment that your client will run against. This value will be read from the _CHALK_ACTIVE_ENVIRONMENT environment variable if not set explicitly. + * The environment that your client will run against. + * This value will be read from the _CHALK_ACTIVE_ENVIRONMENT environment variable if not set explicitly. * * If not specified and unset by your environment, an error will be thrown on client creation */ activeEnvironment?: string; + /** + * If specified, Chalk will route all requests from this client instance to the relevant branch. + * This value will be read from the _CHALK_BRANCH environment variable if not set explicitly. + * + * Some methods allow you to override this instance-level branch configuration by passing in a `branch` argument. + */ + branch?: string; + /** * A custom fetch client that will replace the fetch polyfill used by default. * @@ -106,6 +115,11 @@ function valueWithEnvFallback( } export interface ChalkRequestOptions { + /** + * If specified, Chalk will route this request to the relevant branch. Overrides the branch passed in to the + * client initialization. + */ + branch?: string; /** * The timeout for the request in milliseconds. If not provided, the client will use the default timeout * specified at the client level. @@ -131,6 +145,7 @@ export class ChalkClient> process.env._CHALK_ACTIVE_ENVIRONMENT ?? undefined, apiServer: resolvedApiServer, + branch: opts?.branch ?? process.env._CHALK_BRANCH ?? undefined, clientId: valueWithEnvFallback( "clientId", opts?.clientId, @@ -157,7 +172,7 @@ export class ChalkClient> async whoami(): Promise { return this.http.v1_who_am_i({ baseUrl: this.config.apiServer, - headers: this.getDefaultHeaders(), + headers: this.getHeaders(), credentials: this.credentials, }); } @@ -168,7 +183,7 @@ export class ChalkClient> pathParams: { run_id: runId, }, - headers: this.getDefaultHeaders(), + headers: this.getHeaders(), credentials: this.credentials, }); } @@ -181,7 +196,7 @@ export class ChalkClient> body: { resolver_fqn: request.resolverFqn, }, - headers: this.getDefaultHeaders(), + headers: this.getHeaders(), credentials: this.credentials, }); } @@ -213,7 +228,7 @@ export class ChalkClient> include_meta: !!request.include_meta, planner_options: request.plannerOptions, }, - headers: this.getDefaultHeaders(), + headers: this.getHeaders(requestOptions), credentials: this.credentials, timeout: requestOptions?.timeout, }); @@ -250,7 +265,7 @@ export class ChalkClient> const rawResult = await this.http.v1_query_feather({ baseUrl: this.config.queryServer, body: requestBuffer.buffer, - headers: this.getDefaultHeaders(), + headers: this.getHeaders(requestOptions), credentials: this.credentials, timeout: requestOptions?.timeout, }); @@ -286,7 +301,7 @@ export class ChalkClient> const rawResult = await this.http.v1_query_feather({ baseUrl: this.config.queryServer, body: requestBuffer.buffer, - headers: this.getDefaultHeaders(), + headers: this.getHeaders(requestOptions), credentials: this.credentials, timeout: requestOptions?.timeout, }); @@ -312,7 +327,7 @@ export class ChalkClient> correlation_id: request.correlationId, deployment_id: request.previewDeploymentId, }, - headers: this.getDefaultHeaders(), + headers: this.getHeaders(), credentials: this.credentials, }); @@ -324,10 +339,17 @@ export class ChalkClient> } } - private getDefaultHeaders(): ChalkHttpHeaders { - return { + private getHeaders(requestOptions?: ChalkRequestOptions): ChalkHttpHeaders { + const headers: ChalkHttpHeaders = { "X-Chalk-Env-Id": this.config.activeEnvironment, "User-Agent": "chalk-ts v1.17.0", }; + + const branch = requestOptions?.branch ?? this.config.branch; + if (branch != null) { + headers["X-Chalk-Branch-Id"] = branch; + } + + return headers; } } diff --git a/src/_http.ts b/src/_http.ts index f58bb58..409429a 100644 --- a/src/_http.ts +++ b/src/_http.ts @@ -7,6 +7,7 @@ import { ChalkErrorCategory, ChalkErrorCode } from "./_interface"; export interface ChalkHttpHeaders { "X-Chalk-Env-Id"?: string; + "X-Chalk-Branch-Id"?: string; "User-Agent"?: string; } @@ -239,6 +240,10 @@ export class ChalkHTTPService { headers.set("X-Chalk-Timeout", effectiveTimeout.toString()); } + if (callArgs.headers?.["X-Chalk-Branch-Id"] != null) { + headers.set("X-Chalk-Branch-Id", callArgs.headers["X-Chalk-Branch-Id"]); + } + const body = callArgs.body !== undefined ? !opts.binaryResponseBody diff --git a/src/_types.ts b/src/_types.ts index 9103cd8..123f005 100644 --- a/src/_types.ts +++ b/src/_types.ts @@ -1,9 +1,9 @@ import { TimestampFormat } from "./_interface"; - export interface ChalkClientConfig { activeEnvironment: string | undefined; apiServer: string; + branch: string | undefined; clientId: string; clientSecret: string; queryServer: string; @@ -24,6 +24,7 @@ export interface ChalkEnvironmentVariables { _CHALK_API_SERVER: string; _CHALK_QUERY_SERVER: string; _CHALK_ACTIVE_ENVIRONMENT: string; + _CHALK_BRANCH: string; } /**