import axios, { AxiosInstance, AxiosError, AxiosRequestConfig } from 'axios';
import type * as types from './types/databricks.js';
export class DatabricksClient {
private client: AxiosInstance;
private config: types.DatabricksConfig;
constructor(config: types.DatabricksConfig) {
this.config = config;
// Validate configuration
if (!config.host) {
throw new Error('DATABRICKS_HOST is required');
}
if (!config.token && !config.clientId) {
throw new Error('Either DATABRICKS_TOKEN or DATABRICKS_CLIENT_ID/SECRET is required');
}
// Create axios instance with base configuration
this.client = axios.create({
baseURL: config.host.replace(/\/$/, ''), // Remove trailing slash
headers: {
'Content-Type': 'application/json',
},
timeout: 60000, // 60 second timeout
});
// Add authentication
if (config.token) {
this.client.defaults.headers.common['Authorization'] = `Bearer ${config.token}`;
}
// Add request/response interceptors for error handling
this.setupInterceptors();
}
private setupInterceptors(): void {
// Request interceptor for logging
this.client.interceptors.request.use(
(config) => {
return config;
},
(error) => {
return Promise.reject(error);
}
);
// Response interceptor for error handling
this.client.interceptors.response.use(
(response) => response,
async (error: AxiosError) => {
if (error.response) {
const status = error.response.status;
const data = error.response.data as types.ErrorResponse;
// Handle specific error codes
if (status === 401) {
throw new Error('Authentication failed. Please check your credentials.');
} else if (status === 403) {
throw new Error('Permission denied. You do not have access to this resource.');
} else if (status === 404) {
throw new Error('Resource not found.');
} else if (status === 429) {
// Rate limiting - implement exponential backoff
const retryAfter = error.response.headers['retry-after'];
const delay = retryAfter ? parseInt(retryAfter) * 1000 : 5000;
await new Promise(resolve => setTimeout(resolve, delay));
// Retry the request
return this.client.request(error.config as AxiosRequestConfig);
} else if (status >= 500) {
throw new Error(`Databricks service error: ${data?.message || 'Internal server error'}`);
}
throw new Error(`API Error: ${data?.message || error.message}`);
} else if (error.request) {
throw new Error('No response from Databricks API. Please check your network connection.');
} else {
throw new Error(`Request error: ${error.message}`);
}
}
);
}
// Cluster Management APIs
async listClusters(): Promise<types.ClusterInfo[]> {
const response = await this.client.get('/api/2.0/clusters/list');
return response.data.clusters || [];
}
async getCluster(clusterId: string): Promise<types.ClusterInfo> {
const response = await this.client.get('/api/2.0/clusters/get', {
params: { cluster_id: clusterId },
});
return response.data;
}
async createCluster(config: types.CreateClusterRequest): Promise<{ cluster_id: string }> {
const response = await this.client.post('/api/2.0/clusters/create', config);
return response.data;
}
async startCluster(clusterId: string): Promise<void> {
await this.client.post('/api/2.0/clusters/start', { cluster_id: clusterId });
}
async restartCluster(clusterId: string): Promise<void> {
await this.client.post('/api/2.0/clusters/restart', { cluster_id: clusterId });
}
async terminateCluster(clusterId: string): Promise<void> {
await this.client.post('/api/2.0/clusters/delete', { cluster_id: clusterId });
}
async editCluster(clusterId: string, config: Partial<types.CreateClusterRequest>): Promise<void> {
await this.client.post('/api/2.0/clusters/edit', {
cluster_id: clusterId,
...config,
});
}
async resizeCluster(
clusterId: string,
numWorkers?: number,
autoscale?: types.AutoScale
): Promise<void> {
await this.client.post('/api/2.0/clusters/resize', {
cluster_id: clusterId,
num_workers: numWorkers,
autoscale,
});
}
async getClusterEvents(
clusterId: string,
startTime?: number,
endTime?: number,
limit?: number
): Promise<types.ClusterEvent[]> {
const response = await this.client.post('/api/2.0/clusters/events', {
cluster_id: clusterId,
start_time: startTime,
end_time: endTime,
limit: limit || 50,
});
return response.data.events || [];
}
async pinCluster(clusterId: string): Promise<void> {
await this.client.post('/api/2.0/clusters/pin', { cluster_id: clusterId });
}
async unpinCluster(clusterId: string): Promise<void> {
await this.client.post('/api/2.0/clusters/unpin', { cluster_id: clusterId });
}
// Workspace/Notebook Management APIs
async listNotebooks(path: string): Promise<types.NotebookInfo[]> {
const response = await this.client.get('/api/2.0/workspace/list', {
params: { path },
});
return response.data.objects || [];
}
async getNotebook(
path: string,
format: types.NotebookFormat = 'SOURCE'
): Promise<types.NotebookContent> {
const response = await this.client.get('/api/2.0/workspace/export', {
params: { path, format },
});
return {
path,
language: response.data.language,
content: response.data.content,
format,
};
}
async createNotebook(request: types.CreateNotebookRequest): Promise<void> {
await this.client.post('/api/2.0/workspace/import', {
path: request.path,
language: request.language || 'PYTHON',
content: request.content ? Buffer.from(request.content).toString('base64') : '',
format: request.format || 'SOURCE',
overwrite: false,
});
}
async updateNotebook(path: string, content: string, format: types.NotebookFormat = 'SOURCE'): Promise<void> {
await this.client.post('/api/2.0/workspace/import', {
path,
content: Buffer.from(content).toString('base64'),
format,
overwrite: true,
});
}
async deleteNotebook(path: string, recursive: boolean = false): Promise<void> {
await this.client.post('/api/2.0/workspace/delete', {
path,
recursive,
});
}
async moveNotebook(sourcePath: string, destinationPath: string): Promise<void> {
await this.client.post('/api/2.0/workspace/move', {
source_path: sourcePath,
destination_path: destinationPath,
});
}
async mkdirs(path: string): Promise<void> {
await this.client.post('/api/2.0/workspace/mkdirs', { path });
}
// Jobs API
async listJobs(limit?: number, offset?: number): Promise<types.JobInfo[]> {
const response = await this.client.get('/api/2.1/jobs/list', {
params: { limit: limit || 25, offset: offset || 0 },
});
return response.data.jobs || [];
}
async getJob(jobId: number): Promise<types.JobInfo> {
const response = await this.client.get('/api/2.1/jobs/get', {
params: { job_id: jobId },
});
return response.data;
}
async createJob(settings: types.JobSettings): Promise<{ job_id: number }> {
const response = await this.client.post('/api/2.1/jobs/create', settings);
return response.data;
}
async updateJob(jobId: number, newSettings: Partial<types.JobSettings>): Promise<void> {
await this.client.post('/api/2.1/jobs/update', {
job_id: jobId,
new_settings: newSettings,
});
}
async deleteJob(jobId: number): Promise<void> {
await this.client.post('/api/2.1/jobs/delete', { job_id: jobId });
}
async runJobNow(
jobId: number,
notebookParams?: Record<string, string>,
jarParams?: string[],
pythonParams?: string[]
): Promise<{ run_id: number }> {
const response = await this.client.post('/api/2.1/jobs/run-now', {
job_id: jobId,
notebook_params: notebookParams,
jar_params: jarParams,
python_params: pythonParams,
});
return response.data;
}
async getJobRun(runId: number): Promise<types.JobRun> {
const response = await this.client.get('/api/2.1/jobs/runs/get', {
params: { run_id: runId },
});
return response.data;
}
async listJobRuns(
jobId?: number,
limit?: number,
offset?: number
): Promise<types.JobRun[]> {
const response = await this.client.get('/api/2.1/jobs/runs/list', {
params: {
job_id: jobId,
limit: limit || 25,
offset: offset || 0,
},
});
return response.data.runs || [];
}
async cancelJobRun(runId: number): Promise<void> {
await this.client.post('/api/2.1/jobs/runs/cancel', { run_id: runId });
}
async getJobRunOutput(runId: number): Promise<types.NotebookOutput> {
const response = await this.client.get('/api/2.1/jobs/runs/get-output', {
params: { run_id: runId },
});
return response.data;
}
// User and Group Management APIs (SCIM)
async listUsers(): Promise<types.User[]> {
const response = await this.client.get('/api/2.0/preview/scim/v2/Users');
return response.data.Resources || [];
}
async getUser(userId: string): Promise<types.User> {
const response = await this.client.get(`/api/2.0/preview/scim/v2/Users/${userId}`);
return response.data;
}
async createUser(user: Partial<types.User>): Promise<types.User> {
const response = await this.client.post('/api/2.0/preview/scim/v2/Users', user);
return response.data;
}
async deleteUser(userId: string): Promise<void> {
await this.client.delete(`/api/2.0/preview/scim/v2/Users/${userId}`);
}
async listGroups(): Promise<types.Group[]> {
const response = await this.client.get('/api/2.0/preview/scim/v2/Groups');
return response.data.Resources || [];
}
async getGroup(groupId: string): Promise<types.Group> {
const response = await this.client.get(`/api/2.0/preview/scim/v2/Groups/${groupId}`);
return response.data;
}
async createGroup(group: Partial<types.Group>): Promise<types.Group> {
const response = await this.client.post('/api/2.0/preview/scim/v2/Groups', group);
return response.data;
}
async deleteGroup(groupId: string): Promise<void> {
await this.client.delete(`/api/2.0/preview/scim/v2/Groups/${groupId}`);
}
async addUserToGroup(groupId: string, userId: string): Promise<void> {
await this.client.patch(`/api/2.0/preview/scim/v2/Groups/${groupId}`, {
schemas: ['urn:ietf:params:scim:api:messages:2.0:PatchOp'],
Operations: [
{
op: 'add',
path: 'members',
value: [{ value: userId }],
},
],
});
}
async removeUserFromGroup(groupId: string, userId: string): Promise<void> {
await this.client.patch(`/api/2.0/preview/scim/v2/Groups/${groupId}`, {
schemas: ['urn:ietf:params:scim:api:messages:2.0:PatchOp'],
Operations: [
{
op: 'remove',
path: `members[value eq "${userId}"]`,
},
],
});
}
async listServicePrincipals(): Promise<types.ServicePrincipal[]> {
const response = await this.client.get('/api/2.0/preview/scim/v2/ServicePrincipals');
return response.data.Resources || [];
}
// Unity Catalog APIs
async listCatalogs(): Promise<types.CatalogInfo[]> {
const response = await this.client.get('/api/2.1/unity-catalog/catalogs');
return response.data.catalogs || [];
}
async getCatalog(catalogName: string): Promise<types.CatalogInfo> {
const response = await this.client.get(`/api/2.1/unity-catalog/catalogs/${catalogName}`);
return response.data;
}
async createCatalog(catalog: Partial<types.CatalogInfo>): Promise<types.CatalogInfo> {
const response = await this.client.post('/api/2.1/unity-catalog/catalogs', catalog);
return response.data;
}
async deleteCatalog(catalogName: string, force: boolean = false): Promise<void> {
await this.client.delete(`/api/2.1/unity-catalog/catalogs/${catalogName}`, {
params: { force },
});
}
async listSchemas(catalogName: string): Promise<types.SchemaInfo[]> {
const response = await this.client.get('/api/2.1/unity-catalog/schemas', {
params: { catalog_name: catalogName },
});
return response.data.schemas || [];
}
async getSchema(fullSchemaName: string): Promise<types.SchemaInfo> {
const response = await this.client.get(`/api/2.1/unity-catalog/schemas/${fullSchemaName}`);
return response.data;
}
async createSchema(schema: Partial<types.SchemaInfo>): Promise<types.SchemaInfo> {
const response = await this.client.post('/api/2.1/unity-catalog/schemas', schema);
return response.data;
}
async deleteSchema(fullSchemaName: string, force: boolean = false): Promise<void> {
await this.client.delete(`/api/2.1/unity-catalog/schemas/${fullSchemaName}`, {
params: { force },
});
}
async listTables(catalogName: string, schemaName: string): Promise<types.TableInfo[]> {
const response = await this.client.get('/api/2.1/unity-catalog/tables', {
params: {
catalog_name: catalogName,
schema_name: schemaName,
},
});
return response.data.tables || [];
}
async getTable(fullTableName: string): Promise<types.TableInfo> {
const response = await this.client.get(`/api/2.1/unity-catalog/tables/${fullTableName}`);
return response.data;
}
async deleteTable(fullTableName: string): Promise<void> {
await this.client.delete(`/api/2.1/unity-catalog/tables/${fullTableName}`);
}
async listVolumes(catalogName: string, schemaName: string): Promise<types.VolumeInfo[]> {
const response = await this.client.get('/api/2.1/unity-catalog/volumes', {
params: {
catalog_name: catalogName,
schema_name: schemaName,
},
});
return response.data.volumes || [];
}
async listExternalLocations(): Promise<types.ExternalLocationInfo[]> {
const response = await this.client.get('/api/2.1/unity-catalog/external-locations');
return response.data.external_locations || [];
}
// Permissions APIs
async getPermissions(objectType: string, objectId: string): Promise<types.AccessControlList> {
const response = await this.client.get(`/api/2.0/permissions/${objectType}/${objectId}`);
return response.data;
}
async setPermissions(
objectType: string,
objectId: string,
accessControlList: types.AccessControl[]
): Promise<void> {
await this.client.put(`/api/2.0/permissions/${objectType}/${objectId}`, {
access_control_list: accessControlList,
});
}
async updatePermissions(
objectType: string,
objectId: string,
accessControlList: types.AccessControl[]
): Promise<void> {
await this.client.patch(`/api/2.0/permissions/${objectType}/${objectId}`, {
access_control_list: accessControlList,
});
}
async getPermissionLevels(objectType: string, objectId: string): Promise<types.PermissionLevel[]> {
const response = await this.client.get(`/api/2.0/permissions/${objectType}/${objectId}/permissionLevels`);
return response.data.permission_levels || [];
}
// Unity Catalog Grants
async grantPermissions(
securableType: string,
securableName: string,
principal: string,
privileges: types.Privilege[]
): Promise<void> {
await this.client.post('/api/2.1/unity-catalog/permissions', {
securable_type: securableType,
securable_name: securableName,
principal,
changes: privileges.map(p => ({ add: [p] })),
});
}
async revokePermissions(
securableType: string,
securableName: string,
principal: string,
privileges: types.Privilege[]
): Promise<void> {
await this.client.post('/api/2.1/unity-catalog/permissions', {
securable_type: securableType,
securable_name: securableName,
principal,
changes: privileges.map(p => ({ remove: [p] })),
});
}
async getEffectivePermissions(
securableType: string,
securableName: string,
principal: string
): Promise<types.Privilege[]> {
const response = await this.client.get('/api/2.1/unity-catalog/effective-permissions', {
params: {
securable_type: securableType,
securable_name: securableName,
principal,
},
});
return response.data.privilege_assignments || [];
}
// SQL Execution (for querying tables)
async executeSqlStatement(
warehouseId: string,
statement: string,
catalog?: string,
schema?: string
): Promise<any> {
const response = await this.client.post('/api/2.0/sql/statements', {
warehouse_id: warehouseId,
statement,
catalog,
schema,
wait_timeout: '30s',
});
return response.data;
}
async getSqlStatementResult(statementId: string): Promise<any> {
const response = await this.client.get(`/api/2.0/sql/statements/${statementId}`);
return response.data;
}
// Custom method for making arbitrary API calls
async apiCall(method: string, endpoint: string, data?: any, params?: any): Promise<any> {
const response = await this.client.request({
method,
url: endpoint,
data,
params,
});
return response.data;
}
}
