db.js 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331
  1. import { entityKind } from "../entity.js";
  2. import {
  3. PgDeleteBase,
  4. PgInsertBuilder,
  5. PgSelectBuilder,
  6. PgUpdateBuilder,
  7. QueryBuilder
  8. } from "./query-builders/index.js";
  9. import { SelectionProxyHandler } from "../selection-proxy.js";
  10. import { sql } from "../sql/sql.js";
  11. import { WithSubquery } from "../subquery.js";
  12. import { PgCountBuilder } from "./query-builders/count.js";
  13. import { RelationalQueryBuilder } from "./query-builders/query.js";
  14. import { PgRaw } from "./query-builders/raw.js";
  15. import { PgRefreshMaterializedView } from "./query-builders/refresh-materialized-view.js";
  16. class PgDatabase {
  17. constructor(dialect, session, schema) {
  18. this.dialect = dialect;
  19. this.session = session;
  20. this._ = schema ? {
  21. schema: schema.schema,
  22. fullSchema: schema.fullSchema,
  23. tableNamesMap: schema.tableNamesMap,
  24. session
  25. } : {
  26. schema: void 0,
  27. fullSchema: {},
  28. tableNamesMap: {},
  29. session
  30. };
  31. this.query = {};
  32. if (this._.schema) {
  33. for (const [tableName, columns] of Object.entries(this._.schema)) {
  34. this.query[tableName] = new RelationalQueryBuilder(
  35. schema.fullSchema,
  36. this._.schema,
  37. this._.tableNamesMap,
  38. schema.fullSchema[tableName],
  39. columns,
  40. dialect,
  41. session
  42. );
  43. }
  44. }
  45. this.$cache = { invalidate: async (_params) => {
  46. } };
  47. }
  48. static [entityKind] = "PgDatabase";
  49. query;
  50. /**
  51. * Creates a subquery that defines a temporary named result set as a CTE.
  52. *
  53. * It is useful for breaking down complex queries into simpler parts and for reusing the result set in subsequent parts of the query.
  54. *
  55. * See docs: {@link https://orm.drizzle.team/docs/select#with-clause}
  56. *
  57. * @param alias The alias for the subquery.
  58. *
  59. * Failure to provide an alias will result in a DrizzleTypeError, preventing the subquery from being referenced in other queries.
  60. *
  61. * @example
  62. *
  63. * ```ts
  64. * // Create a subquery with alias 'sq' and use it in the select query
  65. * const sq = db.$with('sq').as(db.select().from(users).where(eq(users.id, 42)));
  66. *
  67. * const result = await db.with(sq).select().from(sq);
  68. * ```
  69. *
  70. * To select arbitrary SQL values as fields in a CTE and reference them in other CTEs or in the main query, you need to add aliases to them:
  71. *
  72. * ```ts
  73. * // Select an arbitrary SQL value as a field in a CTE and reference it in the main query
  74. * const sq = db.$with('sq').as(db.select({
  75. * name: sql<string>`upper(${users.name})`.as('name'),
  76. * })
  77. * .from(users));
  78. *
  79. * const result = await db.with(sq).select({ name: sq.name }).from(sq);
  80. * ```
  81. */
  82. $with = (alias, selection) => {
  83. const self = this;
  84. const as = (qb) => {
  85. if (typeof qb === "function") {
  86. qb = qb(new QueryBuilder(self.dialect));
  87. }
  88. return new Proxy(
  89. new WithSubquery(
  90. qb.getSQL(),
  91. selection ?? ("getSelectedFields" in qb ? qb.getSelectedFields() ?? {} : {}),
  92. alias,
  93. true
  94. ),
  95. new SelectionProxyHandler({ alias, sqlAliasedBehavior: "alias", sqlBehavior: "error" })
  96. );
  97. };
  98. return { as };
  99. };
  100. $count(source, filters) {
  101. return new PgCountBuilder({ source, filters, session: this.session });
  102. }
  103. $cache;
  104. /**
  105. * Incorporates a previously defined CTE (using `$with`) into the main query.
  106. *
  107. * This method allows the main query to reference a temporary named result set.
  108. *
  109. * See docs: {@link https://orm.drizzle.team/docs/select#with-clause}
  110. *
  111. * @param queries The CTEs to incorporate into the main query.
  112. *
  113. * @example
  114. *
  115. * ```ts
  116. * // Define a subquery 'sq' as a CTE using $with
  117. * const sq = db.$with('sq').as(db.select().from(users).where(eq(users.id, 42)));
  118. *
  119. * // Incorporate the CTE 'sq' into the main query and select from it
  120. * const result = await db.with(sq).select().from(sq);
  121. * ```
  122. */
  123. with(...queries) {
  124. const self = this;
  125. function select(fields) {
  126. return new PgSelectBuilder({
  127. fields: fields ?? void 0,
  128. session: self.session,
  129. dialect: self.dialect,
  130. withList: queries
  131. });
  132. }
  133. function selectDistinct(fields) {
  134. return new PgSelectBuilder({
  135. fields: fields ?? void 0,
  136. session: self.session,
  137. dialect: self.dialect,
  138. withList: queries,
  139. distinct: true
  140. });
  141. }
  142. function selectDistinctOn(on, fields) {
  143. return new PgSelectBuilder({
  144. fields: fields ?? void 0,
  145. session: self.session,
  146. dialect: self.dialect,
  147. withList: queries,
  148. distinct: { on }
  149. });
  150. }
  151. function update(table) {
  152. return new PgUpdateBuilder(table, self.session, self.dialect, queries);
  153. }
  154. function insert(table) {
  155. return new PgInsertBuilder(table, self.session, self.dialect, queries);
  156. }
  157. function delete_(table) {
  158. return new PgDeleteBase(table, self.session, self.dialect, queries);
  159. }
  160. return { select, selectDistinct, selectDistinctOn, update, insert, delete: delete_ };
  161. }
  162. select(fields) {
  163. return new PgSelectBuilder({
  164. fields: fields ?? void 0,
  165. session: this.session,
  166. dialect: this.dialect
  167. });
  168. }
  169. selectDistinct(fields) {
  170. return new PgSelectBuilder({
  171. fields: fields ?? void 0,
  172. session: this.session,
  173. dialect: this.dialect,
  174. distinct: true
  175. });
  176. }
  177. selectDistinctOn(on, fields) {
  178. return new PgSelectBuilder({
  179. fields: fields ?? void 0,
  180. session: this.session,
  181. dialect: this.dialect,
  182. distinct: { on }
  183. });
  184. }
  185. /**
  186. * Creates an update query.
  187. *
  188. * Calling this method without `.where()` clause will update all rows in a table. The `.where()` clause specifies which rows should be updated.
  189. *
  190. * Use `.set()` method to specify which values to update.
  191. *
  192. * See docs: {@link https://orm.drizzle.team/docs/update}
  193. *
  194. * @param table The table to update.
  195. *
  196. * @example
  197. *
  198. * ```ts
  199. * // Update all rows in the 'cars' table
  200. * await db.update(cars).set({ color: 'red' });
  201. *
  202. * // Update rows with filters and conditions
  203. * await db.update(cars).set({ color: 'red' }).where(eq(cars.brand, 'BMW'));
  204. *
  205. * // Update with returning clause
  206. * const updatedCar: Car[] = await db.update(cars)
  207. * .set({ color: 'red' })
  208. * .where(eq(cars.id, 1))
  209. * .returning();
  210. * ```
  211. */
  212. update(table) {
  213. return new PgUpdateBuilder(table, this.session, this.dialect);
  214. }
  215. /**
  216. * Creates an insert query.
  217. *
  218. * Calling this method will create new rows in a table. Use `.values()` method to specify which values to insert.
  219. *
  220. * See docs: {@link https://orm.drizzle.team/docs/insert}
  221. *
  222. * @param table The table to insert into.
  223. *
  224. * @example
  225. *
  226. * ```ts
  227. * // Insert one row
  228. * await db.insert(cars).values({ brand: 'BMW' });
  229. *
  230. * // Insert multiple rows
  231. * await db.insert(cars).values([{ brand: 'BMW' }, { brand: 'Porsche' }]);
  232. *
  233. * // Insert with returning clause
  234. * const insertedCar: Car[] = await db.insert(cars)
  235. * .values({ brand: 'BMW' })
  236. * .returning();
  237. * ```
  238. */
  239. insert(table) {
  240. return new PgInsertBuilder(table, this.session, this.dialect);
  241. }
  242. /**
  243. * Creates a delete query.
  244. *
  245. * Calling this method without `.where()` clause will delete all rows in a table. The `.where()` clause specifies which rows should be deleted.
  246. *
  247. * See docs: {@link https://orm.drizzle.team/docs/delete}
  248. *
  249. * @param table The table to delete from.
  250. *
  251. * @example
  252. *
  253. * ```ts
  254. * // Delete all rows in the 'cars' table
  255. * await db.delete(cars);
  256. *
  257. * // Delete rows with filters and conditions
  258. * await db.delete(cars).where(eq(cars.color, 'green'));
  259. *
  260. * // Delete with returning clause
  261. * const deletedCar: Car[] = await db.delete(cars)
  262. * .where(eq(cars.id, 1))
  263. * .returning();
  264. * ```
  265. */
  266. delete(table) {
  267. return new PgDeleteBase(table, this.session, this.dialect);
  268. }
  269. refreshMaterializedView(view) {
  270. return new PgRefreshMaterializedView(view, this.session, this.dialect);
  271. }
  272. authToken;
  273. execute(query) {
  274. const sequel = typeof query === "string" ? sql.raw(query) : query.getSQL();
  275. const builtQuery = this.dialect.sqlToQuery(sequel);
  276. const prepared = this.session.prepareQuery(
  277. builtQuery,
  278. void 0,
  279. void 0,
  280. false
  281. );
  282. return new PgRaw(
  283. () => prepared.execute(void 0, this.authToken),
  284. sequel,
  285. builtQuery,
  286. (result) => prepared.mapResult(result, true)
  287. );
  288. }
  289. transaction(transaction, config) {
  290. return this.session.transaction(transaction, config);
  291. }
  292. }
  293. const withReplicas = (primary, replicas, getReplica = () => replicas[Math.floor(Math.random() * replicas.length)]) => {
  294. const select = (...args) => getReplica(replicas).select(...args);
  295. const selectDistinct = (...args) => getReplica(replicas).selectDistinct(...args);
  296. const selectDistinctOn = (...args) => getReplica(replicas).selectDistinctOn(...args);
  297. const $count = (...args) => getReplica(replicas).$count(...args);
  298. const _with = (...args) => getReplica(replicas).with(...args);
  299. const $with = (arg) => getReplica(replicas).$with(arg);
  300. const update = (...args) => primary.update(...args);
  301. const insert = (...args) => primary.insert(...args);
  302. const $delete = (...args) => primary.delete(...args);
  303. const execute = (...args) => primary.execute(...args);
  304. const transaction = (...args) => primary.transaction(...args);
  305. const refreshMaterializedView = (...args) => primary.refreshMaterializedView(...args);
  306. return {
  307. ...primary,
  308. update,
  309. insert,
  310. delete: $delete,
  311. execute,
  312. transaction,
  313. refreshMaterializedView,
  314. $primary: primary,
  315. $replicas: replicas,
  316. select,
  317. selectDistinct,
  318. selectDistinctOn,
  319. $count,
  320. $with,
  321. with: _with,
  322. get query() {
  323. return getReplica(replicas).query;
  324. }
  325. };
  326. };
  327. export {
  328. PgDatabase,
  329. withReplicas
  330. };
  331. //# sourceMappingURL=db.js.map