Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : 
       3              : import type { EmbeddedComputeDatabase, EmbeddedTableImportConnection } from '../../platform/database/embedded_database.js';
       4              : import { quoteLiteral } from '../../platform/hyperdb/hyperdb_incremental_extract.js';
       5              : import { ConnectorType } from '../../connections/connector_info.js';
       6              : import { utf8Vector } from '../../utils/arrow_vectors.js';
       7              : import type { SnippetRegistry } from './snippet_registry.js';
       8              : 
       9              : /// The table of the snippets in the snippet database, one row per snippet and connector
      10            1 : export const SNIPPET_SEARCH_TABLE = 'snippets';
      11              : 
      12              : /// The rows of the snippet table
      13            1 : function snippetRows(registry: SnippetRegistry): arrow.Table {
      14            1 :     const rows = registry.snippets.flatMap(snippet => snippet.connectors.map(connector => ({
      15            1 :         id: snippet.id,
      16            1 :         category: snippet.category,
      17            1 :         connector: ConnectorType[connector],
      18            1 :         path: snippet.paths.get(connector)!,
      19            1 :         body: snippet.bodies.get(connector)!,
      20            1 :     })));
      21            1 :     return new arrow.Table({
      22            1 :         id: utf8Vector(rows.map(row => row.id)),
      23            1 :         category: utf8Vector(rows.map(row => row.category)),
      24            1 :         connector: utf8Vector(rows.map(row => row.connector)),
      25            1 :         path: utf8Vector(rows.map(row => row.path)),
      26            1 :         body: utf8Vector(rows.map(row => row.body)),
      27            1 :     });
      28            1 : }
      29              : 
      30              : /// The ILIKE pattern of the snippets that contain a text, the wildcards in the text match only themselves
      31            1 : export function snippetSearchPattern(text: string): string {
      32            1 :     return `%${text.replace(/[\\%_]/g, char => `\\${char}`)}%`;
      33            1 : }
      34              : 
      35              : /// The full-text search over the snippets in the snippet database of the embedded Hyper.
      36              : /// The search runs its statements one after the other on one connection, the database serves one statement at a
      37              : /// time per connection.
      38            1 : export class SnippetSearch {
      39              :     /// The registry in the snippet table, null before the first search
      40            1 :     private loaded: SnippetRegistry | null = null;
      41              :     /// The last statement sequence, the next one starts after it
      42            1 :     private operation: Promise<unknown> = Promise.resolve();
      43              : 
      44            1 :     constructor(private readonly connection: EmbeddedTableImportConnection) {}
      45              : 
      46              :     /// Connect to the snippet database of a database instance
      47            1 :     static async connect(database: EmbeddedComputeDatabase): Promise<SnippetSearch> {
      48            1 :         return new SnippetSearch(await database.connect({ defaultDatabase: '__dashql_snippets' }));
      49            1 :     }
      50              : 
      51              :     /// The ids of the snippets of a registry that run on a connector and contain a text in their id or text, ignoring
      52              :     /// case. The text includes the description. The snippet table is filled with the registry first, if it holds
      53              :     /// another one.
      54            1 :     search(registry: SnippetRegistry, connector: ConnectorType, text: string): Promise<Set<string>> {
      55            1 :         return this.enqueue(async () => {
      56            1 :             if (this.loaded !== registry) {
      57            1 :                 this.loaded = null;
      58            1 :                 await this.connection.execute(`DROP TABLE IF EXISTS ${SNIPPET_SEARCH_TABLE}`);
      59            1 :                 await this.connection.insertArrowTable(snippetRows(registry), { name: SNIPPET_SEARCH_TABLE, create: true });
      60            1 :                 this.loaded = registry;
      61            1 :             }
      62            1 :             const pattern = quoteLiteral(snippetSearchPattern(text));
      63            1 :             const matches = (column: string) => `${column} ILIKE ${pattern} ESCAPE '\\'`;
      64            1 :             const result = await this.connection.query(`
      65            1 :                 SELECT id FROM ${SNIPPET_SEARCH_TABLE}
      66            1 :                 WHERE connector = ${quoteLiteral(ConnectorType[connector])}
      67            1 :                     AND (${matches('id')} OR ${matches('body')})
      68            1 :             `);
      69            1 :             const ids = new Set<string>();
      70            1 :             for (const id of result.getChild('id') ?? []) ids.add(String(id));
      71            1 :             return ids;
      72            1 :         });
      73            1 :     }
      74              : 
      75              :     /// Close the connection after the running search
      76            1 :     close(): Promise<void> {
      77            1 :         return this.enqueue(() => this.connection.close());
      78            1 :     }
      79              : 
      80            1 :     private enqueue<T>(operation: () => Promise<T>): Promise<T> {
      81            1 :         const next = this.operation.then(operation, operation);
      82            1 :         this.operation = next.catch(() => {});
      83            1 :         return next;
      84            1 :     }
      85            1 : }
      86              : 
      87              : /// The snippet searches by database instance, one connection per instance
      88            1 : const SEARCHES = new WeakMap<EmbeddedComputeDatabase, Promise<SnippetSearch>>();
      89              : 
      90              : /// The snippet search of a database instance, connected once
      91            1 : export function getSnippetSearch(database: EmbeddedComputeDatabase): Promise<SnippetSearch> {
      92            1 :     let search = SEARCHES.get(database);
      93            1 :     if (search == null) {
      94            1 :         search = SnippetSearch.connect(database);
      95            1 :         search.catch(() => SEARCHES.delete(database));
      96            1 :         SEARCHES.set(database, search);
      97            1 :     }
      98            1 :     return search;
      99            1 : }