diff --git a/data_entity/data_entity.module b/data_entity/data_entity.module
index 0903ac0..704e9bd 100644
--- a/data_entity/data_entity.module
+++ b/data_entity/data_entity.module
@@ -63,17 +63,18 @@ function data_entity_entity_info() {
       'label' => $table->title,
       'controller class' => 'DataEntityController',
       'base table'  => $table_name,
+      'bundle keys' => array('bundle' => 'type'),
       'fieldable' => TRUE,
       'entity keys' => array(
         'id' => $id_field,
       ),
       'bundles' => array(
         $entity_type => array(
+          'type'  => $entity_type,
           'label' => $table->title,
           'admin' => array(
             'path'      => 'admin/structure/data/edit/%data_ui_table',
             'real path' => 'admin/structure/data/edit/' . $table_name,
-            'bundle argument' => 4,
             'access arguments' => array('administer data tables'),
           ),
         ),
@@ -123,19 +124,49 @@ function data_entity_menu() {
 function data_entity_menu_alter(&$items) {
   // Attempt at fixing tab parentage for Field UI items.
   // @see data_entity_entity_info() for more details.
-  $tables = data_entity_get_entity_tables();
-  $info = array();
 
-  /*
-  // This breaks access on these menu items!
-  foreach ($tables as $table_name => $table) {
-    $field_ui_base_path = "admin/structure/data/edit/$table_name";
-    foreach (array('fields', 'fields/%', 'display') as $path_suffix) {
-      $items[$field_ui_base_path . '/' . $path_suffix]['tab_parent']  = 'admin/structure/data/edit/%';
-      $items[$field_ui_base_path . '/' . $path_suffix]['tab_root']    = 'admin/structure/data/edit/%';
+  // Fix Field UI menu items
+  if (module_exists('field_ui')) {
+    $prefix = 'admin/structure/data/edit/%data_ui_table/fields';
+    $prefix_new = 'admin/structure/data/edit/%data_entity_string/fields';
+
+    $argument_to_replace = '__nothing__';
+    $entity_tables = data_entity_get_entity_tables();
+    if (!$entity_tables) {
+      return;
+    }
+
+    // We don't want data_ui_table_load() to run when editing fields,
+    // Entity UI forms accept string names as their arguments.
+    $item = $items[$prefix];
+    $static_name = $item['page arguments'][1];
+    foreach ($item['page arguments'] as $arg_key => $value) {
+      if ($value == $static_name) {
+        $item['page arguments'][$arg_key] = 4;
+      }
+    }
+    $items[$prefix_new] = $item;
+    unset($items[$prefix]);
+
+    foreach (array('/%field_ui_menu', '/%field_ui_menu/edit', '/%field_ui_menu/field-settings',
+      '/%field_ui_menu/widget-type', '/%field_ui_menu/delete') as $suffix) {
+      $item = $items[$prefix . $suffix];
+
+      // Field UI incorrectly replaces dynamic % argument with the last value loaded with %data_ui_table
+      // Fix this.
+      if (!empty($item['load arguments'])) {
+        foreach ($item['load arguments'] as $arg_key => $value) {
+          if ($value == $static_name) {
+            $item['load arguments'][$arg_key] = 4;
+          }
+        }
+      }
+      $new_key = $prefix_new . $suffix;
+      $new_key = str_replace('%field_ui_menu', '%data_entity_field_ui_menu', $new_key);
+      $items[$new_key] = $item;
+      unset($items[$prefix . $suffix]);
     }
   }
-  */
 }
 
 /**
@@ -157,6 +188,27 @@ function data_entity_item_load($deid, $table_name) {
 }
 
 /**
+ * Menu loader callback.
+ * @see data_entity_menu_alter()
+ *
+ */
+function data_entity_string_load($table_name) {
+  return 'data_' . $table_name;
+}
+
+/**
+ * Helper menu loader function to pass correct values to field_ui_menu_load()
+ * @see field_ui_menu_load()
+ * @see data_entity_menu_alter()
+ */
+function data_entity_field_ui_menu_load() {
+  $args = func_get_args();
+  $args[1] = 'data_' . $args[1];
+  $args[2] = 'data_' . $args[2];
+  return call_user_func_array('field_ui_menu_load', $args);
+}
+
+/**
  * Implements hook_permission().
  */
 function data_entity_permission() {
diff --git a/data_entity/data_entity.pages.inc b/data_entity/data_entity.pages.inc
index 32a3e19..e72f09d 100644
--- a/data_entity/data_entity.pages.inc
+++ b/data_entity/data_entity.pages.inc
@@ -93,24 +93,28 @@ function data_entity_entity_edit_form_validate($form, &$form_state) {
  * Form submit handler for saving a data entity.
  */
 function data_entity_entity_edit_form_submit($form, &$form_state) {
-  //dsm($form_state, 'fs');
-
   $data_entity = $form['#entity'];
   $entity_type = $data_entity->entity_type;
 
+  entity_form_submit_build_entity($entity_type, $data_entity, $form, $form_state);
+
+  $tables = data_entity_get_entity_tables();
+  // Entity type is data_$table_name
+  $table_name = substr($entity_type, 5);
+  $key_field = data_entity_get_id_field($tables[$table_name]);
+
+  // Let modules modify entity before it is saved to the database.
+  field_attach_presave($entity_type, $data_entity);
+  module_invoke_all('entity_presave', $data_entity, $entity_type);
+
   $table = $form_state['values']['table'];
 
   $record = $form_state['values']['data'];
   drupal_write_record($table->name, $record, $table->table_schema['primary key']);
 
-  // Build a pseudo entity for FieldAPI field attach.
-  $pseudo_entity = $form_state['values'];
-  unset($pseudo_entity['data'], $pseudo_entity['table']);
-  $pseudo_entity += $form_state['values']['data'];
-  $pseudo_entity = (object) $pseudo_entity;
-
-  field_attach_submit($entity_type, $pseudo_entity, $form, $form_state);
+  $op = empty($data_entity->{$key_field}) ? 'insert' : 'update';
+  $function = "field_attach_$op";
+  $function($entity_type, $data_entity);
 
-  // Save fields.
-  field_attach_update($entity_type, $pseudo_entity);
+  module_invoke_all('entity_' . $op, $data_entity, $entity_type);
 }
diff --git a/includes/DataTable.inc b/includes/DataTable.inc
index 30a85cb..8604b88 100644
--- a/includes/DataTable.inc
+++ b/includes/DataTable.inc
@@ -58,6 +58,7 @@ class DataTable {
   protected function __construct($name) {
     // Set our name after sanitizing it.
     $this->name = db_escape_table($name);
+    $this->type = 'data_' . $this->name;
 
     // Try to load table information.
     if ($table = _data_load_table($name)) {