fix(core): fix leaking raw fragments cache

packages/core/src/utils/RawQueryFragment.ts

@@ -6,6 +6,7 @@ export class RawQueryFragment {
 
   static #rawQueryCache = new Map<string, RawQueryFragment>();
   static #index = 0;
+  static cloneRegistry?: Set<string>;
 
   #used = false;
   readonly #key: string;
@@ -44,6 +45,7 @@ export class RawQueryFragment {
   }
 
   clone(): RawQueryFragment {
+    RawQueryFragment.cloneRegistry?.add(this.#key);
     return new RawQueryFragment(this.sql, this.params);
   }
 
@@ -54,7 +56,11 @@ export class RawQueryFragment {
     return this.#rawQueryCache.size;
   }
 
-  static isKnownFragment(key: string) {
+  static isKnownFragment(key: string | RawQueryFragment) {
+    if (key instanceof RawQueryFragment) {
+      return true;
+    }
+
     return this.#rawQueryCache.has(key);
   }
 
@@ -66,12 +72,16 @@ export class RawQueryFragment {
     const raw = this.#rawQueryCache.get(key);
 
     if (raw && cleanup) {
-      this.#rawQueryCache.delete(key);
+      this.remove(key);
     }
 
     return raw;
   }
 
+  static remove(key: string) {
+    this.#rawQueryCache.delete(key);
+  }
+
   /* istanbul ignore next */
   /** @ignore */
   [inspect.custom]() {

packages/knex/src/AbstractSqlDriver.ts

@@ -898,6 +898,10 @@ export abstract class AbstractSqlDriver<Connection extends AbstractSqlConnection
       for (const field of targetFields) {
         const f = field.toString();
         fields.unshift(f.includes('.') ? field as string : `${targetAlias}.${f}`);
+
+        if (RawQueryFragment.isKnownFragment(field as string)) {
+          qb.rawFragments.add(f);
+        }
       }
 
       // we need to handle 1:1 owner auto-joins explicitly, as the QB type is the pivot table, not the target
@@ -926,6 +930,7 @@ export abstract class AbstractSqlDriver<Connection extends AbstractSqlConnection
 
     const res = owners.length ? await this.rethrow(qb.execute('all', { mergeResults: false, mapResults: false })) : [];
     const items = res.map((row: Dictionary) => super.mapResult(row, prop.targetMeta));
+    qb.clearRawFragmentsCache();
 
     const map: Dictionary<T[]> = {};
     const pkProps = ownerMeta.getPrimaryProps();

packages/knex/src/query/QueryBuilder.ts

@@ -95,6 +95,8 @@ export class QueryBuilder<T extends object = AnyEntity> {
   __populateWhere?: ObjectQuery<T> | PopulateHint | `${PopulateHint}`;
   /** @internal */
   _populateMap: Dictionary<string> = {};
+  /** @internal */
+  readonly rawFragments = new Set<string>();
 
   private aliasCounter = 0;
   private flags: Set<QueryFlag> = new Set([QueryFlag.CONVERT_CUSTOM_TYPES]);
@@ -629,10 +631,19 @@ export class QueryBuilder<T extends object = AnyEntity> {
     }
 
     this.helper.finalize(type, qb, this.mainAlias.metadata, this._data, this._returning);
+    this.clearRawFragmentsCache();
 
     return qb;
   }
 
+  /**
+   * @internal
+   */
+  clearRawFragmentsCache(): void {
+    this.rawFragments.forEach(key => RawQueryFragment.remove(key));
+    this.rawFragments.clear();
+  }
+
   /**
    * Returns the query with parameters as wildcards.
    */
@@ -920,26 +931,29 @@ export class QueryBuilder<T extends object = AnyEntity> {
       return qb;
     }
 
-    Object.assign(qb, this);
     reset = reset || [];
 
     // clone array/object properties
     const properties = [
       'flags', '_populate', '_populateWhere', '__populateWhere', '_populateMap', '_joins', '_joinedProps', '_cond', '_data', '_orderBy',
       '_schema', '_indexHint', '_cache', 'subQueries', 'lockMode', 'lockTables', '_groupBy', '_having', '_returning',
-      '_comments', '_hintComments',
-    ] as const;
+      '_comments', '_hintComments', 'rawFragments',
+    ];
+
+    RawQueryFragment.cloneRegistry = this.rawFragments;
 
-    for (const prop of properties) {
+    for (const prop of Object.keys(this)) {
       if (reset.includes(prop)) {
         continue;
       }
 
-      (qb as any)[prop] = Utils.copy(this[prop]);
+      (qb as any)[prop] = properties.includes(prop) ? Utils.copy(this[prop as keyof this]) : this[prop as keyof this];
     }
 
+    delete RawQueryFragment.cloneRegistry;
+
     /* istanbul ignore else */
-    if (this._fields) {
+    if (this._fields && !reset.includes('_fields')) {
       qb._fields = [...this._fields];
     }
 
@@ -1411,7 +1425,7 @@ export class QueryBuilder<T extends object = AnyEntity> {
 
   private wrapPaginateSubQuery(meta: EntityMetadata): void {
     const pks = this.prepareFields(meta.primaryKeys, 'sub-query') as string[];
-    const subQuery = this.clone(['_orderBy', '_cond']).select(pks).groupBy(pks).limit(this._limit!);
+    const subQuery = this.clone(['_orderBy', '_fields']).select(pks).groupBy(pks).limit(this._limit!);
     // revert the on conditions added via populateWhere, we want to apply those only once
     Object.values(subQuery._joins).forEach(join => join.cond = join.cond_ ?? {});
 
@@ -1428,6 +1442,7 @@ export class QueryBuilder<T extends object = AnyEntity> {
         for (const [field, direction] of Object.entries(orderMap)) {
           if (RawQueryFragment.isKnownFragment(field)) {
             const rawField = RawQueryFragment.getKnownFragment(field, false)!;
+            this.rawFragments.add(field);
             orderBy.push({ [rawField.clone() as any]: direction });
             continue;
           }
@@ -1459,11 +1474,18 @@ export class QueryBuilder<T extends object = AnyEntity> {
             return field.__as === prop;
           }
 
+          if (field instanceof RawQueryFragment) {
+            // not perfect, but should work most of the time, ideally we should check only the alias (`... as alias`)
+            return field.sql.includes(prop);
+          }
+
           // not perfect, but should work most of the time, ideally we should check only the alias (`... as alias`)
           return field.toString().includes(prop);
         });
 
-        if (field) {
+        if (field instanceof RawQueryFragment) {
+          knexQuery.select(this.platform.formatQuery(field.sql, field.params));
+        } else if (field) {
           knexQuery.select(field as string);
         }
       });

tests/EntityManager.mysql.test.ts

@@ -23,6 +23,7 @@ import {
   PopulateHint,
   raw,
   ref,
+  RawQueryFragment,
 } from '@mikro-orm/core';
 import { MySqlDriver, MySqlConnection, ScalarReference } from '@mikro-orm/mysql';
 import { Address2, Author2, Book2, BookTag2, FooBar2, FooBaz2, Publisher2, PublisherType, Test2 } from './entities-sql';
@@ -39,6 +40,7 @@ describe('EntityManagerMySql', () => {
   beforeAll(async () => orm = await initORMMySql('mysql', {}, true));
   beforeEach(async () => orm.schema.clearDatabase());
   afterEach(() => {
+    expect(RawQueryFragment.checkCacheSize()).toBe(0);
     orm.config.set('debug', false);
     Author2Subscriber.log.length = 0;
     EverythingSubscriber.log.length = 0;

tests/EntityManager.postgre.test.ts

@@ -1,5 +1,13 @@
 import { v4 } from 'uuid';
-import { AnyEntity, ChangeSet, DefaultLogger, EventSubscriber, FilterQuery, FlushEventArgs } from '@mikro-orm/core';
+import {
+  AnyEntity,
+  ChangeSet,
+  DefaultLogger,
+  EventSubscriber,
+  FilterQuery,
+  FlushEventArgs,
+  RawQueryFragment,
+} from '@mikro-orm/core';
 import {
   ChangeSetType,
   Collection,
@@ -72,6 +80,7 @@ describe('EntityManagerPostgre', () => {
 
   beforeAll(async () => orm = await initORMPostgreSql());
   beforeEach(async () => orm.schema.clearDatabase());
+  afterEach(() => expect(RawQueryFragment.checkCacheSize()).toBe(0));
   afterAll(async () => {
     await orm.schema.dropDatabase();
     await orm.close(true);
@@ -1136,8 +1145,8 @@ describe('EntityManagerPostgre', () => {
     book2.tags.add(tag1, tag2, tag5);
     book3.tags.add(tag2, tag4, tag5);
 
-    await orm.em.persist(book1);
-    await orm.em.persist(book2);
+    orm.em.persist(book1);
+    orm.em.persist(book2);
     await orm.em.persistAndFlush(book3);
 
     expect(typeof tag1.id).toBe('bigint');

tests/QueryBuilder.test.ts

@@ -1,5 +1,14 @@
 import { inspect } from 'util';
-import { LockMode, MikroORM, QueryFlag, QueryOrder, raw, sql, UnderscoreNamingStrategy } from '@mikro-orm/core';
+import {
+  LockMode,
+  MikroORM,
+  QueryFlag,
+  QueryOrder,
+  raw,
+  RawQueryFragment,
+  sql,
+  UnderscoreNamingStrategy,
+} from '@mikro-orm/core';
 import { CriteriaNode, QueryBuilder, PostgreSqlDriver } from '@mikro-orm/postgresql';
 import { MySqlDriver } from '@mikro-orm/mysql';
 import { v4 } from 'uuid';
@@ -24,6 +33,7 @@ describe('QueryBuilder', () => {
       },
     }, true);
   });
+  afterEach(() => expect(RawQueryFragment.checkCacheSize()).toBe(0));
   afterAll(async () => {
     await orm.schema.dropDatabase();
     await orm.close(true);

tests/issues/GHx6.test.ts

@@ -134,8 +134,8 @@ test('qb.joinAndSelect', async () => {
       [raw('similarity(u."name", ?)', ['def'])]: QueryOrder.DESC_NULLS_LAST,
     })
     .limit(100)
-    .offset(0).
-    getFormattedQuery();
+    .offset(0)
+    .getFormattedQuery();
   expect(query).toMatch('select `u`.*, `a`.`id` as `a__id`, `a`.`DateCompleted` as `a__DateCompleted` ' +
     'from `tag` as `u` ' +
     'left join `tag_jobs` as `t1` on `u`.`id` = `t1`.`tag_id` ' +