session.js 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. import { NoopCache } from "../cache/core/index.js";
  2. import { Column } from "../column.js";
  3. import { entityKind, is } from "../entity.js";
  4. import { NoopLogger } from "../logger.js";
  5. import {
  6. MySqlPreparedQuery,
  7. MySqlSession,
  8. MySqlTransaction
  9. } from "../mysql-core/session.js";
  10. import { fillPlaceholders, sql } from "../sql/sql.js";
  11. import { mapResultRow } from "../utils.js";
  12. const executeRawConfig = { fullResult: true };
  13. const queryConfig = { arrayMode: true };
  14. class 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. }
  70. class 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. returningIds
  97. );
  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.schema
  118. );
  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. }
  128. class 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 + 1
  140. );
  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. }
  152. export {
  153. TiDBServerlessPreparedQuery,
  154. TiDBServerlessSession,
  155. TiDBServerlessTransaction
  156. };
  157. //# sourceMappingURL=session.js.map