From 01ddeaa2be6880490c9a2a3113c12b3436e57cf9 Mon Sep 17 00:00:00 2001
From: Valerii Gorbachev <darkdef@gmail.com>
Date: Wed, 8 Mar 2023 21:36:53 +0300
Subject: [PATCH] Add contexts for profiler (#568)

* Add contexts for profiler

* styleci

* psalm fix

* Profiler tests

* fix

* add more test

* styleci
---
 src/Command/AbstractCommand.php            | 11 ++++--
 src/Driver/PDO/AbstractConnectionPDO.php   |  8 ++--
 src/Profiler/Context/AbstractContext.php   | 34 ++++++++++++++++
 src/Profiler/Context/CommandContext.php    | 35 ++++++++++++++++
 src/Profiler/Context/ConnectionContext.php | 13 ++++++
 src/Profiler/ContextInterface.php          | 15 +++++++
 src/Profiler/ProfilerInterface.php         |  8 ++--
 tests/AbstractCommandTest.php              | 46 +++++++++++++++++++++-
 tests/AbstractConnectionTest.php           | 29 ++++++++++++++
 tests/Db/Command/CommandTest.php           | 10 ++++-
 10 files changed, 195 insertions(+), 14 deletions(-)
 create mode 100644 src/Profiler/Context/AbstractContext.php
 create mode 100644 src/Profiler/Context/CommandContext.php
 create mode 100644 src/Profiler/Context/ConnectionContext.php
 create mode 100644 src/Profiler/ContextInterface.php

diff --git a/src/Command/AbstractCommand.php b/src/Command/AbstractCommand.php
index 59f264548..5fde84195 100644
--- a/src/Command/AbstractCommand.php
+++ b/src/Command/AbstractCommand.php
@@ -11,6 +11,7 @@
 use Yiisoft\Db\Exception\Exception;
 use Yiisoft\Db\Expression\Expression;
 use Yiisoft\Db\Profiler\ProfilerAwareTrait;
+use Yiisoft\Db\Profiler\Context\CommandContext;
 use Yiisoft\Db\Query\Data\DataReaderInterface;
 use Yiisoft\Db\Query\QueryInterface;
 
@@ -540,24 +541,26 @@ protected function queryInternal(int $queryMode): mixed
         $isReadMode = $this->isReadMode($queryMode);
 
         $logCategory = self::class . '::' . ($isReadMode ? 'query' : 'execute');
+        $queryContext = new CommandContext(__METHOD__, $logCategory, $this->getSql(), $this->getParams());
 
         $this->logQuery($rawSql, $logCategory);
         $this->prepare($isReadMode);
 
         try {
-            $this->profiler?->begin($rawSql, [$logCategory]);
+            $this->profiler?->begin($rawSql, $queryContext);
+
             $this->internalExecute($rawSql);
 
             /** @psalm-var mixed $result */
             $result = $this->internalGetQueryResult($queryMode);
-            $this->profiler?->end($rawSql, [$logCategory]);
+
+            $this->profiler?->end($rawSql, $queryContext);
 
             if (!$isReadMode) {
                 $this->refreshTableSchema();
             }
         } catch (Exception $e) {
-            $this->profiler?->end($rawSql, [$logCategory]);
-
+            $this->profiler?->end($rawSql, $queryContext->setException($e));
             throw $e;
         }
 
diff --git a/src/Driver/PDO/AbstractConnectionPDO.php b/src/Driver/PDO/AbstractConnectionPDO.php
index d9135f8a2..b7f52e866 100644
--- a/src/Driver/PDO/AbstractConnectionPDO.php
+++ b/src/Driver/PDO/AbstractConnectionPDO.php
@@ -12,6 +12,7 @@
 use Yiisoft\Db\Exception\Exception;
 use Yiisoft\Db\Exception\InvalidCallException;
 use Yiisoft\Db\Exception\InvalidConfigException;
+use Yiisoft\Db\Profiler\Context\ConnectionContext;
 use Yiisoft\Db\QueryBuilder\QueryBuilderInterface;
 use Yiisoft\Db\Schema\QuoterInterface;
 use Yiisoft\Db\Schema\SchemaInterface;
@@ -78,14 +79,15 @@ public function open(): void
         }
 
         $token = 'Opening DB connection: ' . $this->driver->getDsn();
