From b958f8461f791356d63933be97e42c9c5e7f69bf Mon Sep 17 00:00:00 2001
From: Manfred Touron <m@42.am>
Date: Wed, 28 Feb 2018 16:54:32 +0100
Subject: [PATCH] Fix  commands when using MySQL driver
 ([#45](https://github.com/moul/sshportal/issues/45))

---
 CHANGELOG.md                    |  1 +
 examples/integration/_client.sh |  4 ++++
 shell.go                        | 42 +++++++++++++++++++++++++++------
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9a73773..16ee1e7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 * The default created user now has the same username as the user starting sshportal (was hardcoded "admin")
 * Add Telnet support
 * Add TTY audit feature ([#23](https://github.com/moul/sshportal/issues/23)) by [@sabban](https://github.com/sabban)
+* Fix `--assign-*` commands when using MySQL driver ([#45](https://github.com/moul/sshportal/issues/45))
 
 ## v1.7.1 (2018-01-03)
 
diff --git a/examples/integration/_client.sh b/examples/integration/_client.sh
index 1b383e2..77c1a88 100755
--- a/examples/integration/_client.sh
+++ b/examples/integration/_client.sh
@@ -47,6 +47,10 @@ ssh sshportal -l admin host create test42
 ssh sshportal -l admin host create --name=testtest --comment=test --password=test test@test.test
 ssh sshportal -l admin host create --group=hg1 --group=hg2 hostwithgroups.org
 ssh sshportal -l admin host inspect example test42 testtest hostwithgroups
+ssh sshportal -l admin host update --assign-group=hg1 test42
+ssh sshportal -l admin host update --unassign-group=hg1 test42
+ssh sshportal -l admin host update --assign-group=hg1 test42
+ssh sshportal -l admin host update --assign-group=hg2 --unassign-group=hg2 test42
 ssh sshportal -l admin host ls
 
 # backup/restore
diff --git a/shell.go b/shell.go
index 178d67c..bc480b6 100644
--- a/shell.go
+++ b/shell.go
@@ -271,10 +271,16 @@ GLOBAL OPTIONS:
 								tx.Rollback()
 								return err
 							}
-							if err := model.Association("UserGroups").Append(&appendUserGroups).Delete(deleteUserGroups).Error; err != nil {
+							if err := model.Association("UserGroups").Append(&appendUserGroups).Error; err != nil {
 								tx.Rollback()
 								return err
 							}
+							if len(deleteUserGroups) > 0 {
+								if err := model.Association("UserGroups").Delete(deleteUserGroups).Error; err != nil {
+									tx.Rollback()
+									return err
+								}
+							}
 
 							var appendHostGroups []HostGroup
 							var deleteHostGroups []HostGroup
@@ -286,10 +292,16 @@ GLOBAL OPTIONS:
 								tx.Rollback()
 								return err
 							}
-							if err := model.Association("HostGroups").Append(&appendHostGroups).Delete(deleteHostGroups).Error; err != nil {
+							if err := model.Association("HostGroups").Append(&appendHostGroups).Error; err != nil {
 								tx.Rollback()
 								return err
 							}
+							if len(deleteHostGroups) > 0 {
+								if err := model.Association("HostGroups").Delete(deleteHostGroups).Error; err != nil {
+									tx.Rollback()
+									return err
+								}
+							}
 						}
 
 						return tx.Commit().Error
@@ -898,10 +910,16 @@ GLOBAL OPTIONS:
 								tx.Rollback()
 								return err
 							}
-							if err := model.Association("Groups").Append(&appendGroups).Delete(deleteGroups).Error; err != nil {
+							if err := model.Association("Groups").Append(&appendGroups).Error; err != nil {
 								tx.Rollback()
 								return err
 							}
+							if len(deleteGroups) > 0 {
+								if err := model.Association("Groups").Delete(deleteGroups).Error; err != nil {
+									tx.Rollback()
+									return err
+								}
+							}
 						}
 
 						return tx.Commit().Error
@@ -1525,11 +1543,16 @@ GLOBAL OPTIONS:
 								tx.Rollback()
 								return err
 							}
-							if err := model.Association("Groups").Append(&appendGroups).Delete(deleteGroups).Error; err != nil {
+							if err := model.Association("Groups").Append(&appendGroups).Error; err != nil {
 								tx.Rollback()
 								return err
 							}
-
+							if len(deleteGroups) > 0 {
+								if err := model.Association("Groups").Delete(deleteGroups).Error; err != nil {
+									tx.Rollback()
+									return err
+								}
+							}
 							var appendRoles []UserRole
 							if err := UserRolesByIdentifiers(db, c.StringSlice("assign-role")).Find(&appendRoles).Error; err != nil {
 								tx.Rollback()
@@ -1540,12 +1563,17 @@ GLOBAL OPTIONS:
 								tx.Rollback()
 								return err
 							}
-							if err := model.Association("Roles").Append(&appendRoles).Delete(deleteRoles).Error; err != nil {
+							if err := model.Association("Roles").Append(&appendRoles).Error; err != nil {
 								tx.Rollback()
 								return err
 							}
+							if len(deleteRoles) > 0 {
+								if err := model.Association("Roles").Delete(deleteRoles).Error; err != nil {
+									tx.Rollback()
+									return err
+								}
+							}
 						}
-
 						return tx.Commit().Error
 					},
 				},