diff --git a/data_entity/data_entity.module b/data_entity/data_entity.module
index 0903ac0..d085cfd 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'),
           ),
         ),
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)) {