Repository: xataio/client-ts
Pull request #147 (2022-05-06 09:41:23)

Database branch resolution
As we discussed we want to use the following logic to calculate the current branch:

- If the env variable `XATA_BRANCH` exists, use it.
- Otherwise, calculate the git branch, check if there's a database branch with the same name.
- If it exists, use it.
- If no database name exists with the same name, use `main` (until the server implements a default branch concept). This can be a dangerous if by mistake a developer connects to the production database. This will be fixed with the default/development branch concept and having scopes in API keys to prevent development API keys to have access to production. In the meantime, since we are in private beta, I think that's fine for now. Another option is to add an option to XataOptions to specify the name of the default branch. But, I prefer the logic to care care of that so we can have a constructor without any arguments: `new XataClient()`.

Getting environment variables and calculating the git branch is platform-specific code. I've implemented it for Node and Deno.


Patch:
diff --git a/.changeset/itchy-spies-thank.md b/.changeset/itchy-spies-thank.md
new file mode 100644
index 00000000..afd22f4d
--- /dev/null
+++ b/.changeset/itchy-spies-thank.md
@@ -0,0 +1,5 @@
+---
+'@xata.io/client': minor
+---
+
+Introduced automatic branch resolution mechanism
diff --git a/.github/workflows/build-pr.yml b/.github/workflows/build-pr.yml
index 40db656d..2a285d14 100644
--- a/.github/workflows/build-pr.yml
+++ b/.github/workflows/build-pr.yml
@@ -13,6 +13,13 @@ jobs:
     steps:
       - uses: actions/checkout@v3
 
+      - name: Create git branch
+        run: |
+          # Create a git branch from the detached HEAD.
+          # We need this for some integration tests.
+          # This branch is not pushed, it's just local to the workflow environment.
+          git checkout -b test-branch-${{ github.sha }}
+
       - name: Read .nvmrc
         run: echo ::set-output name=NVMRC::$(cat .nvmrc)
         id: nvm
diff --git a/client/src/api/client.ts b/client/src/api/client.ts
index 225d63c6..8f74d61f 100644
--- a/client/src/api/client.ts
+++ b/client/src/api/client.ts
@@ -1,3 +1,5 @@
+import { getAPIKey } from '../schema/config';
+import { getFetchImplementation } from '../util/fetch';
 import type * as Types from './components';
 import { operationsByTag } from './components';
 import type { FetcherExtraProps, FetchImpl } from './fetcher';
@@ -7,7 +9,7 @@ import type * as Schemas from './schemas';
 
 export interface XataApiClientOptions {
   fetch?: FetchImpl;
-  apiKey: string;
+  apiKey?: string;
   host?: HostProvider;
 }
 
@@ -15,22 +17,18 @@ export class XataApiClient {
   #extraProps: FetcherExtraProps;
 
   constructor(options: XataApiClientOptions) {
-    const globalFetch = typeof fetch !== 'undefined' ? fetch : undefined;
-    const fetchImpl = options.fetch ?? globalFetch;
-    if (!fetchImpl) {
-      /** @todo add a link after docs exist */
-      throw new Error(
-        `The \`fetch\` option passed to the Xata client is resolving to a falsy value and may not be correctly imported.`
-      );
-    }
-
     const provider = options.host ?? 'production';
+    const apiKey = options?.apiKey ?? getAPIKey();
+
+    if (!apiKey) {
+      throw new Error('Could not resolve a valid apiKey');
+    }
 
     this.#extraProps = {
       apiUrl: getHostUrl(provider, 'main'),
       workspacesApiUrl: getHostUrl(provider, 'workspaces'),
-      fetchImpl,
-      apiKey: options.apiKey
+      fetchImpl: getFetchImplementation(options.fetch),
+      apiKey
     };
   }
 
