ResultSetDataSource.ts

/*
 * CloudBeaver - Cloud Database Manager
 * Copyright (C) 2020-2026 DBeaver Corp and others
 *
 * Licensed under the Apache License, Version 2.0.
 * you may not use this file except in compliance with the License.
 */
import { autorun, type IReactionDisposer, makeObservable, observable } from 'mobx';

import type { IConnectionExecutionContext, IConnectionExecutionContextInfo } from '@cloudbeaver/core-connections';
import type { IServiceProvider } from '@cloudbeaver/core-di';
import type { ITask } from '@cloudbeaver/core-executor';
import { AsyncTaskInfoService } from '@cloudbeaver/core-root';
import type { GraphQLService, SqlRowIdentifier, SqlRowIdentifierState } from '@cloudbeaver/core-sdk';

import { DatabaseDataSource } from '../DatabaseDataModel/DatabaseDataSource.js';
import { type IDatabaseDataOptions } from '../DatabaseDataModel/IDatabaseDataOptions.js';
import type { IDatabaseResultSet } from '../DatabaseDataModel/IDatabaseResultSet.js';
import {
  applyPersistedDataFilterConstraints,
  DatabaseDataConstraintAction,
  persistDataFilterConstraints,
} from '../DatabaseDataModel/Actions/DatabaseDataConstraintAction.js';
import { DocumentDataAction } from '../DatabaseDataModel/Actions/Document/DocumentDataAction.js';
import { DocumentEditAction } from '../DatabaseDataModel/Actions/Document/DocumentEditAction.js';
import { IDatabaseDataCacheAction } from '../DatabaseDataModel/Actions/IDatabaseDataCacheAction.js';
import { IDatabaseDataConstraintAction } from '../DatabaseDataModel/Actions/IDatabaseDataConstraintAction.js';
import { IDatabaseDataEditAction } from '../DatabaseDataModel/Actions/IDatabaseDataEditAction.js';
import { IDatabaseDataFormatAction } from '../DatabaseDataModel/Actions/IDatabaseDataFormatAction.js';
import { IDatabaseDataResultAction } from '../DatabaseDataModel/Actions/IDatabaseDataResultAction.js';
import { IDatabaseDataViewAction } from '../DatabaseDataModel/Actions/IDatabaseDataViewAction.js';
import { ResultSetCacheAction } from '../DatabaseDataModel/Actions/ResultSet/ResultSetCacheAction.js';
import { ResultSetDataAction } from '../DatabaseDataModel/Actions/ResultSet/ResultSetDataAction.js';
import { ResultSetEditAction } from '../DatabaseDataModel/Actions/ResultSet/ResultSetEditAction.js';
import { ResultSetFormatAction } from '../DatabaseDataModel/Actions/ResultSet/ResultSetFormatAction.js';
import { ResultSetSelectAction } from '../DatabaseDataModel/Actions/ResultSet/ResultSetSelectAction.js';
import { ResultSetViewAction } from '../DatabaseDataModel/Actions/ResultSet/ResultSetViewAction.js';
import { IDatabaseDataSelectAction } from '../DatabaseDataModel/Actions/IDatabaseDataSelectAction.js';
import { DatabaseDataFeature } from '../DatabaseDataModel/IDatabaseDataSource.js';

export interface IRowIdentifierInfo {
  state: SqlRowIdentifierState | null;
  identifier: SqlRowIdentifier | null;
}

export abstract class ResultSetDataSource<TOptions extends IDatabaseDataOptions = IDatabaseDataOptions> extends DatabaseDataSource<
  TOptions,
  IDatabaseResultSet
