=== modified file 'dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataelement/DataElementGroup.java'
--- dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataelement/DataElementGroup.java	2011-05-05 21:14:56 +0000
+++ dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataelement/DataElementGroup.java	2011-05-30 14:56:19 +0000
@@ -60,6 +60,38 @@
     {
         this.name = name;
     }
+
+    // -------------------------------------------------------------------------
+    // Logic
+    // -------------------------------------------------------------------------
+
+    public void addDataElement( DataElement dataElement )
+    {
+        members.add( dataElement );
+        dataElement.getGroups().add( this );
+    }
+    
+    public void removeDataElement( DataElement dataElement )
+    {
+        members.remove( dataElement );
+        dataElement.getGroups().remove( this );
+    }
+    
+    public void updateDataElements( Set<DataElement> updates )
+    {
+        for ( DataElement dataElement : new HashSet<DataElement>( members ) )
+        {
+            if ( !updates.contains( dataElement ) )
+            {
+                removeDataElement( dataElement );
+            }
+        }
+        
+        for ( DataElement dataElement : updates )
+        {
+            addDataElement( dataElement );
+        }
+    }
     
     // -------------------------------------------------------------------------
     // hashCode and equals

=== modified file 'dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataset/DataSet.java'
--- dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataset/DataSet.java	2011-05-08 22:34:28 +0000
+++ dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataset/DataSet.java	2011-05-30 14:56:19 +0000
@@ -149,6 +149,34 @@
     // Logic
     // -------------------------------------------------------------------------
 
