[9.x] Allow handling cumulative query duration limit per DB connection

src/Illuminate/Database/Connection.php

@@ -2,6 +2,7 @@
 
 namespace Illuminate\Database;
 
+use Carbon\CarbonInterval;
 use Closure;
 use DateTimeInterface;
 use Doctrine\DBAL\Connection as DoctrineConnection;
@@ -19,6 +20,7 @@
 use Illuminate\Database\Query\Processors\Processor;
 use Illuminate\Database\Schema\Builder as SchemaBuilder;
 use Illuminate\Support\Arr;
+use Illuminate\Support\InteractsWithTime;
 use Illuminate\Support\Traits\Macroable;
 use PDO;
 use PDOStatement;
@@ -29,6 +31,7 @@ class Connection implements ConnectionInterface
     use DetectsConcurrencyErrors,
         DetectsLostConnections,
         Concerns\ManagesTransactions,
+        InteractsWithTime,
         Macroable;
 
     /**
@@ -157,6 +160,20 @@ class Connection implements ConnectionInterface
      */
     protected $loggingQueries = false;
 
+    /**
+     * The duration of all executed queries in milliseconds.
+     *
+     * @var float
+     */
+    protected $totalQueryDuration = 0.0;
+
+    /**
+     * All of the registered query duration handlers.
+     *
+     * @var array
+     */
+    protected $queryDurationHandlers = [];
+
     /**
      * Indicates if the connection is in a "dry run".
      *
@@ -755,6 +772,8 @@ protected function runQueryCallback($query, $bindings, Closure $callback)
      */
     public function logQuery($query, $bindings, $time = null)
     {
+        $this->totalQueryDuration += $time ?? 0.0;
+
         $this->event(new QueryExecuted($query, $bindings, $time, $this));
 
         if ($this->loggingQueries) {
@@ -773,6 +792,71 @@ protected function getElapsedTime($start)
         return round((microtime(true) - $start) * 1000, 2);
     }
 
+    /**
+     * Register a callback to be invoked when the connection queries for longer than a given amount of time.
+     *
+     * @param  \DateTimeInterface|\Carbon\CarbonInterval|float|int  $threshold
+     * @param  callable  $handler
+     * @return void
+     */
+    public function whenQueryingForLongerThan($threshold, $handler)
+    {
+        $threshold = $threshold instanceof DateTimeInterface
+            ? $this->secondsUntil($threshold) * 1000
+            : $threshold;
+
+        $threshold = $threshold instanceof CarbonInterval
+            ? $threshold->totalMilliseconds
+            : $threshold;
+
+        $this->queryDurationHandlers[] = [
+            'has_run' => false,
+            'handler' => $handler,
+        ];
+
+        $key = count($this->queryDurationHandlers) - 1;
+
+        $this->listen(function ($event) use ($threshold, $handler, $key) {
+            if (! $this->queryDurationHandlers[$key]['has_run'] && $this->totalQueryDuration() > $threshold) {
+                $handler($this, $event);
+
+                $this->queryDurationHandlers[$key]['has_run'] = true;
+            }
+        });
+    }
+
+    /**
+     * Allow all the query duration handlers to run again, even if they have already run.
+     *
+     * @return void
+     */
+    public function allowQueryDurationHandlersToRunAgain()
+    {
+        foreach ($this->queryDurationHandlers as $key => $queryDurationHandler) {
+            $this->queryDurationHandlers[$key]['has_run'] = false;
+        }
+    }
+
+    /**
+     * Get the duration of all run queries in milliseconds.
+     *
+     * @return float
+     */
+    public function totalQueryDuration()
+    {
+        return $this->totalQueryDuration;
+    }
+
+    /**
+     * Reset the duration of all run queries.
+     *
+     * @return void
+     */
+    public function resetTotalQueryDuration()
+    {
+        $this->totalQueryDuration = 0.0;
+    }
+
     /**
      * Handle a query exception.
      *

src/Illuminate/Queue/QueueServiceProvider.php

@@ -200,6 +200,13 @@ protected function registerWorker()
                     $app['log']->withoutContext();
                 }
 
+                if (method_exists($app['db'], 'getConnections')) {
+                    foreach ($app['db']->getConnections() as $connection) {
+                        $connection->resetTotalQueryDuration();
+                        $connection->allowQueryDurationHandlersToRunAgain();
+                    }
+                }
+
                 return $app->forgetScopedInstances();
             };
 

tests/Database/QueryDurationThresholdTest.php

@@ -0,0 +1,218 @@
+<?php
+
+namespace Illuminate\Tests\Database;
+
+use Carbon\CarbonInterval;
+use Illuminate\Database\Connection;
+use Illuminate\Events\Dispatcher;
+use Illuminate\Support\Arr;
+use PDO;
+use PHPUnit\Framework\TestCase;
+
+class QueryDurationThresholdTest extends TestCase
+{
+    public function testItCanHandleReachingADurationThresholdInTheDb()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+        $called = 0;
+        $connection->whenQueryingForLongerThan(CarbonInterval::milliseconds(1.1), function () use (&$called) {
+            $called++;
+        });
+
+        $connection->logQuery('xxxx', [], 1.0);
+        $connection->logQuery('xxxx', [], 0.1);
+        $this->assertSame(0, $called);
+
+        $connection->logQuery('xxxx', [], 0.1);
+        $this->assertSame(1, $called);
+    }
+
+    public function testItIsOnlyCalledOnce()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+        $called = 0;
+        $connection->whenQueryingForLongerThan(CarbonInterval::milliseconds(1), function () use (&$called) {
+            $called++;
+        });
+
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+
+        $this->assertSame(1, $called);
+    }
+
+    public function testItIsOnlyCalledOnceWhenGivenDateTime()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+        $called = 0;
+        $connection->whenQueryingForLongerThan(now()->addMilliseconds(1), function () use (&$called) {
+            $called++;
+        });
+
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+
+        $this->assertSame(1, $called);
+    }
+
+    public function testItCanSpecifyMultipleHandlersWithTheSameIntervals()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+        $called = [];
+        $connection->whenQueryingForLongerThan(CarbonInterval::milliseconds(1), function () use (&$called) {
+            $called['a'] = true;
+        });
+        $connection->whenQueryingForLongerThan(CarbonInterval::milliseconds(1), function () use (&$called) {
+            $called['b'] = true;
+        });
+
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+
+        $this->assertSame([
+            'a' => true,
+            'b' => true,
+        ], $called);
+    }
+
+    public function testItCanSpecifyMultipleHandlersWithDifferentIntervals()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+        $called = [];
+        $connection->whenQueryingForLongerThan(CarbonInterval::milliseconds(1), function () use (&$called) {
+            $called['a'] = true;
+        });
+        $connection->whenQueryingForLongerThan(CarbonInterval::milliseconds(2), function () use (&$called) {
+            $called['b'] = true;
+        });
+
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+        $this->assertSame([
+            'a' => true,
+        ], $called);
+
+        $connection->logQuery('xxxx', [], 1);
+        $this->assertSame([
+            'a' => true,
+            'b' => true,
+        ], $called);
+    }
+
+    public function testItHasAccessToConnectionInHandler()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'), '', '', ['name' => 'expected-name']);
+        $connection->setEventDispatcher(new Dispatcher());
+        $name = null;
+        $connection->whenQueryingForLongerThan(CarbonInterval::milliseconds(1), function ($connection) use (&$name) {
+            $name = $connection->getName();
+        });
+
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+
+        $this->assertSame('expected-name', $name);
+    }
+
+    public function testItHasSpecifyThresholdWithFloat()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+        $called = false;
+        $connection->whenQueryingForLongerThan(1.1, function () use (&$called) {
+            $called = true;
+        });
+
+        $connection->logQuery('xxxx', [], 1.1);
+        $this->assertFalse($called);
+
+        $connection->logQuery('xxxx', [], 0.1);
+        $this->assertTrue($called);
+    }
+
+    public function testItHasSpecifyThresholdWithInt()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+        $called = false;
+        $connection->whenQueryingForLongerThan(2, function () use (&$called) {
+            $called = true;
+        });
+
+        $connection->logQuery('xxxx', [], 1.1);
+        $this->assertFalse($called);
+
+        $connection->logQuery('xxxx', [], 1.0);
+        $this->assertTrue($called);
+    }
+
+    public function testItCanResetTotalQueryDuration()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+
+        $connection->logQuery('xxxx', [], 1.1);
+        $this->assertSame(1.1, $connection->totalQueryDuration());
+        $connection->logQuery('xxxx', [], 1.1);
+        $this->assertSame(2.2, $connection->totalQueryDuration());
+
+        $connection->resetTotalQueryDuration();
+        $this->assertSame(0.0, $connection->totalQueryDuration());
+    }
+
+    public function testItCanRestoreAlreadyRunHandlers()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+        $called = 0;
+        $connection->whenQueryingForLongerThan(CarbonInterval::milliseconds(1), function () use (&$called) {
+            $called++;
+        });
+
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+        $this->assertSame(1, $called);
+
+        $connection->allowQueryDurationHandlersToRunAgain();
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+        $this->assertSame(2, $called);
+
+        $connection->allowQueryDurationHandlersToRunAgain();
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+        $connection->logQuery('xxxx', [], 1);
+        $this->assertSame(3, $called);
+    }
+
+    public function testItCanAccessAllQueriesWhenQueryLoggingIsActive()
+    {
+        $connection = new Connection(new PDO('sqlite::memory:'));
+        $connection->setEventDispatcher(new Dispatcher());
+        $connection->enableQueryLog();
+        $queries = [];
+        $connection->whenQueryingForLongerThan(CarbonInterval::milliseconds(2), function ($connection, $event) use (&$queries) {
+            $queries = Arr::pluck($connection->getQueryLog(), 'query');
+            $queries[] = $event->sql;
+        });
+
+        $connection->logQuery('foo', [], 1);
+        $connection->logQuery('bar', [], 1);
+        $connection->logQuery('baz', [], 1);
+
+        $this->assertSame([
+            'foo',
+            'bar',
+            'baz',
+        ], $queries);
+    }
+}