lib/distribute_reads.rb in distribute_reads-0.2.3 vs lib/distribute_reads.rb in distribute_reads-0.2.4

@@ -30,11 +30,11 @@
 
     connection ||= ActiveRecord::Base.connection
 
     replica_pool = connection.instance_variable_get(:@slave_pool)
     if replica_pool && replica_pool.connections.size > 1
-      warn "[distribute_reads] Multiple replicas available, lag only reported for one"
+      log "Multiple replicas available, lag only reported for one"
     end
 
     if %w(PostgreSQL PostGIS).include?(connection.adapter_name)
       # cache the version number
       @server_version_num ||= {}
@@ -57,17 +57,34 @@
     elsif %w(MySQL Mysql2 Mysql2Spatial Mysql2Rgeo).include?(connection.adapter_name)
       replica_value = Thread.current[:distribute_reads][:replica]
       begin
         # makara doesn't send SHOW queries to replica, so we must force it
         Thread.current[:distribute_reads][:replica] = true
-        status = connection.exec_query("SHOW SLAVE STATUS").to_hash.first
-        status ? status["Seconds_Behind_Master"].to_f : 0.0
+
+        @aurora_mysql ||= {}
+        cache_key = connection.pool.object_id
+
+        unless @aurora_mysql.key?(cache_key)
+          @aurora_mysql[cache_key] = connection.exec_query("SHOW VARIABLES LIKE 'aurora_version'").to_hash.any?
+        end
+
+        if @aurora_mysql[cache_key]
+          status = connection.exec_query("SELECT Replica_lag_in_msec FROM mysql.ro_replica_status WHERE Server_id = @@aurora_server_id").to_hash.first
+          status ? status["Replica_lag_in_msec"].to_f / 1000.0 : 0.0
+        else
+          status = connection.exec_query("SHOW SLAVE STATUS").to_hash.first
+          status ? status["Seconds_Behind_Master"].to_f : 0.0
+        end
       ensure
         Thread.current[:distribute_reads][:replica] = replica_value
       end
     else
       raise DistributeReads::Error, "Option not supported with this adapter"
     end
+  end
+
+  def self.log(message)
+    warn "[distribute_reads] #{message}"
   end
 
   # private
   def self.makara3?
     unless defined?(@makara3)