CoolFace
Apppublic

AK-21/Graphite-Industrial-Intelligence

sourceHugging Faceupdated 3mo agoView on Hugging Face
0likes
session.js157 linesDownload Raw Back to tidb-serverless
1import { NoopCache } from "../cache/core/index.js";2import { Column } from "../column.js";3import { entityKind, is } from "../entity.js";4import { NoopLogger } from "../logger.js";5import {6  MySqlPreparedQuery,7  MySqlSession,8  MySqlTransaction9} from "../mysql-core/session.js";10import { fillPlaceholders, sql } from "../sql/sql.js";11import { mapResultRow } from "../utils.js";12const executeRawConfig = { fullResult: true };13const queryConfig = { arrayMode: true };14class TiDBServerlessPreparedQuery extends MySqlPreparedQuery {15  constructor(client, queryString, params, logger, cache, queryMetadata, cacheConfig, fields, customResultMapper, generatedIds, returningIds) {16    super(cache, queryMetadata, cacheConfig);17    this.client = client;18    this.queryString = queryString;19    this.params = params;20    this.logger = logger;21    this.fields = fields;22    this.customResultMapper = customResultMapper;23    this.generatedIds = generatedIds;24    this.returningIds = returningIds;25  }26  static [entityKind] = "TiDBPreparedQuery";27  async execute(placeholderValues = {}) {28    const params = fillPlaceholders(this.params, placeholderValues);29    this.logger.logQuery(this.queryString, params);30    const { fields, client, queryString, joinsNotNullableMap, customResultMapper, returningIds, generatedIds } = this;31    if (!fields && !customResultMapper) {32      const res = await this.queryWithCache(queryString, params, async () => {33        return await client.execute(queryString, params, executeRawConfig);34      });35      const insertId = res.lastInsertId ?? 0;36      const affectedRows = res.rowsAffected ?? 0;37      if (returningIds) {38        const returningResponse = [];39        let j = 0;40        for (let i = insertId; i < insertId + affectedRows; i++) {41          for (const column of returningIds) {42            const key = returningIds[0].path[0];43            if (is(column.field, Column)) {44              if (column.field.primary && column.field.autoIncrement) {45                returningResponse.push({ [key]: i });46              }47              if (column.field.defaultFn && generatedIds) {48                returningResponse.push({ [key]: generatedIds[j][key] });49              }50            }51          }52          j++;53        }54        return returningResponse;55      }56      return res;57    }58    const rows = await this.queryWithCache(queryString, params, async () => {59      return await client.execute(queryString, params, queryConfig);60    });61    if (customResultMapper) {62      return customResultMapper(rows);63    }64    return rows.map((row) => mapResultRow(fields, row, joinsNotNullableMap));65  }66  iterator(_placeholderValues) {67    throw new Error("Streaming is not supported by the TiDB Cloud Serverless driver");68  }69}70class TiDBServerlessSession extends MySqlSession {71  constructor(baseClient, dialect, tx, schema, options = {}) {72    super(dialect);73    this.baseClient = baseClient;74    this.schema = schema;75    this.options = options;76    this.client = tx ?? baseClient;77    this.logger = options.logger ?? new NoopLogger();78    this.cache = options.cache ?? new NoopCache();79  }80  static [entityKind] = "TiDBServerlessSession";81  logger;82  client;83  cache;84  prepareQuery(query, fields, customResultMapper, generatedIds, returningIds, queryMetadata, cacheConfig) {85    return new TiDBServerlessPreparedQuery(86      this.client,87      query.sql,88      query.params,89      this.logger,90      this.cache,91      queryMetadata,92      cacheConfig,93      fields,94      customResultMapper,95      generatedIds,96      returningIds97    );98  }99  all(query) {100    const querySql = this.dialect.sqlToQuery(query);101    this.logger.logQuery(querySql.sql, querySql.params);102    return this.client.execute(querySql.sql, querySql.params);103  }104  async count(sql2) {105    const res = await this.execute(sql2);106    return Number(107      res["rows"][0]["count"]108    );109  }110  async transaction(transaction) {111    const nativeTx = await this.baseClient.begin();112    try {113      const session = new TiDBServerlessSession(this.baseClient, this.dialect, nativeTx, this.schema, this.options);114      const tx = new TiDBServerlessTransaction(115        this.dialect,116        session,117        this.schema118      );119      const result = await transaction(tx);120      await nativeTx.commit();121      return result;122    } catch (err) {123      await nativeTx.rollback();124      throw err;125    }126  }127}128class TiDBServerlessTransaction extends MySqlTransaction {129  static [entityKind] = "TiDBServerlessTransaction";130  constructor(dialect, session, schema, nestedIndex = 0) {131    super(dialect, session, schema, nestedIndex, "default");132  }133  async transaction(transaction) {134    const savepointName = `sp${this.nestedIndex + 1}`;135    const tx = new TiDBServerlessTransaction(136      this.dialect,137      this.session,138      this.schema,139      this.nestedIndex + 1140    );141    await tx.execute(sql.raw(`savepoint ${savepointName}`));142    try {143      const result = await transaction(tx);144      await tx.execute(sql.raw(`release savepoint ${savepointName}`));145      return result;146    } catch (err) {147      await tx.execute(sql.raw(`rollback to savepoint ${savepointName}`));148      throw err;149    }150  }151}152export {153  TiDBServerlessPreparedQuery,154  TiDBServerlessSession,155  TiDBServerlessTransaction156};157//# sourceMappingURL=session.js.map