session.js 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201
  1. import {
  2. BeginTransactionCommand,
  3. CommitTransactionCommand,
  4. ExecuteStatementCommand,
  5. RollbackTransactionCommand
  6. } from "@aws-sdk/client-rds-data";
  7. import { NoopCache } from "../../cache/core/cache.js";
  8. import { entityKind } from "../../entity.js";
  9. import {
  10. PgPreparedQuery,
  11. PgSession,
  12. PgTransaction
  13. } from "../../pg-core/index.js";
  14. import { fillPlaceholders, sql } from "../../sql/sql.js";
  15. import { mapResultRow } from "../../utils.js";
  16. import { getValueFromDataApi, toValueParam } from "../common/index.js";
  17. class AwsDataApiPreparedQuery extends PgPreparedQuery {
  18. constructor(client, queryString, params, typings, options, cache, queryMetadata, cacheConfig, fields, transactionId, _isResponseInArrayMode, customResultMapper) {
  19. super({ sql: queryString, params }, cache, queryMetadata, cacheConfig);
  20. this.client = client;
  21. this.queryString = queryString;
  22. this.params = params;
  23. this.typings = typings;
  24. this.options = options;
  25. this.fields = fields;
  26. this.transactionId = transactionId;
  27. this._isResponseInArrayMode = _isResponseInArrayMode;
  28. this.customResultMapper = customResultMapper;
  29. this.rawQuery = new ExecuteStatementCommand({
  30. sql: queryString,
  31. parameters: [],
  32. secretArn: options.secretArn,
  33. resourceArn: options.resourceArn,
  34. database: options.database,
  35. transactionId,
  36. includeResultMetadata: !fields && !customResultMapper
  37. });
  38. }
  39. static [entityKind] = "AwsDataApiPreparedQuery";
  40. rawQuery;
  41. async execute(placeholderValues = {}) {
  42. const { fields, joinsNotNullableMap, customResultMapper } = this;
  43. const result = await this.values(placeholderValues);
  44. if (!fields && !customResultMapper) {
  45. const { columnMetadata, rows } = result;
  46. if (!columnMetadata) {
  47. return result;
  48. }
  49. const mappedRows = rows.map((sourceRow) => {
  50. const row = {};
  51. for (const [index, value] of sourceRow.entries()) {
  52. const metadata = columnMetadata[index];
  53. if (!metadata) {
  54. throw new Error(
  55. `Unexpected state: no column metadata found for index ${index}. Please report this issue on GitHub: https://github.com/drizzle-team/drizzle-orm/issues/new/choose`
  56. );
  57. }
  58. if (!metadata.name) {
  59. throw new Error(
  60. `Unexpected state: no column name for index ${index} found in the column metadata. Please report this issue on GitHub: https://github.com/drizzle-team/drizzle-orm/issues/new/choose`
  61. );
  62. }
  63. row[metadata.name] = value;
  64. }
  65. return row;
  66. });
  67. return Object.assign(result, { rows: mappedRows });
  68. }
  69. return customResultMapper ? customResultMapper(result.rows) : result.rows.map((row) => mapResultRow(fields, row, joinsNotNullableMap));
  70. }
  71. async all(placeholderValues) {
  72. const result = await this.execute(placeholderValues);
  73. if (!this.fields && !this.customResultMapper) {
  74. return result.rows;
  75. }
  76. return result;
  77. }
  78. async values(placeholderValues = {}) {
  79. const params = fillPlaceholders(this.params, placeholderValues ?? {});
  80. this.rawQuery.input.parameters = params.map((param, index) => ({
  81. name: `${index + 1}`,
  82. ...toValueParam(param, this.typings[index])
  83. }));
  84. this.options.logger?.logQuery(this.rawQuery.input.sql, this.rawQuery.input.parameters);
  85. const result = await this.queryWithCache(this.queryString, params, async () => {
  86. return await this.client.send(this.rawQuery);
  87. });
  88. const rows = result.records?.map((row) => {
  89. return row.map((field) => getValueFromDataApi(field));
  90. }) ?? [];
  91. return {
  92. ...result,
  93. rows
  94. };
  95. }
  96. /** @internal */
  97. mapResultRows(records, columnMetadata) {
  98. return records.map((record) => {
  99. const row = {};
  100. for (const [index, field] of record.entries()) {
  101. const { name } = columnMetadata[index];
  102. row[name ?? index] = getValueFromDataApi(field);
  103. }
  104. return row;
  105. });
  106. }
  107. /** @internal */
  108. isResponseInArrayMode() {
  109. return this._isResponseInArrayMode;
  110. }
  111. }
  112. class AwsDataApiSession extends PgSession {
  113. constructor(client, dialect, schema, options, transactionId) {
  114. super(dialect);
  115. this.client = client;
  116. this.schema = schema;
  117. this.options = options;
  118. this.transactionId = transactionId;
  119. this.rawQuery = {
  120. secretArn: options.secretArn,
  121. resourceArn: options.resourceArn,
  122. database: options.database
  123. };
  124. this.cache = options.cache ?? new NoopCache();
  125. }
  126. static [entityKind] = "AwsDataApiSession";
  127. /** @internal */
  128. rawQuery;
  129. cache;
  130. prepareQuery(query, fields, name, isResponseInArrayMode, customResultMapper, queryMetadata, cacheConfig, transactionId) {
  131. return new AwsDataApiPreparedQuery(
  132. this.client,
  133. query.sql,
  134. query.params,
  135. query.typings ?? [],
  136. this.options,
  137. this.cache,
  138. queryMetadata,
  139. cacheConfig,
  140. fields,
  141. transactionId ?? this.transactionId,
  142. isResponseInArrayMode,
  143. customResultMapper
  144. );
  145. }
  146. execute(query) {
  147. return this.prepareQuery(
  148. this.dialect.sqlToQuery(query),
  149. void 0,
  150. void 0,
  151. false,
  152. void 0,
  153. void 0,
  154. void 0,
  155. this.transactionId
  156. ).execute();
  157. }
  158. async transaction(transaction, config) {
  159. const { transactionId } = await this.client.send(new BeginTransactionCommand(this.rawQuery));
  160. const session = new AwsDataApiSession(this.client, this.dialect, this.schema, this.options, transactionId);
  161. const tx = new AwsDataApiTransaction(this.dialect, session, this.schema);
  162. if (config) {
  163. await tx.setTransaction(config);
  164. }
  165. try {
  166. const result = await transaction(tx);
  167. await this.client.send(new CommitTransactionCommand({ ...this.rawQuery, transactionId }));
  168. return result;
  169. } catch (e) {
  170. await this.client.send(new RollbackTransactionCommand({ ...this.rawQuery, transactionId }));
  171. throw e;
  172. }
  173. }
  174. }
  175. class AwsDataApiTransaction extends PgTransaction {
  176. static [entityKind] = "AwsDataApiTransaction";
  177. async transaction(transaction) {
  178. const savepointName = `sp${this.nestedIndex + 1}`;
  179. const tx = new AwsDataApiTransaction(
  180. this.dialect,
  181. this.session,
  182. this.schema,
  183. this.nestedIndex + 1
  184. );
  185. await this.session.execute(sql.raw(`savepoint ${savepointName}`));
  186. try {
  187. const result = await transaction(tx);
  188. await this.session.execute(sql.raw(`release savepoint ${savepointName}`));
  189. return result;
  190. } catch (e) {
  191. await this.session.execute(sql.raw(`rollback to savepoint ${savepointName}`));
  192. throw e;
  193. }
  194. }
  195. }
  196. export {
  197. AwsDataApiPreparedQuery,
  198. AwsDataApiSession,
  199. AwsDataApiTransaction
  200. };
  201. //# sourceMappingURL=session.js.map