diff --git a/client/src/schema/config.ts b/client/src/schema/config.ts
new file mode 100644
index 00000000..0c1f18fd
--- /dev/null
+++ b/client/src/schema/config.ts
@@ -0,0 +1,71 @@
+import { getBranchDetails } from '../api';
+import { FetcherExtraProps } from '../api/fetcher';
+import { getEnvVariable, getGitBranch } from '../util/environment';
+import { isObject } from '../util/lang';
+
+const envBranchNames = [
+  'XATA_BRANCH',
+  'VERCEL_GIT_COMMIT_REF', // Vercel
+  'CF_PAGES_BRANCH', // Cloudflare Pages
+  'BRANCH' // Netlify. Putting it the last one because it is more ambiguous
+];
+
+export async function getBranch(fetchProps: Omit<FetcherExtraProps, 'workspacesApiUrl'>): Promise<string | undefined> {
+  const env = getBranchByEnvVariable();
+  if (env) return env;
+
+  const branch = await getGitBranch();
+  if (!branch) return undefined;
+
+  // TODO: in the future, call /resolve endpoint
+  // For now, call API to see if the branch exists. If not, use a default value.
+
+  const [protocol, , host, , database] = fetchProps.apiUrl.split('/');
+  const [workspace] = host.split('.');
+  const dbBranchName = `${database}:${branch}`;
+  try {
+    await getBranchDetails({
+      ...fetchProps,
+      workspacesApiUrl: `${protocol}//${host}`,
+      pathParams: {
+        dbBranchName,
+        workspace
+      }
+    });
+  } catch (err) {
+    if (isObject(err) && err.status === 404) return 'main';
+    throw err;
+  }
+
+  return branch;
+}
+
+function getBranchByEnvVariable(): string | undefined {
+  for (const name of envBranchNames) {
+    const value = getEnvVariable(name);
+    if (value) {
+      return value;
+    }
+  }
+  try {
+    return XATA_BRANCH;
+  } catch (err) {
+    // Ignore ReferenceError. Only CloudFlare workers set env variables as global variables
+  }
+}
+
+export function getDatabaseUrl() {
+  try {
+    return getEnvVariable('XATA_DATABASE_URL') ?? XATA_DATABASE_URL;
+  } catch (err) {
+    return undefined;
+  }
+}
+
+export function getAPIKey() {
+  try {
+    return getEnvVariable('XATA_API_KEY') ?? XATA_API_KEY;
+  } catch (err) {
+    return undefined;
+  }
+}
diff --git a/client/src/schema/repository.ts b/client/src/schema/repository.ts
index 3148bf2b..b5aa438d 100644
--- a/client/src/schema/repository.ts
+++ b/client/src/schema/repository.ts
@@ -9,7 +9,9 @@ import {
   upsertRecordWithID
 } from '../api';
 import { FetcherExtraProps, FetchImpl } from '../api/fetcher';
+import { getFetchImplementation } from '../util/fetch';
 import { isObject, isString } from '../util/lang';
+import { getAPIKey, getBranch, getDatabaseUrl } from './config';
 import { Page } from './pagination';
 import { Query } from './query';
 import { BaseData, EditableData, Identifiable, isIdentifiable, XataRecord } from './record';
@@ -144,25 +146,21 @@ export class RestRepository<Data extends BaseData, Record extends XataRecord = D
     super(null, table, {});
     this.#client = client;
     this.#table = table;
-
-    // TODO: Remove when integrating with API client
-    const globalFetch = typeof fetch !== 'undefined' ? fetch : undefined;
-    const fetchImpl = this.#client.options.fetch ?? globalFetch;
-    if (!fetchImpl) {
-      /** @todo add a link after docs exist */
-      throw new Error(
-        `The \`fetch\` option passed to the Xata client is resolving to a falsy value and may not be correctly imported.`
-      );
-    }
-    this.#fetch = fetchImpl;
+    this.#fetch = getFetchImplementation(this.#client.options.fetch);
   }
 
   async #getFetchProps(): Promise<FetcherExtraProps> {
     const branch = await this.#client.getBranch();
 