> {
  executionContext: IConnectionExecutionContext | null;
  totalCountRequestTask: ITask<number> | null;
  private keepExecutionContextOnDispose: boolean;
  private readonly persistConstraintsDisposer: IReactionDisposer;

  constructor(
    override readonly serviceProvider: IServiceProvider,
    protected graphQLService: GraphQLService,
    protected asyncTaskInfoService: AsyncTaskInfoService,
  ) {
    super(serviceProvider);
    this.totalCountRequestTask = null;
    this.executionContext = null;
    this.keepExecutionContextOnDispose = false;
    this.setFeature(DatabaseDataFeature.ResultSet);

    this.actions
      .registerAction(IDatabaseDataResultAction, DocumentDataAction)
      .registerAction(IDatabaseDataEditAction, DocumentEditAction)

      .registerAction(IDatabaseDataResultAction, ResultSetDataAction)
      .registerAction(IDatabaseDataEditAction, ResultSetEditAction)
      .registerAction(IDatabaseDataViewAction, ResultSetViewAction)
      .registerAction(IDatabaseDataSelectAction, ResultSetSelectAction)
      .registerAction(IDatabaseDataFormatAction, ResultSetFormatAction)
      .registerAction(IDatabaseDataCacheAction, ResultSetCacheAction)
      .registerAction(IDatabaseDataConstraintAction, DatabaseDataConstraintAction);

    makeObservable(this, {
      totalCountRequestTask: observable.ref,
      executionContext: observable,
    });

    this.persistConstraintsDisposer = autorun(() => persistDataFilterConstraints(this));
  }

  protected override onPersistedStateLoaded(): void {
    applyPersistedDataFilterConstraints(this);
  }

  override isReadonly(resultIndex: number): boolean {
    return super.isReadonly(resultIndex) || !this.executionContext?.context || !!this.getResult(resultIndex)?.data?.readOnly;
  }

  override async cancel(): Promise<void> {
    await super.cancel();
    await this.cancelLoadTotalCount();
  }

  override async saveData(): Promise<void> {
    await super.saveData();
    // TODO: Remove this when we have virtual keys. We need to refresh the data in tables without a primary key to avoid UI glitch #5140.
    if (!this.hasElementIdentifier(0)) {
      this.setOutdated();
    }
  }

  async cancelLoadTotalCount(): Promise<ITask<number> | null> {
    await this.totalCountRequestTask?.cancel();

    return this.totalCountRequestTask;
  }

  async loadTotalCount(resultIndex: number): Promise<ITask<number>> {
    const executionContext = this.executionContext;
    const executionContextInfo = this.executionContext?.context;

    if (!executionContext || !executionContextInfo) {
      throw new Error('Context must be provided');
    }

    const result = this.getResult(resultIndex);

    if (!result?.id) {
      throw new Error('Result id must be provided');
    }

    const asyncTask = this.asyncTaskInfoService.create(async () => {
      const { taskInfo } = await this.graphQLService.sdk.asyncSqlRowDataCount({
        resultsId: result.id!,
        connectionId: executionContextInfo.connectionId,
        contextId: executionContextInfo.id,
        projectId: executionContextInfo.projectId,
      });

      return taskInfo;
    });

    const task = executionContext.run(
      async () => {
        const info = await this.asyncTaskInfoService.run(asyncTask);

        const { count } = await this.graphQLService.sdk.getSqlRowDataCountResult({ taskId: info.id });

        return count;
      },
      () => this.asyncTaskInfoService.cancel(asyncTask.id),
      () => this.asyncTaskInfoService.remove(asyncTask.id),
    );

    this.totalCountRequestTask = task;

    const count = await task;
    this.setTotalCount(resultIndex, count);

    return this.totalCountRequestTask;
  }

  override setResults(results: IDatabaseResultSet[]): this {
    this.closeResults(this.results.filter(result => !results.some(r => r.id === result.id)));
    return super.setResults(results);
  }

  override async dispose(): Promise<void> {
    this.persistConstraintsDisposer();
    await super.dispose();
    if (this.keepExecutionContextOnDispose) {
      await this.closeResults(this.results);
    } else {
      await this.executionContext?.destroy();
    }
  }

  setKeepExecutionContextOnDispose(keep: boolean): this {
    this.keepExecutionContextOnDispose = keep;
    return this;
  }

  setExecutionContext(context: IConnectionExecutionContext | null): this {
    this.executionContext = context;
    this.setOutdated();
    return this;
  }

  hasElementIdentifier(resultIndex: number): boolean {
    return this.getResult(resultIndex)?.data?.hasRowIdentifier === true;
  }

  getRowIdentifierInfo(resultIndex: number): IRowIdentifierInfo {
    const data = this.getResult(resultIndex)?.data;
    return {
      state: (data?.rowIdentifierState as SqlRowIdentifierState | undefined) ?? null,
      identifier: (data?.rowIdentifier as SqlRowIdentifier | undefined) ?? null,
    };
  }

  protected getPreviousResultId(prevResults: IDatabaseResultSet[], context: IConnectionExecutionContextInfo) {
    let resultId: string | undefined;

    if (
      prevResults.length === 1 &&
      prevResults[0]!.contextId === context.id &&
      prevResults[0]!.connectionId === context.connectionId &&
      prevResults[0]!.id !== null
    ) {
      resultId = prevResults[0]!.id;
    }

    return resultId;
  }

  private setTotalCount(resultIndex: number, count: number): this {
    const result = this.getResult(resultIndex);

    if (result) {
      result.totalCount = count;
    }
    return this;
  }

  private async closeResults(results: IDatabaseResultSet[]): Promise<void> {
    if (!this.executionContext?.context) {
      return;
    }

    for (const result of results) {
      // TODO: it's better to track that context is closed with subscription
      if (result.id === null || result.contextId !== this.executionContext.context.id) {
        continue;
      }
      try {
        await this.graphQLService.sdk.closeResult({
          projectId: result.projectId,
          connectionId: result.connectionId,
          contextId: result.contextId,
          resultId: result.id,
        });
      } catch (exception: any) {
        console.log(`Error closing result (${result.id}):`, exception);
      }
    }
  }
}

export function isResultSetDataSource<T extends IDatabaseDataOptions = IDatabaseDataOptions>(
  dataSource: unknown,
): dataSource is ResultSetDataSource<T> {
  return dataSource instanceof ResultSetDataSource;
}