+        $connectionContext = new ConnectionContext(__METHOD__);
 
         try {
             $this->logger?->log(LogLevel::INFO, $token);
-            $this->profiler?->begin($token, [__METHOD__]);
+            $this->profiler?->begin($token, $connectionContext);
             $this->initConnection();
-            $this->profiler?->end($token, [__METHOD__]);
+            $this->profiler?->end($token, $connectionContext);
         } catch (PDOException $e) {
-            $this->profiler?->end($token, [__METHOD__]);
+            $this->profiler?->end($token, $connectionContext->setException($e));
             $this->logger?->log(LogLevel::ERROR, $token);
 
             throw new Exception($e->getMessage(), (array) $e->errorInfo, $e);
diff --git a/src/Profiler/Context/AbstractContext.php b/src/Profiler/Context/AbstractContext.php
new file mode 100644
index 000000000..7534f55d3
--- /dev/null
+++ b/src/Profiler/Context/AbstractContext.php
@@ -0,0 +1,34 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Yiisoft\Db\Profiler\Context;
+
+use Throwable;
+use Yiisoft\Db\Profiler\ContextInterface;
+
+abstract class AbstractContext implements ContextInterface
+{
+    protected const METHOD = 'method';
+    protected const EXCEPTION = 'exception';
+
+    private Throwable|null $exception = null;
+
+    public function __construct(private string $method)
+    {
+    }
+
+    public function setException(Throwable $e): static
+    {
+        $this->exception = $e;
+        return $this;
+    }
+
+    public function asArray(): array
+    {
+        return [
+            self::METHOD => $this->method,
+            self::EXCEPTION => $this->exception,
+        ];
+    }
+}
diff --git a/src/Profiler/Context/CommandContext.php b/src/Profiler/Context/CommandContext.php
new file mode 100644
index 000000000..b5e4a9cb3
--- /dev/null
+++ b/src/Profiler/Context/CommandContext.php
@@ -0,0 +1,35 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Yiisoft\Db\Profiler\Context;
+
+final class CommandContext extends AbstractContext
+{
+    private const LOG_CONTEXT = 'logContext';
+    private const SQL = 'sql';
+    private const PARAMS = 'params';
+
+    public function __construct(
+        private string $method,
+        private string $logContext,
+        private string $sql,
+        private array $params,
+    ) {
+        parent::__construct($this->method);
+    }
+
+    public function getType(): string
+    {
+        return 'command';
+    }
+
+    public function asArray(): array
+    {
+        return parent::asArray() + [
+            self::LOG_CONTEXT => $this->logContext,
+            self::SQL => $this->sql,
+            self::PARAMS => $this->params,
+        ];
+    }
+}
diff --git a/src/Profiler/Context/ConnectionContext.php b/src/Profiler/Context/ConnectionContext.php
new file mode 100644
index 000000000..b7ee9e17f
--- /dev/null
+++ b/src/Profiler/Context/ConnectionContext.php
@@ -0,0 +1,13 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Yiisoft\Db\Profiler\Context;
+
+final class ConnectionContext extends AbstractContext
+{
+    public function getType(): string
+    {
+        return 'connection';
+    }
+}
diff --git a/src/Profiler/ContextInterface.php b/src/Profiler/ContextInterface.php
new file mode 100644
index 000000000..9383b7658
--- /dev/null
+++ b/src/Profiler/ContextInterface.php
@@ -0,0 +1,15 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Yiisoft\Db\Profiler;
+
+interface ContextInterface
+{
+    /**
+     * @return string Type of the context
+     */
+    public function getType(): string;
+
+    public function asArray(): array;
+}
diff --git a/src/Profiler/ProfilerInterface.php b/src/Profiler/ProfilerInterface.php
index a9c9556d4..24e52b98c 100644
--- a/src/Profiler/ProfilerInterface.php
+++ b/src/Profiler/ProfilerInterface.php
@@ -19,9 +19,9 @@ interface ProfilerInterface
      * The begin and end calls must also be nested.
      *
      * @param string $token Token for the code block.
-     * @param array $context The context data of this profile block.
+     * @param array|ContextInterface $context The context data of this profile block.
      */
-    public function begin(string $token, array $context = []): void;
+    public function begin(string $token, array|ContextInterface $context = []): void;
 
     /**
      * Marks the end of a code block for profiling.
@@ -29,9 +29,9 @@ public function begin(string $token, array $context = []): void;
      * This has to be matched with an earlier call to {@see begin()} with the same category name.
      *
      * @param string $token Token for the code block.
-     * @param array $context The context data of this profile block.
+     * @param array|ContextInterface $context The context data of this profile block.
      *
      * {@see begin()}
      */
-    public function end(string $token, array $context = []): void;
+    public function end(string $token, array|ContextInterface $context = []): void;
 }
diff --git a/tests/AbstractCommandTest.php b/tests/AbstractCommandTest.php
index fa7a59bad..d0607c2b1 100644
--- a/tests/AbstractCommandTest.php
+++ b/tests/AbstractCommandTest.php
@@ -11,6 +11,8 @@
 use Yiisoft\Db\Exception\Exception;
 use Yiisoft\Db\Exception\InvalidConfigException;
 use Yiisoft\Db\Exception\NotSupportedException;
