ConnectionResolver.php 2.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. <?php
  2. declare(strict_types=1);
  3. /**
  4. * This file is part of Hyperf.
  5. *
  6. * @link https://www.hyperf.io
  7. * @document https://hyperf.wiki
  8. * @contact group@hyperf.io
  9. * @license https://github.com/hyperf/hyperf/blob/master/LICENSE
  10. */
  11. namespace Hyperf\DbConnection;
  12. use Hyperf\Context\Context;
  13. use Hyperf\Coroutine\Coroutine;
  14. use Hyperf\Database\ConnectionInterface;
  15. use Hyperf\Database\ConnectionResolverInterface;
  16. use Hyperf\DbConnection\Pool\PoolFactory;
  17. use Psr\Container\ContainerInterface;
  18. use function Hyperf\Coroutine\defer;
  19. class ConnectionResolver implements ConnectionResolverInterface
  20. {
  21. /**
  22. * The default connection name.
  23. */
  24. protected string $default = 'default';
  25. protected PoolFactory $factory;
  26. public function __construct(protected ContainerInterface $container)
  27. {
  28. $this->factory = $container->get(PoolFactory::class);
  29. }
  30. /**
  31. * Get a database connection instance.
  32. */
  33. public function connection(?string $name = null): ConnectionInterface
  34. {
  35. if (is_null($name)) {
  36. $name = $this->getDefaultConnection();
  37. }
  38. $connection = null;
  39. $id = $this->getContextKey($name);
  40. if (Context::has($id)) {
  41. $connection = Context::get($id);
  42. }
  43. if (! $connection instanceof ConnectionInterface) {
  44. $pool = $this->factory->getPool($name);
  45. $connection = $pool->get();
  46. try {
  47. // PDO is initialized as an anonymous function, so there is no IO exception,
  48. // but if other exceptions are thrown, the connection will not return to the connection pool properly.
  49. $connection = $connection->getConnection();
  50. Context::set($id, $connection);
  51. } finally {
  52. if (Coroutine::inCoroutine()) {
  53. defer(function () use ($connection, $id) {
  54. Context::set($id, null);
  55. $connection->release();
  56. });
  57. }
  58. }
  59. }
  60. return $connection;
  61. }
  62. /**
  63. * Get the default connection name.
  64. */
  65. public function getDefaultConnection(): string
  66. {
  67. return $this->default;
  68. }
  69. /**
  70. * Set the default connection name.
  71. */
  72. public function setDefaultConnection(string $name): void
  73. {
  74. $this->default = $name;
  75. }
  76. /**
  77. * The key to identify the connection object in coroutine context.
  78. * @param mixed $name
  79. */
  80. private function getContextKey($name): string
  81. {
  82. return sprintf('database.connection.%s', $name);
  83. }
  84. }