src/services/api/streaming.ts - main - githubmirror/ydb-embedded-ui

src/services/api/streaming.ts
202 строки7 KB
Konstantin Kozlov
fix: send database query param for post requests (#4168)
27 июл 2026, 17:47
27 июл 2026, 17:474ee8923
import {parseMultipart} from '@mjackson/multipart-parser';
import qs from 'qs';

import {
    isErrorChunk,
    isKeepAliveChunk,
    isQueryResponseChunk,
    isSessionChunk,
    isStreamDataChunk,
} from '../../store/reducers/query/utils';
import {readSettingValueFromLS} from '../../store/reducers/settings/utils';
import type {Actions, StreamQueryParams} from '../../types/api/query';
import type {
    QueryResponseChunk,
    SessionChunk,
    StreamDataChunk,
    StreamingChunk,
} from '../../types/store/streaming';
import {DEV_ENABLE_TRACING_FOR_ALL_REQUESTS} from '../../utils/constants';
import {EXTRACTABLE_RESPONSE_HEADERS} from '../../utils/errors/extractErrorDetails';
import {isRedirectToAuth} from '../../utils/response';

import {BaseYdbAPI, CSRF_TOKEN_HEADER_NAME} from './base';
import {readPartText} from './streamingPartReader';
import {isNeedResetResponse, processNeedReset} from './utils/needReset';

/**
 * Extracts useful headers from a fetch Response into a plain object.
 * Only includes headers that extractErrorDetails knows how to use.
 */
function extractResponseHeaders(response: Response): Record<string, string> {
    const result: Record<string, string> = {};
    for (const name of EXTRACTABLE_RESPONSE_HEADERS) {
        const value = response.headers.get(name);
        if (value) {
            result[name] = value;
        }
    }
    return result;
}

/**
 * Creates a structured error from a fetch Response.
 * Returns a real Error (with stack trace and instanceof support) enriched
 * with the fields that extractErrorDetails expects (status, statusText, headers, data, config).
 */
function createStreamingResponseError(response: Response, responseData: unknown): Error {
    const message = `${response.status} ${response.statusText || 'Error'}`;
    const error = new Error(message);
    return Object.assign(error, {
        status: response.status,
        statusText: response.statusText,
        data: responseData,
        headers: extractResponseHeaders(response),
        config: {url: response.url, method: 'POST'},
    });
}

const BOUNDARY = 'boundary';

export interface StreamQueryOptions {
    signal?: AbortSignal;
    onStreamDataChunk: (chunk: StreamDataChunk) => void;
    onQueryResponseChunk: (chunk: QueryResponseChunk) => void;
    onSessionChunk: (chunk: SessionChunk) => void;
}

export class StreamingAPI extends BaseYdbAPI {
    async streamQuery<Action extends Actions>(
        params: StreamQueryParams<Action>,
        options: StreamQueryOptions,
    ) {
        const base64 = params.base64;

        const queryParams = qs.stringify(
            {database: params.database, timeout: params.timeout, base64, schema: 'multipart'},
            {encoder: encodeURIComponent},
        );

        const body = {...params, base64, schema: 'multipart'};
        const headers = new Headers({
            Accept: 'multipart/form-data',
            'Content-Type': 'application/json',
        });

        const csrfToken = this.getCsrfToken();
        if (csrfToken) {
            headers.set(CSRF_TOKEN_HEADER_NAME, csrfToken);
        }

        if (params.tracingLevel) {
            headers.set('X-Trace-Verbosity', String(params.tracingLevel));
        }

        const enableTracing = readSettingValueFromLS(DEV_ENABLE_TRACING_FOR_ALL_REQUESTS);

        if (enableTracing) {
            headers.set('X-Want-Trace', '1');
        }

        const url = `${this.getPath('/viewer/query')}?${queryParams}`;
        let response: Response;
        try {
            response = await fetch(url, {
                method: 'POST',
                signal: options.signal,
                headers,
                credentials: this._axios.defaults.withCredentials ? 'include' : 'same-origin',
                body: JSON.stringify(body),
            });
        } catch (fetchError) {
            const enriched =
                fetchError instanceof Error ? fetchError : new Error(String(fetchError));
            Object.assign(enriched, {
                config: {url, method: 'POST'},
                errorPhase: 'connection',
                networkOnline: navigator.onLine,
            });
            throw enriched;
        }

        if (!response.ok) {
            const responseData = await response
                .text()
                .then((text) => {
                    try {
                        return JSON.parse(text) as unknown;
                    } catch {
                        return text || undefined;
                    }
                })
                .catch(() => undefined);
            if (isRedirectToAuth({status: response.status, data: responseData})) {
                const data = responseData as {authUrl: string};
                window.location.assign(data.authUrl);
                return;
            }
            if (isNeedResetResponse(responseData) && document.visibilityState === 'visible') {
                processNeedReset();
            }
            throw createStreamingResponseError(response, responseData);
        }

        if (!response.body) {
            const error = new Error('Empty response body');
            Object.assign(error, {
                status: response.status,
                statusText: response.statusText,
                config: {url, method: 'POST'},
                headers: extractResponseHeaders(response),
            });
            throw error;
        }

        const traceId = response.headers.get('traceresponse')?.split('-')[1];

        try {
            await parseMultipart(response.body, {boundary: BOUNDARY}, async (part) => {
                const text = await readPartText(part);

                let chunk: unknown;
                try {
                    chunk = JSON.parse(text);
                } catch (e) {
                    const preview = text.length > 200 ? text.slice(0, 200) + '…' : text;
                    throw new Error(`Error parsing chunk: ${e}\nRaw: ${preview}`);
                }

                if (isErrorChunk(chunk)) {
                    await response.body?.cancel().catch(() => {});
                    throw chunk;
                }

                const streamingChunk = chunk as StreamingChunk;

                if (isSessionChunk(streamingChunk)) {
                    const sessionChunk = streamingChunk;
                    sessionChunk.meta.trace_id = traceId;
                    options.onSessionChunk(streamingChunk);
                } else if (isStreamDataChunk(streamingChunk)) {
                    options.onStreamDataChunk(streamingChunk);
                } else if (isQueryResponseChunk(streamingChunk)) {
                    options.onQueryResponseChunk(streamingChunk);
                } else if (isKeepAliveChunk(streamingChunk)) {
                    console.info('Received keep alive chunk');
                }
            });
        } catch (streamError) {
            if (isErrorChunk(streamError)) {
                throw streamError;
            }
            const enriched =
                streamError instanceof Error ? streamError : new Error(String(streamError));
            Object.assign(enriched, {
                config: {url, method: 'POST'},
                headers: extractResponseHeaders(response),
                errorPhase: 'stream',
            });
            throw enriched;
        }
    }
}