diff --git a/src/Helpers/QueryBuilderHelper.php b/src/Helpers/QueryBuilderHelper.php
index 9b14e0e..78e4b3c 100644
--- a/src/Helpers/QueryBuilderHelper.php
+++ b/src/Helpers/QueryBuilderHelper.php
@@ -116,7 +116,6 @@ public function delete($key = null)
 
         return $this->getClient()->deleteBatch($keys);
     }
-
     /**
      * @inheritdoc
      */
@@ -149,11 +148,21 @@ public function insert(array $values, $key = '', $options = [])
 		return null;
     }
 
+    /**
+     * @inheritdoc
+     */
+    public function update(array $values, $key = '', $options = [])
+    {
+        $this->applyBeforeQueryCallbacks();
+        return $this->upsert($values, $key, $options);
+    }
+
     /**
      * @inheritdoc
      */
     public function upsert(array $values, $key = '', $options = [])
     {
+        
         if (empty($this->from)) {
             throw new \LogicException('No kind/table specified');
         }
@@ -161,10 +170,23 @@ public function upsert(array $values, $key = '', $options = [])
         if (empty($values)) {
             return true;
         }
-
-        $key = $key ? $this->getClient()->key($this->from, $key) : $this->getClient()->key($this->from);
-
-        $entity = $this->getClient()->entity($key, $values, $options);
+        foreach($this->wheres as $where):
+            if($where['value'] instanceof \Google\Cloud\Datastore\Key):
+                $key = $where['value'];
+                break;
+            endif;
+        endforeach;
+
+        if($key instanceof \Google\Cloud\Datastore\Key):
+            $entity = $this->getClient()->lookup($key) ?? $this->getClient()->entity($key, $values, $options);
+
+            foreach($values as $key=>$value):
+                $entity->$key = $value;
+            endforeach;
+        else:
+            $key = $key ? $this->getClient()->key($this->from, $key) : $this->getClient()->key($this->from);
+            $entity = $this->getClient()->entity($key, $values, $options);
+        endif;
 
         return $this->getClient()->upsert($entity);
     }