customTokenProvider.js

"use strict";
/**
 * Example: Custom Token Provider Implementation
 *
 * This example demonstrates how to create a custom token provider by
 * implementing the ITokenProvider interface. This gives you full control
 * over token management, including custom caching, refresh logic, and
 * error handling.
 */
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
    if (mod && mod.__esModule) return mod;
    var result = {};
    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
    __setModuleDefault(result, mod);
    return result;
};
Object.defineProperty(exports, "__esModule", { value: true });
const sql_1 = require("@databricks/sql");
const tokenProvider_1 = require("../../lib/connection/auth/tokenProvider");
/**
 * Custom token provider that refreshes tokens from a custom OAuth server.
 */
class CustomOAuthTokenProvider {
    constructor(oauthServerUrl, clientId, clientSecret) {
        this.oauthServerUrl = oauthServerUrl;
        this.clientId = clientId;
        this.clientSecret = clientSecret;
    }
    async getToken() {
        var _a;
        console.log('Fetching token from custom OAuth server...');
        // Example: Fetch token using client credentials grant
        const response = await fetch(`${this.oauthServerUrl}/oauth/token`, {
            method: 'POST',
            headers: {
                'Content-Type': 'application/x-www-form-urlencoded',
            },
            body: new URLSearchParams({
                grant_type: 'client_credentials',
                client_id: this.clientId,
                client_secret: this.clientSecret,
                scope: 'sql',
            }).toString(),
        });
        if (!response.ok) {
            throw new Error(`OAuth token request failed: ${response.status}`);
        }
        const data = await response.json();
        // Calculate expiration
        let expiresAt;
        if (typeof data.expires_in === 'number') {
            expiresAt = new Date(Date.now() + data.expires_in * 1000);
        }
        return new tokenProvider_1.Token(data.access_token, {
            tokenType: (_a = data.token_type) !== null && _a !== void 0 ? _a : 'Bearer',
            expiresAt,
        });
    }
    getName() {
        return 'CustomOAuthTokenProvider';
    }
}
/**
 * Simple token provider that reads from a file (for development/testing).
 */
class FileTokenProvider {
    constructor(filePath) {
        this.filePath = filePath;
    }
    async getToken() {
        const fs = await Promise.resolve().then(() => __importStar(require('fs/promises')));
        const tokenData = await fs.readFile(this.filePath, 'utf-8');
        const parsed = JSON.parse(tokenData);
        return tokenProvider_1.Token.fromJWT(parsed.access_token, {
            refreshToken: parsed.refresh_token,
        });
    }
    getName() {
        return 'FileTokenProvider';
    }
}
async function main() {
    const host = process.env.DATABRICKS_HOST;
    const path = process.env.DATABRICKS_HTTP_PATH;
    const client = new sql_1.DBSQLClient();
    // Option 1: Use a custom OAuth token provider
    const oauthProvider = new CustomOAuthTokenProvider(process.env.OAUTH_SERVER_URL, process.env.OAUTH_CLIENT_ID, process.env.OAUTH_CLIENT_SECRET);
    await client.connect({
        host,
        path,
        authType: 'token-provider',
        tokenProvider: oauthProvider,
        // Optionally enable federation if your OAuth server issues non-Databricks tokens
        enableTokenFederation: true,
    });
    console.log('Connected successfully with custom token provider');
    // Open a session and run a query
    const session = await client.openSession();
    const operation = await session.executeStatement('SELECT 1 AS result');
    const result = await operation.fetchAll();
    console.log('Query result:', result);
    await operation.close();
    await session.close();
    await client.close();
}
main().catch(console.error);
//# sourceMappingURL=customTokenProvider.js.map