feat(belongsToMany): add option to disable pivotPrimaryKey

fix #287
adonisjs · Jul 14, 2018 · 0cb13bf · 0cb13bf
1 parent 7aaf45e
commit 0cb13bf
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 4 deletions.
diff --git a/src/Lucid/Model/PivotModel.js b/src/Lucid/Model/PivotModel.js
@@ -39,9 +39,11 @@ class PivotModel extends BaseModel {
   _instantiate () {
     super._instantiate()
     this.__setters__.push('$withTimestamps')
+    this.__setters__.push('$primaryKey')
     this.__setters__.push('$table')
     this.__setters__.push('$connection')
     this.$withTimestamps = false
+    this.$primaryKey = 'id'
   }
 
   /**
@@ -162,9 +164,11 @@ class PivotModel extends BaseModel {
       query.transacting(trx)
     }
 
-    const result = await query
-      .returning('id')
-      .insert(this.$attributes)
+    if (this.$primaryKey) {
+      query.returning(typeof (this.$primaryKey) === 'string' ? this.$primaryKey : 'id')
+    }
+
+    const result = await query.insert(this.$attributes)
 
     this.primaryKeyValue = result[0]
     this.$persisted = true

diff --git a/src/Lucid/Relations/BelongsToMany.js b/src/Lucid/Relations/BelongsToMany.js
@@ -78,7 +78,8 @@ class BelongsToMany extends BaseRelation {
     this._pivot = {
       table: util.makePivotTableName(parentInstance.constructor.name, relatedModel.name),
       withTimestamps: false,
-      withFields: []
+      withFields: [],
+      pivotPrimaryKey: 'id'
     }
 
     this._relatedFields = []
@@ -349,6 +350,7 @@ class BelongsToMany extends BaseRelation {
       pivotModel.$table = this.$pivotTable
       pivotModel.$connection = this.RelatedModel.connection
       pivotModel.$withTimestamps = this._pivot.withTimestamps
+      pivotModel.$primaryKey = this._pivot.pivotPrimaryKey
     }
 
     /**
@@ -473,6 +475,25 @@ class BelongsToMany extends BaseRelation {
     return this
   }
 
+  /**
+   * Define the primary key to be selected for the
+   * pivot table.
+   *
+   * @method pivotPrimaryKey
+   *
+   * @param  {String}        key
+   *
+   * @chainable
+   */
+  pivotPrimaryKey (key) {
+    if (this._PivotModel) {
+      throw CE.ModelRelationException.pivotModelIsDefined('pivotPrimaryKey')
+    }
+
+    this._pivot.pivotPrimaryKey = key
+    return this
+  }
+
   /**
    * Make sure `created_at` and `updated_at` timestamps
    * are being used

diff --git a/test/unit/lucid-belongs-to-many.spec.js b/test/unit/lucid-belongs-to-many.spec.js
@@ -2616,4 +2616,28 @@ test.group('Relations | Belongs To Many', (group) => {
 
     assert.equal(helpers.formatQuery(postsQuery.sql), helpers.formatQuery('select "posts".*, "post_user"."post_id" as "pivot_post_id", "post_user"."user_id" as "pivot_user_id" from "posts" inner join "post_user" on "posts"."id" = "post_user"."post_id" where "post_user"."user_id" in (?) and "post_user"."is_published" = ?'))
   })
+
+  test('on save do not add returning statement when withPrimaryKey is false (pg only)', async (assert) => {
+    class Post extends Model {
+    }
+
+    class User extends Model {
+      posts () {
+        return this.belongsToMany(Post).pivotPrimaryKey(false)
+      }
+    }
+
+    User._bootIfNotBooted()
+    Post._bootIfNotBooted()
+
+    let postsQuery = null
+    ioc.use('Database').on('query', (query) => (postsQuery = query))
+
+    await ioc.use('Database').table('users').insert({ username: 'virk' })
+
+    const user = await User.find(1)
+    await user.posts().create({ title: 'Adonis 101' })
+
+    assert.equal(postsQuery.sql, helpers.formatQuery('insert into "post_user" ("post_id", "user_id") values (?, ?)'))
+  })
 })