+    const apiKey = this.#client.options.apiKey ?? getAPIKey();
+
+    if (!apiKey) {
+      throw new Error('Could not resolve a valid apiKey');
+    }
+
     return {
       fetchImpl: this.#fetch,
-      apiKey: this.#client.options.apiKey,
+      apiKey,
       apiUrl: '',
       // Instead of using workspace and dbBranch, we inject a probably CNAME'd URL
       workspacesApiUrl: (path, params) => {
@@ -470,14 +468,38 @@ export type XataClientOptions = {
    */
   fetch?: FetchImpl;
   databaseURL?: string;
-  branch: BranchStrategyOption;
+  branch?: BranchStrategyOption;
   /**
    * API key to be used. You can create one in your account settings at https://app.xata.io/settings.
    */
-  apiKey: string;
+  apiKey?: string;
   repositoryFactory?: RepositoryFactory;
 };
 
+function resolveXataClientOptions(options?: Partial<XataClientOptions>): XataClientOptions {
+  const databaseURL = options?.databaseURL || getDatabaseUrl() || '';
+  const apiKey = options?.apiKey || getAPIKey() || '';
+  const branch =
+    options?.branch ||
+    (() =>
+      getBranch({
+        apiKey,
+        apiUrl: databaseURL,
+        fetchImpl: getFetchImplementation(options?.fetch)
+      }));
+
+  if (!databaseURL || !apiKey) {
+    throw new Error('Options databaseURL and apiKey are required');
+  }
+
+  return {
+    ...options,
+    databaseURL,
+    apiKey,
+    branch
+  };
+}
+
 export class BaseClient<D extends Record<string, Repository<any>> = Record<string, Repository<any>>> {
   #links: Links;
   #branch: BranchStrategyValue;
@@ -485,15 +507,13 @@ export class BaseClient<D extends Record<string, Repository<any>> = Record<strin
   options: XataClientOptions;
   public db!: D;
 
-  constructor(options: XataClientOptions, links: Links = {}) {
-    if (!options.databaseURL || !options.apiKey || !options.branch) {
-      throw new Error('Options databaseURL, apiKey and branch are required');
-    }
-
-    this.options = options;
+  constructor(options: XataClientOptions = {}, links: Links = {}) {
+    this.options = resolveXataClientOptions(options);
+    // Make this property not enumerable so it doesn't show up in console.dir, etc.
+    Object.defineProperty(this.options, 'apiKey', { enumerable: false });
     this.#links = links;
 
-    const factory = options.repositoryFactory || new RestRespositoryFactory();
+    const factory = this.options.repositoryFactory || new RestRespositoryFactory();
 
     this.db = new Proxy({} as D, {
       get: (_target, prop) => {
diff --git a/client/src/types/global-deno.d.ts b/client/src/types/global-deno.d.ts
new file mode 100644
index 00000000..e7ebf94f
--- /dev/null
+++ b/client/src/types/global-deno.d.ts
@@ -0,0 +1,7 @@
+declare namespace Deno {
+  const env: {
+    get(name: string): string | undefined;
+  };
+
+  function run(options: { cmd: string[]; stdout?: string; stderr?: string }): { output(): Promise<BufferSource> };
+}
diff --git a/client/src/types/global-node.d.ts b/client/src/types/global-node.d.ts
new file mode 100644
index 00000000..bd2d275b
--- /dev/null
+++ b/client/src/types/global-node.d.ts
@@ -0,0 +1,5 @@
+declare namespace process {
+  const env: Record<string, string>;
+}
+
+declare function require(module: string): any;
diff --git a/client/src/types/global-variables.d.ts b/client/src/types/global-variables.d.ts
new file mode 100644
index 00000000..7b314803
--- /dev/null
+++ b/client/src/types/global-variables.d.ts
@@ -0,0 +1,3 @@
+declare const XATA_DATABASE_URL: string | undefined;
+declare const XATA_API_KEY: string | undefined;
+declare const XATA_BRANCH: string | undefined;
diff --git a/client/src/util/environment.ts b/client/src/util/environment.ts
new file mode 100644
index 00000000..5a723a24
--- /dev/null
+++ b/client/src/util/environment.ts
@@ -0,0 +1,52 @@
+// eslint-disable-next-line @typescript-eslint/triple-slash-reference
+///<reference path="../types/global-node.d.ts"/>
+// eslint-disable-next-line @typescript-eslint/triple-slash-reference
+///<reference path="../types/global-variables.d.ts"/>
+// eslint-disable-next-line @typescript-eslint/triple-slash-reference
+///<reference path="../types/global-deno.d.ts"/>
+
+import { isObject, isString } from './lang';
+
+export function getEnvVariable(name: string): string | undefined {
+  // Node.js: process.env
+  try {
+    if (isObject(process) && isString(process?.env?.[name])) {
+      return process.env[name];
+    }
+  } catch (err) {
+    // Ignore: Should never happen
+  }
+
+  try {
+    // Deno: Deno.env.get
+    if (isObject(Deno) && isString(Deno?.env?.get(name))) {
+      return Deno.env.get(name);
+    }
+  } catch (err) {
+    // Ignore: Will fail if not using --allow-env
+  }
+}
+
+export async function getGitBranch(): Promise<string | undefined> {
+  // Node.js: child_process.execSync
+  try {
+    // eslint-disable-next-line @typescript-eslint/no-var-requires
+    return require('child_process').execSync('git branch --show-current', { encoding: 'utf-8' }).trim();
+  } catch (err) {
+    // Ignore
+  }
+
+  // Deno: Deno.run
+  try {
+    if (isObject(Deno)) {
+      const process = Deno.run({
+        cmd: ['git', 'branch', '--show-current'],
+        stdout: 'piped',
+        stderr: 'piped'
+      });
+      return new TextDecoder().decode(await process.output()).trim();
+    }
+  } catch (err) {
+    // Ignore: Will fail if not using --allow-run
+  }
+}
diff --git a/client/src/util/fetch.ts b/client/src/util/fetch.ts
new file mode 100644
index 00000000..7499cbc5
--- /dev/null
+++ b/client/src/util/fetch.ts
@@ -0,0 +1,13 @@
+import { FetchImpl } from '../api/fetcher';
+
+export function getFetchImplementation(userFetch?: FetchImpl) {
+  const globalFetch = typeof fetch !== 'undefined' ? fetch : undefined;
+  const fetchImpl = userFetch ?? globalFetch;
+  if (!fetchImpl) {
+    /** @todo add a link after docs exist */
+    throw new Error(
+      `The \`fetch\` option passed to the Xata client is resolving to a falsy value and may not be correctly imported.`
+    );
+  }
+  return fetchImpl;
+}
diff --git a/codegen/example/types.d.ts b/codegen/example/types.d.ts
index fe55bc8e..d14b5e5b 100644
--- a/codegen/example/types.d.ts
+++ b/codegen/example/types.d.ts
@@ -19,5 +19,5 @@ export declare class XataClient extends BaseClient<{
     "teams": Repository<Team>;
     "users": Repository<User>;
 }> {
-    constructor(options: XataClientOptions);
+    constructor(options?: XataClientOptions);
 }
diff --git a/codegen/example/xata.js b/codegen/example/xata.js
index 0bdf69c4..c675c41e 100644
--- a/codegen/example/xata.js
+++ b/codegen/example/xata.js
@@ -12,7 +12,7 @@ export class XataClient extends BaseClient {
       { databaseURL: "https://test-r5vcv5.xata.sh/db/test", ...options },
       links
     );
-    const factory = options.repositoryFactory || new RestRespositoryFactory();
+    const factory = options?.repositoryFactory || new RestRespositoryFactory();
     /** @type {{ "teams": TeamRepository; "users": UserRepository }} */
     this.db = {
       teams: factory.createRepository(this, "teams"),
diff --git a/codegen/example/xata.ts b/codegen/example/xata.ts
index 27e5c65d..634b454b 100644
--- a/codegen/example/xata.ts
+++ b/codegen/example/xata.ts
@@ -1,10 +1,4 @@
-import {
-  BaseClient,
-  Repository,
-  RestRespositoryFactory,
-  XataClientOptions,
-  XataRecord,
-} from "../../client/src";
+import { BaseClient, Repository, RestRespositoryFactory, XataClientOptions, XataRecord } from '../../client/src';
 
 export interface Team {
   name?: string | null;
@@ -23,23 +17,20 @@ export interface User {
 
 export type UserRecord = User & XataRecord;
 
-const links = { teams: [["owner", "users"]], users: [["team", "teams"]] };
+const links = { teams: [['owner', 'users']], users: [['team', 'teams']] };
 
 export class XataClient extends BaseClient<{
   teams: Repository<Team>;
   users: Repository<User>;
 }> {
-  constructor(options: XataClientOptions) {
-    super(
-      { databaseURL: "https://test-r5vcv5.xata.sh/db/test", ...options },
-      links
-    );
+  constructor(options?: XataClientOptions) {
+    super({ databaseURL: 'https://test-r5vcv5.xata.sh/db/test', ...options }, links);
 
-    const factory = options.repositoryFactory || new RestRespositoryFactory();
+    const factory = options?.repositoryFactory || new RestRespositoryFactory();
 
     this.db = {
-      teams: factory.createRepository(this, "teams"),
-      users: factory.createRepository(this, "users"),
+      teams: factory.createRepository(this, 'teams'),
+      users: factory.createRepository(this, 'users')
     };
   }
 }
diff --git a/codegen/src/codegen.ts b/codegen/src/codegen.ts
index dc874043..5c31c3b5 100644
--- a/codegen/src/codegen.ts
+++ b/codegen/src/codegen.ts
@@ -123,10 +123,10 @@ export async function generate({
     export class XataClient extends BaseClient<{
       ${tables.map((table) => `"${table.name}": Repository<${getTypeName(table.name)}>;`).join('\n')}
     }> {
-      constructor(options: XataClientOptions) {
+      constructor(options?: XataClientOptions) {
         super({ databaseURL: "https://${config.workspaceID}.xata.sh/db/${config.dbName}", ...options}, links);
-        
-        const factory = options.repositoryFactory || new RestRespositoryFactory();
+
+        const factory = options?.repositoryFactory || new RestRespositoryFactory();
         ${language === 'javascript' ? generateJSDocInternalType(tables) : ''}
         this.db = {
           ${tables.map((table) => `"${table.name}": factory.createRepository(this, "${table.name}"),`).join('\n')}


Test patch:
diff --git a/client/src/schema/index.test.ts b/client/src/schema/index.test.ts
index 87e82871..8d9f50bc 100644
--- a/client/src/schema/index.test.ts
+++ b/client/src/schema/index.test.ts
@@ -5,7 +5,7 @@ interface User extends XataRecord {
   name: string;
 }
 
-const buildClient = (options: Partial<XataClientOptions> = {}) => {
+const buildClient = (options: XataClientOptions = {}) => {
   const {
     apiKey = '1234',
     databaseURL = 'https://my-workspace-5df34do.staging.xatabase.co/db/xata',
@@ -27,14 +27,18 @@ describe('client options', () => {
   });
 
   test('throws if mandatory options are missing', () => {
-    // @ts-expect-error Options are mandatory in TypeScript
-    expect(() => buildClient({ apiKey: null }, {})).toThrow('Options databaseURL, apiKey and branch are required');
+    const { XATA_DATABASE_URL, XATA_API_KEY } = process.env;
+    process.env.XATA_API_KEY = '';
+    process.env.XATA_DATABASE_URL = '';
 
     // @ts-expect-error Options are mandatory in TypeScript
-    expect(() => buildClient({ databaseURL: null }, {})).toThrow('Options databaseURL, apiKey and branch are required');
+    expect(() => buildClient({ apiKey: null }, {})).toThrow('Options databaseURL and apiKey are required');
 
     // @ts-expect-error Options are mandatory in TypeScript
-    expect(() => buildClient({ branch: null }, {})).toThrow('Options databaseURL, apiKey and branch are required');
+    expect(() => buildClient({ databaseURL: null }, {})).toThrow('Options databaseURL and apiKey are required');
+
+    process.env.XATA_API_KEY = XATA_API_KEY;
+    process.env.XATA_DATABASE_URL = XATA_DATABASE_URL;
   });
 
   test('throws if branch cannot be resolved', async () => {
diff --git a/importer/src/processor.test.ts b/importer/src/processor.test.ts
index 31b83e6f..8f4054c9 100644
--- a/importer/src/processor.test.ts
+++ b/importer/src/processor.test.ts
@@ -152,7 +152,7 @@ const dumbTableInfo: TableInfo = {
 
 describe('createProcessor', () => {
   test('fails when the number of types and columns does not match', () => {
-    const xata = new XataApiClient({ fetch: {} as any, apiKey: '' });
+    const xata = new XataApiClient({ fetch: {} as any, apiKey: 'anything' });
     const shouldContinue = jest.fn();
 
     expect(() => createProcessor(xata, dumbTableInfo, { shouldContinue, types: [], columns: ['a'] })).toThrowError(
@@ -162,7 +162,7 @@ describe('createProcessor', () => {
   });
 
   test('needs to receive a header or receive specific column names', async () => {
-    const xata = new XataApiClient({ fetch: {} as any, apiKey: '' });
+    const xata = new XataApiClient({ fetch: {} as any, apiKey: 'anything' });
     const shouldContinue = jest.fn();
 
     const { callback } = createProcessor(xata, dumbTableInfo, { shouldContinue });
@@ -174,7 +174,7 @@ describe('createProcessor', () => {
   });
 
   test('calls shuldContinue and stops the parsing if it returns false', async () => {
-    const xata = new XataApiClient({ fetch: {} as any, apiKey: '' });
+    const xata = new XataApiClient({ fetch: {} as any, apiKey: 'anything' });
     const shouldContinue = jest.fn().mockImplementation(() => false);
 
     Object.defineProperty(xata, 'branches', {
@@ -194,7 +194,7 @@ describe('createProcessor', () => {
   });
 
   test('calls shuldContinue, continues, creates a table and inserts the records', async () => {
-    const xata = new XataApiClient({ fetch: {} as any, apiKey: '' });
+    const xata = new XataApiClient({ fetch: {} as any, apiKey: 'anything' });
     const shouldContinue = jest.fn().mockImplementation(() => true);
 
     // Mock branches API
@@ -248,7 +248,7 @@ describe('createProcessor', () => {
   });
 
   test('calls shuldContinue, continues, updates a table and inserts the records', async () => {
-    const xata = new XataApiClient({ fetch: {} as any, apiKey: '' });
+    const xata = new XataApiClient({ fetch: {} as any, apiKey: 'anything' });
     const shouldContinue = jest.fn().mockImplementation(() => true);
 
     // Mock branches API
diff --git a/test/__snapshots__/codegen.test.ts.snap b/test/__snapshots__/codegen.test.ts.snap
index 33f75edf..f4646e05 100644
--- a/test/__snapshots__/codegen.test.ts.snap
+++ b/test/__snapshots__/codegen.test.ts.snap
@@ -21,7 +21,7 @@ export class XataClient extends BaseClient {
       },
       links
     );
-    const factory = options.repositoryFactory || new RestRespositoryFactory();
+    const factory = options?.repositoryFactory || new RestRespositoryFactory();
     /** @type {{ \\"authors\\": AuthorRepository; \\"posts\\": PostRepository; \\"people\\": PersonRepository }} */
     this.db = {
       authors: factory.createRepository(this, \\"authors\\"),
@@ -82,7 +82,7 @@ export class XataClient extends BaseClient<{
   posts: Repository<Post>;
   people: Repository<Person>;
 }> {
-  constructor(options: XataClientOptions) {
+  constructor(options?: XataClientOptions) {
     super(
       {
         databaseURL: \\"https://my-workspace-gfdd12.xata.sh/db/hello-xata\\",
@@ -91,7 +91,7 @@ export class XataClient extends BaseClient<{
       links
     );
 
-    const factory = options.repositoryFactory || new RestRespositoryFactory();
+    const factory = options?.repositoryFactory || new RestRespositoryFactory();
 
     this.db = {
       authors: factory.createRepository(this, \\"authors\\"),
diff --git a/test/integration.test.ts b/test/integration.test.ts
index 1943bba4..916fd7be 100644
--- a/test/integration.test.ts
+++ b/test/integration.test.ts
@@ -2,6 +2,8 @@ import fetch from 'cross-fetch';
 import dotenv from 'dotenv';
 import { join } from 'path';
 import { BaseClient, contains, isXataRecord, lt, Repository, XataApiClient } from '../client/src';
+import { FetchImpl } from '../client/src/api/fetcher';
+import { getBranch } from '../client/src/schema/config';
 import {
   Paginable,
   PAGINATION_DEFAULT_SIZE,
@@ -10,6 +12,7 @@ import {
 } from '../client/src/schema/pagination';
 import { User, UserRecord, XataClient } from '../codegen/example/xata';
 import { mockUsers, teamColumns, userColumns } from './mock_data';
+import { execSync } from 'child_process';
 
 // Get environment variables before reading them
 dotenv.config({ path: join(process.cwd(), '.envrc') });
@@ -751,3 +754,66 @@ describe('record create or update', () => {
     expect(apiTeams[1].name).toBe('Team boats');
   });
 });
+
+describe('getBranch', () => {
+  const envValues = { ...process.env };
+  const gitBranch = execSync('git branch --show-current', { encoding: 'utf-8' }).trim();
+
+  afterAll(() => {
+    // Reset env variable values
+    process.env = envValues;
+  });
+
+  test('uses an env variable if it is set', async () => {
+    const branchName = 'using-env-variable';
+
+    const getBranchOptions = { apiKey: '', apiUrl: '', fetchImpl: {} as FetchImpl };
+
+    process.env = { XATA_BRANCH: branchName };
+    expect(await getBranch(getBranchOptions)).toEqual(branchName);
+
+    process.env = { VERCEL_GIT_COMMIT_REF: branchName };
+    expect(await getBranch(getBranchOptions)).toEqual(branchName);
+
+    process.env = { CF_PAGES_BRANCH: branchName };
+    expect(await getBranch(getBranchOptions)).toEqual(branchName);
+
+    process.env = { BRANCH: branchName };
+    expect(await getBranch(getBranchOptions)).toEqual(branchName);
+  });
+
+  test('uses the git branch if no env variable is set', async () => {
+    process.env = {};
+    const fetchImpl = jest.fn(() => ({
+      ok: true,
+      json() {
+        return { branchName: gitBranch };
+      }
+    }));
+    const branch = await getBranch({
+      apiKey: '',
+      apiUrl: 'https://workspace-id-1234.xata.sh/db/test:main',
+      fetchImpl: fetchImpl as unknown as FetchImpl
+    });
+
+    expect(branch).toEqual(gitBranch);
+  });
+
+  test('uses `main` if no env variable is set is not set and there is not associated git branch', async () => {
+    process.env = {};
+    const fetchImpl = jest.fn(() => ({
+      ok: false,
+      status: 404,
+      json() {
+        return {};
+      }
+    }));
+    const branch = await getBranch({
+      apiKey: '',
+      apiUrl: 'https://workspace-id-1234.xata.sh/db/test:main',
+      fetchImpl: fetchImpl as unknown as FetchImpl
+    });
+
+    expect(branch).toEqual('main');
+  });
+});