+use Yiisoft\Db\Profiler\Context\CommandContext;
+use Yiisoft\Db\Profiler\ContextInterface;
 use Yiisoft\Db\Profiler\ProfilerInterface;
 use Yiisoft\Db\Tests\Support\DbHelper;
 use Yiisoft\Db\Tests\Support\TestTrait;
@@ -217,9 +219,9 @@ public function testSetSql(): void
      * @throws Throwable
      * @throws \PHPUnit\Framework\MockObject\Exception
      */
-    public function testProfiler(): void
+    public function testProfiler(string $sql = null): void
     {
-        $sql = 'SELECT 123';
+        $sql = $sql ??  'SELECT 123';
 
         $db = $this->getConnection();
         $db->open();
@@ -237,4 +239,44 @@ public function testProfiler(): void
 
         $db->createCommand($sql)->execute();
     }
+
+    /**
+     * @throws Exception
+     * @throws InvalidConfigException
+     * @throws Throwable
+     * @throws \PHPUnit\Framework\MockObject\Exception
+     */
+    public function testProfilerData(string $sql = null): void
+    {
+        $sql = $sql ?? 'SELECT 123';
+
+        $db = $this->getConnection();
+        $db->open();
+
+        $profiler = new class ($this, $sql) implements ProfilerInterface {
+            public function __construct(private TestCase $test, private string $sql)
+            {
+            }
+
+            public function begin(string $token, ContextInterface|array $context = []): void
+            {
+                $this->test->assertSame($this->sql, $token);
+                $this->test->assertInstanceOf(CommandContext::class, $context);
+                $this->test->assertSame('command', $context->getType());
+                $this->test->assertIsArray($context->asArray());
+            }
+
+            public function end(string $token, ContextInterface|array $context = []): void
+            {
+                $this->test->assertSame($this->sql, $token);
+                $this->test->assertInstanceOf(CommandContext::class, $context);
+                $this->test->assertSame('command', $context->getType());
+                $this->test->assertIsArray($context->asArray());
+            }
+        };
+
+        $db->setProfiler($profiler);
+
+        $db->createCommand($sql)->execute();
+    }
 }
diff --git a/tests/AbstractConnectionTest.php b/tests/AbstractConnectionTest.php
index 367862844..b603d2799 100644
--- a/tests/AbstractConnectionTest.php
+++ b/tests/AbstractConnectionTest.php
@@ -10,6 +10,8 @@
 use Yiisoft\Db\Driver\PDO\ConnectionPDOInterface;
 use Yiisoft\Db\Exception\InvalidConfigException;
 use Yiisoft\Db\Exception\NotSupportedException;
+use Yiisoft\Db\Profiler\Context\ConnectionContext;
+use Yiisoft\Db\Profiler\ContextInterface;
 use Yiisoft\Db\Query\BatchQueryResult;
 use Yiisoft\Db\Query\Query;
 use Yiisoft\Db\Tests\Support\Assert;
@@ -110,6 +112,33 @@ public function testNotProfiler(): void
         $this->assertNull(Assert::getInaccessibleProperty($db, 'profiler'));
     }
 
+    public function testProfiler(): void
+    {
+        $db = $this->getConnection();
+
+        $profiler = new class ($this) implements ProfilerInterface {
+            public function __construct(private TestCase $test)
+            {
+            }
+
+            public function begin(string $token, ContextInterface|array $context = []): void
+            {
+                $this->test->assertInstanceOf(ConnectionContext::class, $context);
+                $this->test->assertSame('connection', $context->getType());
+                $this->test->assertIsArray($context->asArray());
+            }
+
+            public function end(string $token, ContextInterface|array $context = []): void
+            {
+                $this->test->assertInstanceOf(ConnectionContext::class, $context);
+                $this->test->assertSame('connection', $context->getType());
+                $this->test->assertIsArray($context->asArray());
+            }
+        };
+        $db->setProfiler($profiler);
+        $db->open();
+    }
+
     public function testSetTablePrefix(): void
     {
         $db = $this->getConnection();
diff --git a/tests/Db/Command/CommandTest.php b/tests/Db/Command/CommandTest.php
index 46568bce9..f1b033e61 100644
--- a/tests/Db/Command/CommandTest.php
+++ b/tests/Db/Command/CommandTest.php
@@ -712,11 +712,19 @@ public function testUpsert(): void
         $command->upsert('{{table}}', []);
     }
 
-    public function testProfiler(): void
+    public function testProfiler(string $sql = null): void
     {
         $this->expectExceptionMessage(
             'Yiisoft\Db\Tests\Support\Stub\Command::internalExecute is not supported by this DBMS.'
         );
         parent::testProfiler();
     }
+
+    public function testProfilerData(string $sql = null): void
+    {
+        $this->expectExceptionMessage(
+            'Yiisoft\Db\Tests\Support\Stub\Command::internalExecute is not supported by this DBMS.'
+        );
+        parent::testProfilerData();
+    }
 }