diff --git a/view_unpublished.install b/view_unpublished.install
index 90331a0..7324da0 100644
--- a/view_unpublished.install
+++ b/view_unpublished.install
@@ -16,4 +16,12 @@ function view_unpublished_enable() {
  */
 function view_unpublished_disable() {
   node_access_needs_rebuild(TRUE);
-}
\ No newline at end of file
+}
+
+/**
+ * Rebuild the node access table and flush caches.
+ */
+function view_unpublished_update_7001() {
+  node_access_needs_rebuild(TRUE);
+  drupal_flush_all_caches();
+}
diff --git a/view_unpublished.module b/view_unpublished.module
index fafc85e..f8d9078 100644
--- a/view_unpublished.module
+++ b/view_unpublished.module
@@ -6,9 +6,9 @@
  */
 
 /**
- * Implements hook_permission().
+ * Helper function to retrieve list of permissions.
  */
-function view_unpublished_permission() {
+function view_unpublished_get_permissions_list() {
   $perms = array(
     'view any unpublished content' => array(
       'title' => t('View any unpublished content'),
@@ -26,51 +26,67 @@ function view_unpublished_permission() {
 }
 
 /**
- * Implements hook_node_access_records().
+ * Implements hook_permission().
  */
-function view_unpublished_node_access_records($node) {
-  // We only care about the node if is unpublished. If not, it is
-  // treated just like any other node and we completely ignore it.
-  if ($node->status == 0) {
-    $grants = array();
-    // Unpublished nodes should be viewable to all editors.
-    $grants[] = array(
-      'realm' => 'view_unpublished_content',
-      'gid' => 1,
-      'grant_view' => 1,
-      'grant_update' => 0,
-      'grant_delete' => 0,
-      'priority' => 0,
-    );
-    $grants[] = array(
-      'realm' => 'view_unpublished_' . $node->type . '_content',
-      'gid' => 1,
-      'grant_view' => 1,
-      'grant_update' => 0,
-      'grant_delete' => 0,
-      'priority' => 0,
-    );
-    return $grants;
+function view_unpublished_permission() {
+  return view_unpublished_get_permissions_list();
+}
+
+/**
+ * Implements hook_og_permission().
+ */
+function view_unpublished_og_permission() {
+  if (module_exists('og_context')) {
+    return view_unpublished_get_permissions_list();
   }
 }
 
 /**
- * Implements hook_node_grants().
+ * Implements hook_node_access().
  */
-function view_unpublished_node_grants($account, $op) {
-  $grants = array();
+function view_unpublished_node_access($node, $op, $account) {
   if ($op == 'view') {
+    // Set a good default response
+    $access = NODE_ACCESS_IGNORE;
+
+    // Check to see if we're going to check for OG permissions as well.
+    $check_og = FALSE;
+    if (module_exists('og_context')) {
+      $og = og_context();
+      if (!empty($og)) {
+        $check_og = TRUE;
+        $gid = $og->gid;
+      }
+    }
+
+    // If a user has the appropriate permission, then allow. If the user doesn't
+    // have the appropriate permission, but we're checking OG permissions and
+    // we're inside an organic group, use the group permission instead.
     if (user_access('view any unpublished content', $account)) {
-      $grants['view_unpublished_content'] = array(1);
-      return $grants;
+      $access = NODE_ACCESS_ALLOW;
     }
-    foreach(node_permissions_get_configured_types() as $type) {
-      if (user_access("view any unpublished $type content", $account)) {
-        $grants['view_unpublished_' . $type . '_content'] = array(1);
-      }
+    if ($check_og && isset($gid) && og_user_access($gid, 'view any unpublished content', $account)) {
+      $access = NODE_ACCESS_ALLOW;
+    }
+
+    // If it's allowed at this point, then we can bail out - we know that the user
+    // will be allowed.
+    if ($access == NODE_ACCESS_ALLOW) {
+      return $access;
+    }
+
+    // If a user has the appropriate permission, then allow. If the user doesn't
+    // have the appropriate permission, but we're checking OG permissions and
+    // we're inside an organic group, use the group permission instead.
+    if (user_access('view any unpublished ' . $node->type . ' content', $account)) {
+      $access = NODE_ACCESS_ALLOW;
     }
+    if ($check_og && isset($gid) && og_user_access($gid, 'view any unpublished ' . $node->type . ' content', $account)) {
+      $access = NODE_ACCESS_ALLOW;
+    }
+
+    return $access;
   }
-  return $grants;
 }
 
 /**