From eaff89f8bd0d11d7b21905ec26de7c84f53d91b0 Mon Sep 17 00:00:00 2001
From: ozh <ozh@ozh.org>
Date: Fri, 7 Feb 2014 11:16:44 +0100
Subject: [PATCH] Check if DB server is dead

---
 includes/class-mysql.php | 54 ++++++++++++++++++++++++++++++++++++----
 includes/functions.php   | 12 ++++++---
 2 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/includes/class-mysql.php b/includes/class-mysql.php
index 072a36f..03364ff 100644
--- a/includes/class-mysql.php
+++ b/includes/class-mysql.php
@@ -45,6 +45,7 @@ function yourls_set_DB_driver( ) {
 	yourls_do_action( 'set_DB_driver', $driver );
 		
 	$ydb = new $class( YOURLS_DB_USER, YOURLS_DB_PASS, YOURLS_DB_NAME, YOURLS_DB_HOST );
+    $ydb->DB_driver = $driver;
 
 	yourls_debug_log( "DB driver: $driver" );
 }
@@ -71,12 +72,55 @@ function yourls_db_connect() {
 		yourls_set_DB_driver();
 	}
 	
-	// Check if connection attempt raised an error. It seems that only PDO does, though.
-	if ( $ydb->last_error )
-		yourls_die( $ydb->last_error, yourls__( 'Fatal error' ), 503 );
-
-	
 	return $ydb;
 }
 
+/**
+ * Return true if DB server is responding
+ *
+ * This function is supposed to be called right after yourls_get_all_options() has fired. It is not designed (yet) to
+ * check for a responding server after several successful operation to check if the server has gone MIA
+ *
+ * @since 1.7.1
+ */
+function yourls_is_db_alive() {
+    global $ydb;
+    
+    $alive = false;
+    switch( $ydb->DB_driver ) {
+        case 'pdo' :
+            $alive = isset( $ydb->dbh );
+            break;
+    
+        case 'mysql' :
+            $alive = ( isset( $ydb->dbh ) && false !== $ydb->dbh );
+            break;
+    
+        case 'mysqli' :
+            $alive = ( null == mysqli_connect_error() );
+            break;
+        
+        // Custom DB driver & class : delegate check
+        default:
+            $alive = yourls_apply_filter( 'is_db_alive_custom', false );
+    }
+    
+    return $alive;
+}
+
+/**
+ * Die with a DB error message
+ *
+ * @TODO in version 1.8 : use a new localized string, specific to the problem (ie: "DB is dead")
+ *
+ * @since 1.7.1
+ */
+function yourls_db_dead() {
+    // Use any /user/db_error.php file
+    if( file_exists( YOURLS_USERDIR . '/db_error.php' ) ) {
+        include_once( YOURLS_USERDIR . '/db_error.php' );
+        die();
+    }
 
+    yourls_die( yourls__( 'Incorrect DB config, or could not connect to DB' ), yourls__( 'Fatal error' ), 503 );
+}
\ No newline at end of file
diff --git a/includes/functions.php b/includes/functions.php
index ac55ebe..e0b4d78 100644
--- a/includes/functions.php
+++ b/includes/functions.php
@@ -986,9 +986,10 @@ function yourls_get_option( $option_name, $default = false ) {
 /**
  * Read all options from DB at once
  *
- * The goal is to read all option at once and then populate array $ydb->option, to prevent further
+ * The goal is to read all options at once and then populate array $ydb->option, to prevent further
  * SQL queries if we need to read an option value later.
- * It's also a simple check whether YOURLS is installed or not (no option = assuming not installed)
+ * It's also a simple check whether YOURLS is installed or not (no option = assuming not installed) after
+ * a check for DB server reachability has been performed
  *
  * @since 1.4
  */
@@ -1012,8 +1013,11 @@ function yourls_get_all_options() {
 		$ydb->option = yourls_apply_filter( 'get_all_options', $ydb->option );
 		$ydb->installed = true;
 	} else {
-		// Zero option found: assume YOURLS is not installed
-		$ydb->installed = false;
+		// Zero option found: either YOURLS is not installed or DB server is dead
+        if( !yourls_is_db_alive() ) {
+            yourls_db_dead(); // YOURLS will die here
+        }
+        $ydb->installed = false;
 	}
 }
 
-- 
2.45.0