+    public void addOrganisationUnit( OrganisationUnit unit )
+    {
+        sources.add( unit );
+        unit.getDataSets().add( this );
+    }
+    
+    public void removeOrganisationUnit( OrganisationUnit unit )
+    {
+        sources.remove( unit );
+        unit.getDataSets().remove( this );
+    }
+    
+    public void updateOrganisationUnits( Set<OrganisationUnit> updates )
+    {
+        for ( OrganisationUnit unit : new HashSet<OrganisationUnit>( sources ) )
+        {
+            if ( !updates.contains( unit ) )
+            {
+                removeOrganisationUnit( unit );
+            }
+        }
+        
+        for ( OrganisationUnit unit : updates )
+        {
+            addOrganisationUnit( unit );
+        }
+    }
+    
     public boolean hasDataEntryForm()
     {
         return dataEntryForm != null;

=== modified file 'dhis-2/dhis-api/src/main/java/org/hisp/dhis/organisationunit/OrganisationUnit.java'
--- dhis-2/dhis-api/src/main/java/org/hisp/dhis/organisationunit/OrganisationUnit.java	2011-05-13 10:42:19 +0000
+++ dhis-2/dhis-api/src/main/java/org/hisp/dhis/organisationunit/OrganisationUnit.java	2011-05-30 14:56:19 +0000
@@ -173,6 +173,34 @@
     // Logic
     // -------------------------------------------------------------------------
     
+    public void addDataSet( DataSet dataSet )
+    {
+        dataSets.add( dataSet );
+        dataSet.getSources().add( this );
+    }
+    
+    public void removeDataSet( DataSet dataSet )
+    {
+        dataSets.remove( dataSet );
+        dataSet.getSources().remove( this );        
+    }
+    
+    public void updateDataSets( Set<DataSet> updates )
+    {
+        for ( DataSet dataSet : new HashSet<DataSet>( dataSets ) )
+        {
+            if ( !updates.contains( dataSet ) )
+            {
+                removeDataSet( dataSet );
+            }
+        }
+        
+        for ( DataSet dataSet : updates )
+        {
+            addDataSet( dataSet );
+        }
+    }
+    
     public List<OrganisationUnit> getSortedChildren()
     {
         List<OrganisationUnit> sortedChildren = new ArrayList<OrganisationUnit>( children );

=== modified file 'dhis-2/dhis-api/src/main/java/org/hisp/dhis/organisationunit/OrganisationUnitGroup.java'
--- dhis-2/dhis-api/src/main/java/org/hisp/dhis/organisationunit/OrganisationUnitGroup.java	2011-05-30 12:43:38 +0000
+++ dhis-2/dhis-api/src/main/java/org/hisp/dhis/organisationunit/OrganisationUnitGroup.java	2011-05-30 14:56:19 +0000
@@ -78,17 +78,17 @@
         unit.getGroups().remove( this );
     }
     
-    public void updateOrganisationUnits( Set<OrganisationUnit> units )
+    public void updateOrganisationUnits( Set<OrganisationUnit> updates )
     {
-        for ( OrganisationUnit unit : members )
+        for ( OrganisationUnit unit : new HashSet<OrganisationUnit>( members ) )
         {
-            if ( !units.contains( unit ) )
+            if ( !updates.contains( unit ) )
             {
                 removeOrganisationUnit( unit );
             }
         }
         
-        for ( OrganisationUnit unit : units )
+        for ( OrganisationUnit unit : updates )
         {
             addOrganisationUnit( unit );
         }

=== modified file 'dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-datadictionary/src/main/java/org/hisp/dhis/dd/action/dataelementgroup/AddDataElementGroupAction.java'
--- dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-datadictionary/src/main/java/org/hisp/dhis/dd/action/dataelementgroup/AddDataElementGroupAction.java	2011-05-05 21:15:45 +0000
+++ dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-datadictionary/src/main/java/org/hisp/dhis/dd/action/dataelementgroup/AddDataElementGroupAction.java	2011-05-30 14:56:19 +0000
@@ -30,7 +30,6 @@
 import java.util.HashSet;
 import java.util.Set;
 
-import org.hisp.dhis.dataelement.DataElement;
 import org.hisp.dhis.dataelement.DataElementGroup;
 import org.hisp.dhis.dataelement.DataElementService;
 
@@ -66,9 +65,9 @@
         this.name = name;
     }
 
-    private Set<Integer> groupMembers = new HashSet<Integer>();
+    private Set<String> groupMembers = new HashSet<String>();
 
-    public void setGroupMembers( Set<Integer> groupMembers )
+    public void setGroupMembers( Set<String> groupMembers )
     {
         this.groupMembers = groupMembers;
     }
@@ -92,18 +91,11 @@
     {
         dataElementGroup = new DataElementGroup( name );
 
-        if ( !groupMembers.isEmpty() )
+        for ( String id : groupMembers )
         {
-            Set<DataElement> members = new HashSet<DataElement>( groupMembers.size() );
-
-            for ( Integer id : groupMembers )
-            {
-                members.add( dataElementService.getDataElement( id ) );
-            }
-
-            dataElementGroup.setMembers( members );
+            dataElementGroup.addDataElement( dataElementService.getDataElement( Integer.parseInt( id ) ) );
         }
-
+        
         dataElementService.addDataElementGroup( dataElementGroup );
 
         return SUCCESS;

=== modified file 'dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-datadictionary/src/main/java/org/hisp/dhis/dd/action/dataelementgroup/UpdateDataElementGroupAction.java'
--- dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-datadictionary/src/main/java/org/hisp/dhis/dd/action/dataelementgroup/UpdateDataElementGroupAction.java	2011-05-05 21:15:45 +0000
+++ dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-datadictionary/src/main/java/org/hisp/dhis/dd/action/dataelementgroup/UpdateDataElementGroupAction.java	2011-05-30 14:56:19 +0000
@@ -73,9 +73,9 @@
         this.name = name;
     }
 
-    private Set<Integer> groupMembers = new HashSet<Integer>();
+    private Set<String> groupMembers = new HashSet<String>();
 
-    public void setGroupMembers( Set<Integer> groupMembers )
+    public void setGroupMembers( Set<String> groupMembers )
     {
         this.groupMembers = groupMembers;
     }
@@ -108,12 +108,12 @@
         {
             Set<DataElement> members = new HashSet<DataElement>();
 
-            for ( Integer id : groupMembers )
+            for ( String id : groupMembers )
             {
-                members.add( dataElementService.getDataElement( id  ) );
+                members.add( dataElementService.getDataElement( Integer.parseInt( id ) ) );
             }
 
-            dataElementGroup.setMembers( members );
+            dataElementGroup.updateDataElements( members );
         }
 
         dataElementService.updateDataElementGroup( dataElementGroup );

=== modified file 'dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-dataset/src/main/java/org/hisp/dhis/dataset/action/DefineDataSetAssociationsAction.java'
--- dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-dataset/src/main/java/org/hisp/dhis/dataset/action/DefineDataSetAssociationsAction.java	2011-04-22 17:45:58 +0000
+++ dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-dataset/src/main/java/org/hisp/dhis/dataset/action/DefineDataSetAssociationsAction.java	2011-05-30 14:56:19 +0000
@@ -27,8 +27,11 @@
  * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
+import java.util.HashSet;
+
 import org.hisp.dhis.dataset.DataSet;
 import org.hisp.dhis.dataset.DataSetService;
+import org.hisp.dhis.organisationunit.OrganisationUnit;
 import org.hisp.dhis.oust.manager.SelectionTreeManager;
 
 import com.opensymphony.xwork2.Action;
@@ -80,8 +83,7 @@
     {
         DataSet dataSet = dataSetService.getDataSet( dataSetId );
 
-        dataSet.getSources().clear();
-        dataSet.getSources().addAll( selectionTreeManager.getReloadedSelectedOrganisationUnits() );
+        dataSet.updateOrganisationUnits( new HashSet<OrganisationUnit>( selectionTreeManager.getReloadedSelectedOrganisationUnits() ) );
         
         dataSetService.updateDataSet( dataSet );
 

=== modified file 'dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-organisationunit/src/main/java/org/hisp/dhis/oum/action/organisationunit/AddOrganisationUnitAction.java'
--- dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-organisationunit/src/main/java/org/hisp/dhis/oum/action/organisationunit/AddOrganisationUnitAction.java	2011-05-30 12:43:38 +0000
+++ dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-organisationunit/src/main/java/org/hisp/dhis/oum/action/organisationunit/AddOrganisationUnitAction.java	2011-05-30 14:56:19 +0000
@@ -33,7 +33,6 @@
 import java.util.Date;
 import java.util.HashSet;
 
-import org.hisp.dhis.dataset.DataSet;
 import org.hisp.dhis.dataset.DataSetService;
 import org.hisp.dhis.i18n.I18nFormat;
 import org.hisp.dhis.organisationunit.OrganisationUnit;
@@ -257,20 +256,13 @@
             parent.getChildren().add( organisationUnit );
         }
 
+        for ( String id : dataSets )
+        {
+            organisationUnit.addDataSet( dataSetService.getDataSet( Integer.parseInt( id ) ) );
+        }
+        
         organisationUnitId = organisationUnitService.addOrganisationUnit( organisationUnit );
 
-        for ( String id : dataSets )
-        {
-            DataSet dataSet = dataSetService.getDataSet( Integer.parseInt( id ) );
-            
-            if ( dataSet != null )
-            {
-                dataSet.getSources().add( organisationUnit );
-                organisationUnit.getDataSets().add( dataSet );
-                dataSetService.updateDataSet( dataSet );
-            }
-        }
-
         for ( String id : selectedGroups )
         {
             OrganisationUnitGroup group = organisationUnitGroupService.getOrganisationUnitGroup( Integer.parseInt( id ) );

=== modified file 'dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-organisationunit/src/main/java/org/hisp/dhis/oum/action/organisationunit/UpdateOrganisationUnitAction.java'
--- dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-organisationunit/src/main/java/org/hisp/dhis/oum/action/organisationunit/UpdateOrganisationUnitAction.java	2011-05-30 12:43:38 +0000
+++ dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-organisationunit/src/main/java/org/hisp/dhis/oum/action/organisationunit/UpdateOrganisationUnitAction.java	2011-05-30 14:56:19 +0000
@@ -34,6 +34,7 @@
 import java.util.Date;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 import org.hisp.dhis.dataset.DataSet;
 import org.hisp.dhis.dataset.DataSetService;
@@ -267,25 +268,15 @@
         organisationUnit.setEmail( email );
         organisationUnit.setPhoneNumber( phoneNumber );
 
-        for ( DataSet dataSet : organisationUnit.getDataSets() ) // Remove current
-        {
-            dataSet.getSources().remove( organisationUnit );
-        }
-
-        organisationUnit.getDataSets().clear();
-
-        for ( String id : dataSets ) // Add selected
-        {
-            DataSet dataSet = dataSetService.getDataSet( Integer.parseInt( id ) );
-            
-            if ( dataSet != null )
-            {
-                dataSet.getSources().add( organisationUnit );
-                organisationUnit.getDataSets().add( dataSet );
-                dataSetService.updateDataSet( dataSet );
-            }
-        }
-
+        Set<DataSet> sets = new HashSet<DataSet>();
+        
+        for ( String id : dataSets )
+        {
+            sets.add( dataSetService.getDataSet( Integer.parseInt( id ) ) );
+        }
+
+        organisationUnit.updateDataSets( sets );
+        
         organisationUnitService.updateOrganisationUnit( organisationUnit );
 
         for ( int i = 0; i < orgUnitGroupSets.size(); i++ )