ydb-qdrant 7.0.1 → 8.1.0

package/dist/ydb/client.js

@@ -1,194 +1,50 @@
-import { Driver } from "@ydbjs/core";
-import { query } from "@ydbjs/query";
-import { CredentialsProvider } from "@ydbjs/auth";
-import { AnonymousCredentialsProvider } from "@ydbjs/auth/anonymous";
-import { AccessTokenCredentialsProvider } from "@ydbjs/auth/access-token";
-import { MetadataCredentialsProvider } from "@ydbjs/auth/metadata";
-import { readFile } from "node:fs/promises";
-import crypto from "node:crypto";
-import path from "node:path";
-import { createRequire } from "node:module";
-import { pathToFileURL } from "node:url";
-import { YDB_DATABASE, YDB_ENDPOINT, STARTUP_PROBE_SESSION_TIMEOUT_MS, } from "../config/env.js";
+import { createRequire } from "module";
+import { YDB_DATABASE, YDB_ENDPOINT, SESSION_POOL_MIN_SIZE, SESSION_POOL_MAX_SIZE, SESSION_KEEPALIVE_PERIOD_MS, STARTUP_PROBE_SESSION_TIMEOUT_MS, } from "../config/env.js";
 import { logger } from "../logging/logger.js";
-import { SessionPool } from "./SessionPool.js";
+const require = createRequire(import.meta.url);
+const { Driver, getCredentialsFromEnv, Types, TypedValues, TableDescription, Column, ExecuteQuerySettings, BulkUpsertSettings, OperationParams, Ydb, } = require("ydb-sdk");
+export { Types, TypedValues, TableDescription, Column, ExecuteQuerySettings, BulkUpsertSettings, Ydb, };
+export function createExecuteQuerySettings(options) {
+    const { keepInCache = true, idempotent = true } = options ?? {};
+    const settings = new ExecuteQuerySettings();
+    if (keepInCache) {
+        settings.withKeepInCache(true);
+    }
+    if (idempotent) {
+        settings.withIdempotent(true);
+    }
+    return settings;
+}
+export function createExecuteQuerySettingsWithTimeout(options) {
+    const settings = createExecuteQuerySettings(options);
+    const op = new OperationParams();
+    const seconds = Math.max(1, Math.ceil(options.timeoutMs / 1000));
+    // Limit both overall operation processing time and cancellation time on the
+    // server side so the probe fails fast instead of hanging for the default.
+    op.withOperationTimeoutSeconds(seconds);
+    op.withCancelAfterSeconds(seconds);
+    settings.withOperationParams(op);
+    return settings;
+}
+export function createBulkUpsertSettingsWithTimeout(options) {
+    const settings = new BulkUpsertSettings();
+    const op = new OperationParams();
+    const seconds = Math.max(1, Math.ceil(options.timeoutMs / 1000));
+    op.withOperationTimeoutSeconds(seconds);
+    op.withCancelAfterSeconds(seconds);
+    settings.withOperationParams(op);
+    return settings;
+}
 const DRIVER_READY_TIMEOUT_MS = 15000;
+const TABLE_SESSION_TIMEOUT_MS = 20000;
 const YDB_HEALTHCHECK_READY_TIMEOUT_MS = 5000;
 const DRIVER_REFRESH_COOLDOWN_MS = 30000;
-// Bound `withSession()` so `d.ready()` / session acquire can't hang forever under YDB stalls.
-// Matches historical table client session timeout behavior (~20s).
-const WITH_SESSION_TIMEOUT_MS = 20000;
 let overrideConfig;
 let driver;
-let sqlClient;
-let sessionPool;
 let lastDriverRefreshAt = 0;
 let driverRefreshInFlight = null;
-// Test-only: allows injecting a mock Driver factory.
+// Test-only: allows injecting a mock Driver factory
 let driverFactoryOverride;
