diff --git a/ui/src/app/components/datasource-entity.directive.js b/ui/src/app/components/datasource-entity.directive.js
index f9074c7066..3b1dd19844 100644
--- a/ui/src/app/components/datasource-entity.directive.js
+++ b/ui/src/app/components/datasource-entity.directive.js
@@ -192,9 +192,9 @@ function DatasourceEntity($compile, $templateCache, $q, $mdDialog, $window, $doc
                 dataKey.color = color;
                 ngModelCtrl.$setDirty();
             });
-        }
+        };
 
-        scope.editDataKey = function (event, dataKey, index) {
+        scope.editDataKey = function (event, dataKey) {
 
             $mdDialog.show({
                 controller: 'DatakeyConfigDialogController',
@@ -214,11 +214,12 @@ function DatasourceEntity($compile, $templateCache, $q, $mdDialog, $window, $doc
                     var w = angular.element($window);
                     w.triggerHandler('resize');
                 }
-            }).then(function (dataKey) {
-                if ((dataKey.type === types.dataKeyType.timeseries) || (dataKey.type === types.dataKeyType.attribute)) {
-                    scope.dataKeys[index] = dataKey;
-                } else if (dataKey.type === types.dataKeyType.alarm) {
-                    scope.alarmDataKeys[index] = dataKey;
+            }).then(function (newDataKey) {
+                let index = scope.dataKeys.indexOf(dataKey);
+                if ((newDataKey.type === types.dataKeyType.timeseries) || (newDataKey.type === types.dataKeyType.attribute)) {
+                    scope.dataKeys[index] = newDataKey;
+                } else if (newDataKey.type === types.dataKeyType.alarm) {
+                    scope.alarmDataKeys[index] = newDataKey;
                 }
                 ngModelCtrl.$setDirty();
             }, function () {
diff --git a/ui/src/app/components/datasource-entity.scss b/ui/src/app/components/datasource-entity.scss
index 75a2eed632..086f6f691f 100644
--- a/ui/src/app/components/datasource-entity.scss
+++ b/ui/src/app/components/datasource-entity.scss
@@ -24,6 +24,7 @@
     display: block;
     height: 48px;
     line-height: 1.5;
+    white-space: normal !important;
 
     .tb-no-entries {
       line-height: 48px;
@@ -32,7 +33,6 @@
 
   li {
     height: auto !important;
-    white-space: normal !important;
   }
 }
 
diff --git a/ui/src/app/components/datasource-entity.tpl.html b/ui/src/app/components/datasource-entity.tpl.html
index b0461a9edf..9a58ecdb66 100644
--- a/ui/src/app/components/datasource-entity.tpl.html
+++ b/ui/src/app/components/datasource-entity.tpl.html
@@ -76,7 +76,7 @@
 							  <div class="tb-color-preview" ng-click="showColorPicker($event, $chip, $index)" style="margin-right: 5px;">
 								  <div class="tb-color-result" ng-style="{background: $chip.color}"></div>
 							  </div>
-							  <div layout="row">
+							  <div layout="row" flex>
 								  <div class="tb-chip-label">
                                       <span ng-show="$chip.type==types.dataKeyType.attribute">
                                           <md-tooltip>{{'datakey.attributes' | translate }}</md-tooltip>
@@ -94,7 +94,7 @@
 									  <strong ng-if="$chip.postFuncBody">f({{$chip.name}})</strong>
 								  </div>
 							  </div>
-							  <md-button ng-click="editDataKey($event, $chip, $index)" class="md-icon-button tb-md-32">
+							  <md-button ng-click="editDataKey($event, $chip)" class="md-icon-button tb-md-32">
 								  <md-icon aria-label="edit" class="material-icons tb-md-20">edit</md-icon>
 							  </md-button>
 						  </div>
@@ -139,7 +139,7 @@
 						   <div class="tb-color-preview" ng-click="showColorPicker($event, $chip, $index)" style="margin-right: 5px;">
 							   <div class="tb-color-result" ng-style="{background: $chip.color}"></div>
 						   </div>
-						   <div layout="row">
+						   <div layout="row" flex>
 							   <div class="tb-chip-label">
 								   {{$chip.label}}
 							   </div>
@@ -149,7 +149,7 @@
 								   <strong ng-if="$chip.postFuncBody">f({{$chip.name}})</strong>
 							   </div>
 						   </div>
-						   <md-button ng-click="editDataKey($event, $chip, $index)" class="md-icon-button tb-md-32">
+						   <md-button ng-click="editDataKey($event, $chip)" class="md-icon-button tb-md-32">
 							   <md-icon aria-label="edit" class="material-icons tb-md-20">edit</md-icon>
 						   </md-button>
 					   </div>
diff --git a/ui/src/app/components/datasource-func.directive.js b/ui/src/app/components/datasource-func.directive.js
index c42cb9344e..b305eeedcc 100644
--- a/ui/src/app/components/datasource-func.directive.js
+++ b/ui/src/app/components/datasource-func.directive.js
@@ -162,7 +162,7 @@ function DatasourceFunc($compile, $templateCache, $mdDialog, $window, $document,
             });
         }
 
-        scope.editDataKey = function (event, dataKey, index) {
+        scope.editDataKey = function (event, dataKey) {
 
             $mdDialog.show({
                 controller: 'DatakeyConfigDialogController',
@@ -182,11 +182,12 @@ function DatasourceFunc($compile, $templateCache, $mdDialog, $window, $document,
                     var w = angular.element($window);
                     w.triggerHandler('resize');
                 }
-            }).then(function (dataKey) {
-                if (dataKey.type === types.dataKeyType.function) {
-                    scope.funcDataKeys[index] = dataKey;
-                } else if (dataKey.type === types.dataKeyType.alarm) {
-                    scope.alarmDataKeys[index] = dataKey;
+            }).then(function (newDataKey) {
+                let index = scope.funcDataKeys.indexOf(dataKey);
+                if (newDataKey.type === types.dataKeyType.function) {
+                    scope.funcDataKeys[index] = newDataKey;
+                } else if (newDataKey.type === types.dataKeyType.alarm) {
+                    scope.alarmDataKeys[index] = newDataKey;
                 }
                 ngModelCtrl.$setDirty();
             }, function () {
diff --git a/ui/src/app/components/datasource-func.tpl.html b/ui/src/app/components/datasource-func.tpl.html
index a16d978e70..43cb963767 100644
--- a/ui/src/app/components/datasource-func.tpl.html
+++ b/ui/src/app/components/datasource-func.tpl.html
@@ -64,7 +64,7 @@
                           <div class="tb-color-preview" ng-click="showColorPicker($event, $chip, $index)" style="margin-right: 5px;">
 							  <div class="tb-color-result" ng-style="{background: $chip.color}"></div>
 						  </div>
-						  <div layout="row">
+						  <div layout="row" flex>
 							  <div class="tb-chip-label">
 								  {{$chip.label}}
 							  </div>
@@ -73,7 +73,7 @@
 								  <strong>{{$chip.name}}</strong>
 							  </div>
 						  </div>
-						  <md-button ng-click="editDataKey($event, $chip, $index)" class="md-icon-button tb-md-32">
+						  <md-button ng-click="editDataKey($event, $chip)" class="md-icon-button tb-md-32">
 							  <md-icon aria-label="edit" class="material-icons tb-md-20">edit</md-icon>
 						  </md-button>
 					  </div>
@@ -118,7 +118,7 @@
                     <div class="tb-color-preview" ng-click="showColorPicker($event, $chip, $index)" style="margin-right: 5px;">
 						<div class="tb-color-result" ng-style="{background: $chip.color}"></div>
 					</div>
-					<div layout="row">
+					<div layout="row" flex>
 						<div class="tb-chip-label">
 							{{$chip.label}}
 						</div>
@@ -128,7 +128,7 @@
 							<strong ng-if="$chip.postFuncBody">f({{$chip.name}})</strong>
 						</div>
 					</div>
-					<md-button ng-click="editDataKey($event, $chip, $index)" class="md-icon-button tb-md-32">
+					<md-button ng-click="editDataKey($event, $chip)" class="md-icon-button tb-md-32">
 						<md-icon aria-label="edit" class="material-icons tb-md-20">edit</md-icon>
 					</md-button>
 				</div>