From c19708667197baf57e6fe22719b701b44382548a Mon Sep 17 00:00:00 2001
From: jlenon7 <lenonSec7@gmail.com>
Date: Fri, 7 Nov 2025 14:10:13 -0300
Subject: [PATCH] feat(model): add way to set transaction for models

---
 package-lock.json                        |  4 ++--
 package.json                             |  2 +-
 src/models/builders/ModelQueryBuilder.ts | 19 ++++++++++++++++++-
 3 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 695e615..cf5fd21 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "@athenna/database",
-  "version": "5.30.0",
+  "version": "5.31.0",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@athenna/database",
-      "version": "5.30.0",
+      "version": "5.31.0",
       "license": "MIT",
       "dependencies": {
         "@faker-js/faker": "^8.4.1"
diff --git a/package.json b/package.json
index ba27d79..62d1320 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@athenna/database",
-  "version": "5.30.0",
+  "version": "5.31.0",
   "description": "The Athenna database handler for SQL/NoSQL.",
   "license": "MIT",
   "author": "João Lenon <lenon@athenna.io>",
diff --git a/src/models/builders/ModelQueryBuilder.ts b/src/models/builders/ModelQueryBuilder.ts
index 7872c88..25200b6 100644
--- a/src/models/builders/ModelQueryBuilder.ts
+++ b/src/models/builders/ModelQueryBuilder.ts
@@ -26,13 +26,14 @@ import type { Driver } from '#src/database/drivers/Driver'
 import { QueryBuilder } from '#src/database/builders/QueryBuilder'
 import type { ModelSchema } from '#src/models/schemas/ModelSchema'
 import { ModelGenerator } from '#src/models/factories/ModelGenerator'
+import type { Transaction } from '#src/database/transactions/Transaction'
 import { UniqueValueException } from '#src/exceptions/UniqueValueException'
 import { HasOneRelation } from '#src/models/relations/HasOne/HasOneRelation'
 import { NotFoundDataException } from '#src/exceptions/NotFoundDataException'
 import { HasManyRelation } from '#src/models/relations/HasMany/HasManyRelation'
 import { NullableValueException } from '#src/exceptions/NullableValueException'
 import { BelongsToRelation } from '#src/models/relations/BelongsTo/BelongsToRelation'
-import { BelongsToManyRelation } from '../relations/BelongsToMany/BelongsToManyRelation.js'
+import { BelongsToManyRelation } from '#src/models/relations/BelongsToMany/BelongsToManyRelation'
 
 export class ModelQueryBuilder<
   M extends BaseModel = any,
@@ -73,6 +74,22 @@ export class ModelQueryBuilder<
     this.setPrimaryKey(this.primaryKeyName)
   }
 
+  /**
+   * Define a transaction to be used by the model query builder.
+   */
+  public setTransaction(trx: Transaction) {
+    return this.setDriver(trx.driver, this.Model.table())
+  }
+
+  /**
+   * Set a different driver to the model query builder.
+   */
+  public setDriver(driver: Driver, tableName?: string) {
+    super.setDriver(driver, tableName)
+
+    return this
+  }
+
   /**
    * Calculate the average of a given column.
    */