-let queryCtxPromise = null;
-async function getQueryCtx() {
-    if (queryCtxPromise) {
-        return await queryCtxPromise;
-    }
-    const require = createRequire(import.meta.url);
-    const entry = require.resolve("@ydbjs/query");
-    const ctxPath = path.join(path.dirname(entry), "ctx.js");
-    queryCtxPromise = import(pathToFileURL(ctxPath).href);
-    return await queryCtxPromise;
-}
-const IAM_TOKEN_URL = "https://iam.api.cloud.yandex.net/iam/v1/tokens";
-const SA_JWT_MAX_AGE_SECONDS = 3600;
-const SA_TOKEN_REFRESH_SKEW_MS = 60_000;
-// Yandex Cloud IAM tokens: lifetime does not exceed 12 hours, but the docs recommend
-// requesting a token more often (e.g., every hour). We refresh a bit earlier as a
-// best-effort fallback when the API response doesn't include expiresAt.
-// Source: https://cloud.yandex.com/en/docs/iam/operations/iam-token/create-for-sa
-const SA_TOKEN_REFRESH_FALLBACK_MS = 55 * 60_000;
-function parseTruthyEnv(value) {
-    if (value === undefined)
-        return false;
-    const normalized = value.trim().toLowerCase();
-    if (normalized === "" ||
-        normalized === "0" ||
-        normalized === "false" ||
-        normalized === "no" ||
-        normalized === "off") {
-        return false;
-    }
-    return true;
-}
-function toBase64Url(input) {
-    const buf = Buffer.isBuffer(input) ? input : Buffer.from(input);
-    return buf
-        .toString("base64")
-        .replace(/\+/g, "-")
-        .replace(/\//g, "_")
-        .replace(/=+$/g, "");
-}
-function jsonToBase64Url(obj) {
-    return toBase64Url(Buffer.from(JSON.stringify(obj), "utf8"));
-}
-function normalizeServiceAccountPrivateKeyPem(raw) {
-    const trimmed = raw.trim();
-    if (trimmed.startsWith("PLEASE DO NOT REMOVE THIS LINE!")) {
-        return trimmed.split("\n").slice(1).join("\n").trim();
-    }
-    return trimmed;
-}
-async function readServiceAccountKeyFile(absPath) {
-    const raw = await readFile(absPath, "utf8");
-    const parsed = JSON.parse(raw);
-    if (typeof parsed.id !== "string" ||
-        typeof parsed.service_account_id !== "string" ||
-        typeof parsed.private_key !== "string") {
-        throw new Error("Invalid service account key JSON: expected fields id, service_account_id, private_key");
-    }
-    return {
-        id: parsed.id,
-        service_account_id: parsed.service_account_id,
-        private_key: parsed.private_key,
-    };
-}
-function createServiceAccountJwt(args) {
-    // Per Yandex Cloud IAM docs, the service account JWT must be PS256 with kid
-    // and audience set to the IAM token endpoint.
-    const header = { typ: "JWT", alg: "PS256", kid: args.keyId };
-    const payload = {
-        iss: args.serviceAccountId,
-        aud: IAM_TOKEN_URL,
-        iat: args.nowSec,
-        exp: args.nowSec + SA_JWT_MAX_AGE_SECONDS,
-    };
-    const encodedHeader = jsonToBase64Url(header);
-    const encodedPayload = jsonToBase64Url(payload);
-    const signingInput = `${encodedHeader}.${encodedPayload}`;
-    const signature = crypto.sign("RSA-SHA256", Buffer.from(signingInput), {
-        key: normalizeServiceAccountPrivateKeyPem(args.privateKeyPem),
-        padding: crypto.constants.RSA_PKCS1_PSS_PADDING,
-        saltLength: 32,
-    });
-    return `${signingInput}.${toBase64Url(signature)}`;
-}
-async function exchangeJwtForIamToken(args) {
-    const res = await fetch(IAM_TOKEN_URL, {
-        method: "POST",
-        headers: { "content-type": "application/json" },
-        body: JSON.stringify({ jwt: args.jwt }),
-        signal: args.signal,
-    });
-    if (!res.ok) {
-        const text = await res.text().catch(() => "");
-        throw new Error(`Failed to exchange JWT for IAM token: HTTP ${res.status}${text ? `: ${text}` : ""}`);
-    }
-    const body = (await res.json());
-    const token = body.iamToken;
-    if (typeof token !== "string" || token.length === 0) {
-        throw new Error("IAM token exchange response is missing iamToken");
-    }
-    const expiresAtRaw = body.expiresAt ?? body.expires_at;
-    const expiresAtMs = typeof expiresAtRaw === "string" ? Date.parse(expiresAtRaw) : NaN;
-    return {
-        iamToken: token,
-        expiresAtMs: Number.isFinite(expiresAtMs) ? expiresAtMs : null,
-    };
-}
-class ServiceAccountKeyFileCredentialsProvider extends CredentialsProvider {
-    keyFilePath;
-    cached = null;
-    constructor(keyFilePath) {
-        super();
-        this.keyFilePath = keyFilePath;
-    }
-    async getToken(force = false, signal) {
-        const nowMs = Date.now();
-        if (!force && this.cached) {
-            const { token, expiresAtMs, cachedAtMs } = this.cached;
-            if (expiresAtMs !== null) {
-                if (nowMs + SA_TOKEN_REFRESH_SKEW_MS < expiresAtMs) {
-                    return token;
-                }
-            }
-            else {
-                // If the API doesn't provide expiresAt, refresh hourly (best-effort).
-                if (nowMs - cachedAtMs < SA_TOKEN_REFRESH_FALLBACK_MS) {
-                    return token;
-                }
-            }
-        }
-        const key = await readServiceAccountKeyFile(this.keyFilePath);
-        const nowSec = Math.floor(nowMs / 1000);
-        const jwt = createServiceAccountJwt({
-            keyId: key.id,
-            serviceAccountId: key.service_account_id,
-            privateKeyPem: key.private_key,
-            nowSec,
-        });
-        const { iamToken, expiresAtMs } = await exchangeJwtForIamToken({
-            jwt,
-            signal,
-        });
-        this.cached = { token: iamToken, expiresAtMs, cachedAtMs: nowMs };
-        return iamToken;
-    }
-}
-function resolveCredentialsProviderFromEnv() {
-    const saKeyPath = process.env.YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS?.trim();
-    if (saKeyPath) {
-        return new ServiceAccountKeyFileCredentialsProvider(saKeyPath);
-    }
-    if (parseTruthyEnv(process.env.YDB_METADATA_CREDENTIALS)) {
-        return new MetadataCredentialsProvider();
-    }
-    const accessToken = process.env.YDB_ACCESS_TOKEN_CREDENTIALS?.trim();
-    if (accessToken) {
-        return new AccessTokenCredentialsProvider({ token: accessToken });
-    }
-    if (parseTruthyEnv(process.env.YDB_ANONYMOUS_CREDENTIALS)) {
-        return new AnonymousCredentialsProvider();
-    }
-    throw new Error("No YDB credentials configured. Set one of: YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS, YDB_METADATA_CREDENTIALS=1, YDB_ACCESS_TOKEN_CREDENTIALS, YDB_ANONYMOUS_CREDENTIALS=1");
-}
 export function isCompilationTimeoutError(error) {
     if (!(error instanceof Error)) {
         return false;
@@ -207,54 +63,6 @@ export function isCompilationTimeoutError(error) {
     }
     return false;
 }
-function getErrorCause(error) {
-    return error.cause;
-}
-function getCauseName(cause) {
-    if (cause instanceof Error) {
-        return cause.name;
-    }
-    if (typeof cause === "object" && cause !== null) {
-        const n = cause.name;
-        if (typeof n === "string") {
-            return n;
-        }
-    }
-    return undefined;
-}
-export function isTimeoutAbortError(error) {
-    // `@ydbjs/query` implements `.timeout(ms)` by adding `AbortSignal.timeout(ms)` to the query
-    // signal chain. When that fires, @ydbjs/retry wraps it via `@ydbjs/abortable` and we end up
-    // observing a DOMException `{ name: "AbortError" }`, often with an underlying `cause`
-    // `{ name: "TimeoutError" }`.
-    if (!(error instanceof Error)) {
-        return false;
-    }
-    if (error.name === "TimeoutError") {
-        return true;
-    }
-    if (error.name !== "AbortError") {
-        return false;
-    }
-    const cause = getErrorCause(error);
-    const causeName = getCauseName(cause);
-    if (causeName === "TimeoutError") {
-        return true;
-    }
-    // Node versions differ in how abort reasons are propagated through `AbortSignal.any(...)`.
-    // In our codebase we currently only use AbortError signals for query timeouts / cancellations,
-    // so treat an AbortError without a usable cause as timeout-like.
-    return true;
-}
-export function getAbortErrorCause(error) {
-    if (!(error instanceof Error)) {
-        return undefined;
-    }
-    if (error.name !== "AbortError") {
-        return undefined;
-    }
-    return getErrorCause(error);
-}
 function shouldTriggerDriverRefresh(error) {
     if (!(error instanceof Error)) {
         return false;
@@ -269,8 +77,9 @@ function shouldTriggerDriverRefresh(error) {
     if (/SessionExpired|SESSION_EXPIRED|session.*expired/i.test(msg)) {
         return true;
     }
-    // Compilation timeout is treated as a signal to refresh the driver so that
-    // subsequent attempts use a fresh connection state.
+    // YDB query compilation timeout (TIMEOUT code 400090) – treat as a signal
+    // to refresh the driver/session pool so that subsequent attempts use a
+    // fresh connection state.
     if (isCompilationTimeoutError(error)) {
         return true;
     }
@@ -306,12 +115,6 @@ async function maybeRefreshDriverOnSessionError(error) {
 }
 export function __setDriverForTests(fake) {
     driver = fake;
-    sqlClient = undefined;
-    sessionPool = undefined;
-    queryCtxPromise = null;
-}
-export function __setSessionPoolForTests(fake) {
-    sessionPool = fake;
 }
 export function __setDriverFactoryForTests(factory) {
     driverFactoryOverride = factory;
@@ -327,207 +130,81 @@ export function configureDriver(config) {
     }
     overrideConfig = config;
 }
-function buildConnectionString(args) {
-    const endpoint = args.endpoint.trim();
-    const database = args.database.trim();
-    if (!endpoint) {
-        throw new Error("YDB endpoint is empty; set YDB_ENDPOINT");
-    }
-    if (!database) {
-        throw new Error("YDB database is empty; set YDB_DATABASE");
-    }
-    const url = new URL(endpoint);
-    const dbPath = database.startsWith("/") ? database : `/${database}`;
-    url.pathname = dbPath;
-    return url.toString();
-}
 function getOrCreateDriver() {
     if (driver) {
         return driver;
     }
-    const endpoint = overrideConfig?.endpoint ?? YDB_ENDPOINT;
-    const database = overrideConfig?.database ?? YDB_DATABASE;
-    let connectionString;
-    if (overrideConfig?.connectionString) {
-        connectionString = overrideConfig.connectionString;
-    }
-    else if (driverFactoryOverride) {
-        // Tests may inject a driver factory without setting env vars.
-        // If endpoint/database are available, still build a real connection string.
-        const hasEndpoint = endpoint.trim().length > 0;
-        const hasDatabase = database.trim().length > 0;
-        connectionString =
-            hasEndpoint && hasDatabase
-                ? buildConnectionString({ endpoint, database })
-                : "";
-    }
-    else {
-        connectionString = buildConnectionString({ endpoint, database });
-    }
-    const credentialsProvider = overrideConfig?.credentialsProvider
-        ? overrideConfig.credentialsProvider
-        : driverFactoryOverride
-            ? undefined
-            : resolveCredentialsProviderFromEnv();
-    const driverOptions = {
-        ...(credentialsProvider ? { credentialsProvider } : {}),
-        "ydb.sdk.ready_timeout_ms": DRIVER_READY_TIMEOUT_MS,
+    const base = overrideConfig?.connectionString != null
+        ? { connectionString: overrideConfig.connectionString }
+        : {
+            endpoint: overrideConfig?.endpoint ?? YDB_ENDPOINT,
+            database: overrideConfig?.database ?? YDB_DATABASE,
+        };
+    const driverConfig = {
+        ...base,
+        authService: overrideConfig?.authService ?? getCredentialsFromEnv(),
+        poolSettings: {
+            minLimit: SESSION_POOL_MIN_SIZE,
+            maxLimit: SESSION_POOL_MAX_SIZE,
+            keepAlivePeriod: SESSION_KEEPALIVE_PERIOD_MS,
+        },
     };
     driver = driverFactoryOverride
-        ? driverFactoryOverride(connectionString, driverOptions)
-        : new Driver(connectionString, driverOptions);
-    logger.info({ hasCredentialsProvider: true, connectionStringMasked: true }, "YDB driver created");
+        ? driverFactoryOverride(driverConfig)
+        : new Driver(driverConfig);
+    logger.info({
+        poolMinSize: SESSION_POOL_MIN_SIZE,
+        poolMaxSize: SESSION_POOL_MAX_SIZE,
+        keepAlivePeriodMs: SESSION_KEEPALIVE_PERIOD_MS,
+    }, "YDB driver created with session pool settings");
     return driver;
 }
-/**
- * @internal
- * Exposes the singleton YDB `Driver` for internal modules that need non-QueryService RPCs
- * (e.g. TableService BulkUpsert).
- */
-export function __getDriverForInternalUse() {
-    return getOrCreateDriver();
-}
-function getOrCreateQueryClient() {
-    if (sqlClient) {
-        return sqlClient;
-    }
+export async function readyOrThrow() {
     const d = getOrCreateDriver();
-    sqlClient = query(d);
-    return sqlClient;
-}
-function getOrCreateSessionPool() {
-    if (sessionPool) {
-        return sessionPool;
+    const ok = await d.ready(DRIVER_READY_TIMEOUT_MS);
+    if (!ok) {
+        throw new Error(`YDB driver is not ready in ${DRIVER_READY_TIMEOUT_MS / 1000}s. Check connectivity and credentials.`);
     }
-    const d = getOrCreateDriver();
-    sessionPool = new SessionPool(d);
-    // Keepalive is best-effort; don't block startup.
-    sessionPool.start();
-    return sessionPool;
 }
-async function withTimeoutSignal(timeoutMs, fn) {
-    const ac = new AbortController();
-    const t = setTimeout(() => ac.abort(), timeoutMs);
+export async function withSession(fn) {
+    const d = getOrCreateDriver();
     try {
-        return await fn(ac.signal);
+        return await d.tableClient.withSession(fn, TABLE_SESSION_TIMEOUT_MS);
     }
-    finally {
-        clearTimeout(t);
+    catch (err) {
+        void maybeRefreshDriverOnSessionError(err);
+        throw err;
     }
 }
-export async function readyOrThrow() {
-    const d = getOrCreateDriver();
-    await withTimeoutSignal(DRIVER_READY_TIMEOUT_MS, async (signal) => {
-        await d.ready(signal);
-    });
-}
-export async function withSession(fn) {
-    const sql = getOrCreateQueryClient();
+export async function withQuerySession(fn, options) {
     const d = getOrCreateDriver();
-    const pool = getOrCreateSessionPool();
     try {
-        // Bound only `d.ready()` and `pool.acquire()` so stalls in driver readiness or session
-        // availability can't hang forever. We intentionally do NOT abort the user callback,
-        // because request operations can legitimately exceed this bound (e.g. multi-batch upserts).
-        await withTimeoutSignal(WITH_SESSION_TIMEOUT_MS, async (signal) => {
-            await d.ready(signal);
+        return await d.queryClient.do({
+            fn,
+            timeout: options?.timeoutMs ?? TABLE_SESSION_TIMEOUT_MS,
+            idempotent: options?.idempotent,
         });
-        // Ensure we have a few sessions ready to reduce cold-start session churn.
-        // Best-effort: ignore failures/timeouts.
-        void withTimeoutSignal(WITH_SESSION_TIMEOUT_MS, async (signal) => {
-            await pool.warmup(signal);
-        }).catch(() => undefined);
-        // @ydbjs/query QueryClient.do() is not implemented in the currently published builds.
-        // Cancellation is cooperative: we provide an AbortSignal to `fn()` and also store it in
-        // the @ydbjs/query async-local context so query operations can observe it via `.signal(...)`.
-        // (Note: this signal is not auto-aborted by `withSession()`; query-level `.timeout(...)`
-        // is the primary guard for long-running YQL statements.)
-        const operationSignal = new AbortController().signal;
-        const ctxMod = await getQueryCtx();
-        let leased = null;
-        const MAX_ATTEMPTS = 3;
-        for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt += 1) {
-            leased = await withTimeoutSignal(WITH_SESSION_TIMEOUT_MS, async (signal) => {
-                return await pool.acquire(signal);
-            });
-            try {
-                const store = {
-                    nodeId: leased.nodeId,
-                    sessionId: leased.sessionId,
-                    signal: operationSignal,
-                };
-                const result = await ctxMod.ctx.run(store, async () => {
-                    return await fn(sql, operationSignal);
-                });
-                pool.release(leased);
-                return result;
-            }
-            catch (err) {
-                const shouldRetryWithFreshSession = err instanceof Error &&
-                    /BAD_SESSION|SESSION_EXPIRED|SessionExpired|No session became available/i.test(err.message ?? "");
-                if (shouldRetryWithFreshSession) {
-                    await pool.discard(leased);
-                    leased = null;
-                    if (attempt === MAX_ATTEMPTS - 1) {
-                        throw err;
-                    }
-                    continue;
-                }
-                pool.release(leased);
-                throw err;
-            }
-        }
-        throw new Error("withSession: exhausted attempts to acquire a healthy session");
     }
     catch (err) {
+        // Query sessions can also get stuck/busy; reuse the same driver refresh path.
         void maybeRefreshDriverOnSessionError(err);
         throw err;
     }
 }
 export async function withStartupProbeSession(fn) {
-    const sql = getOrCreateQueryClient();
     const d = getOrCreateDriver();
-    const pool = getOrCreateSessionPool();
-    const ac = new AbortController();
-    const t = setTimeout(() => ac.abort(), STARTUP_PROBE_SESSION_TIMEOUT_MS);
     try {
-        // Ensure the startup probe timeout also applies to driver readiness.
-        await d.ready(ac.signal);
-        void pool.warmup(ac.signal);
-        const ctxMod = await getQueryCtx();
-        const leased = await pool.acquire(ac.signal);
-        try {
-            const store = {
-                nodeId: leased.nodeId,
-                sessionId: leased.sessionId,
-                signal: ac.signal,
-            };
-            const result = await ctxMod.ctx.run(store, async () => {
-                return await fn(sql, ac.signal);
-            });
-            pool.release(leased);
-            return result;
-        }
-        catch (err) {
-            pool.release(leased);
-            throw err;
-        }
+        return await d.tableClient.withSession(fn, STARTUP_PROBE_SESSION_TIMEOUT_MS);
     }
     catch (err) {
         void maybeRefreshDriverOnSessionError(err);
         throw err;
     }
-    finally {
-        clearTimeout(t);
-    }
 }
 export async function isYdbAvailable(timeoutMs = YDB_HEALTHCHECK_READY_TIMEOUT_MS) {
     const d = getOrCreateDriver();
     try {
-        await withTimeoutSignal(timeoutMs, async (signal) => {
-            await d.ready(signal);
-        });
-        return true;
+        return await d.ready(timeoutMs);
     }
     catch {
         return false;
@@ -541,27 +218,9 @@ export async function destroyDriver() {
     if (!driver) {
         return;
     }
-    logger.info("Destroying YDB driver");
+    logger.info("Destroying YDB driver and session pool");
     try {
-        if (sessionPool) {
-            try {
-                await sessionPool.close();
-            }
-            catch (err) {
-                logger.warn({ err }, "Error during session pool close (ignored)");
-            }
-            sessionPool = undefined;
-        }
-        if (sqlClient) {
-            try {
-                await sqlClient[Symbol.asyncDispose]();
-            }
-            catch (err) {
-                logger.warn({ err }, "Error during query client disposal (ignored)");
-            }
-            sqlClient = undefined;
-        }
-        driver.close();
+        await driver.destroy();
     }
     catch (err) {
         logger.warn({ err }, "Error during driver destruction (ignored)");