From a9d59cd6fe413b59a19378d0cd971468fa9d2a35 Mon Sep 17 00:00:00 2001
From: Rix Beck <rix@boltcms.hu>
Date: Wed, 1 Apr 2020 19:58:46 +0200
Subject: [PATCH] EM based table prefixing

---
 src/Doctrine/TablePrefix.php                  | 33 ++++++++------
 src/Doctrine/TablePrefixTrait.php             | 45 +++++++++++++++++++
 .../Subscriber/TimedPublishSubscriber.php     | 43 +++++++++++-------
 3 files changed, 91 insertions(+), 30 deletions(-)
 create mode 100644 src/Doctrine/TablePrefixTrait.php

diff --git a/src/Doctrine/TablePrefix.php b/src/Doctrine/TablePrefix.php
index 5b2ccd00..66143f66 100644
--- a/src/Doctrine/TablePrefix.php
+++ b/src/Doctrine/TablePrefix.php
@@ -4,33 +4,38 @@ declare(strict_types=1);
 
 namespace Bolt\Doctrine;
 
-use Bolt\Common\Str;
 use Doctrine\ORM\Event\LoadClassMetadataEventArgs;
 use Doctrine\ORM\Mapping\ClassMetadataInfo;
+use Symfony\Bridge\Doctrine\ManagerRegistry;
 
 class TablePrefix
 {
-    private $tablePrefix;
+    use TablePrefixTrait;
 
-    public function __construct(string $tablePrefix)
+    public function __construct($tablePrefixes, ManagerRegistry $managerRegistry)
     {
-        $this->tablePrefix = Str::ensureEndsWith($tablePrefix, '_');
+        $this->setTablePrefixes($tablePrefixes, $managerRegistry);
     }
 
     public function loadClassMetadata(LoadClassMetadataEventArgs $eventArgs): void
     {
-        $classMetadata = $eventArgs->getClassMetadata();
+        if ($tablePrefix = $this->getTablePrefix($eventArgs->getEntityManager())) {
+            $classMetadata = $eventArgs->getClassMetadata();
 
-        if (! $classMetadata->isInheritanceTypeSingleTable() || $classMetadata->getName() === $classMetadata->rootEntityName) {
-            $classMetadata->setPrimaryTable([
-                'name' => $this->tablePrefix . $classMetadata->getTableName(),
-            ]);
-        }
+            if (!$classMetadata->isInheritanceTypeSingleTable()
+                || $classMetadata->getName() === $classMetadata->rootEntityName) {
+                $classMetadata->setPrimaryTable(
+                    [
+                        'name' => $tablePrefix.$classMetadata->getTableName(),
+                    ]
+                );
+            }
 
-        foreach ($classMetadata->getAssociationMappings() as $fieldName => $mapping) {
-            if ($mapping['type'] === ClassMetadataInfo::MANY_TO_MANY && $mapping['isOwningSide']) {
-                $mappedTableName = $mapping['joinTable']['name'];
-                $classMetadata->associationMappings[$fieldName]['joinTable']['name'] = $this->tablePrefix . $mappedTableName;
+            foreach ($classMetadata->getAssociationMappings() as $fieldName => $mapping) {
+                if ($mapping['type'] === ClassMetadataInfo::MANY_TO_MANY && $mapping['isOwningSide']) {
+                    $mappedTableName = $mapping['joinTable']['name'];
+                    $classMetadata->associationMappings[$fieldName]['joinTable']['name'] = $tablePrefix.$mappedTableName;
+                }
             }
         }
     }
diff --git a/src/Doctrine/TablePrefixTrait.php b/src/Doctrine/TablePrefixTrait.php
new file mode 100644
index 00000000..37129a73
--- /dev/null
+++ b/src/Doctrine/TablePrefixTrait.php
@@ -0,0 +1,45 @@
+<?php
+/**
+ * @author Rix Beck <rix@neologik.hu>
+ */
+
+namespace Bolt\Doctrine;
+
+use Bolt\Common\Str;
+use Doctrine\Persistence\ObjectManager;
+use Symfony\Bridge\Doctrine\ManagerRegistry;
+
+trait TablePrefixTrait
+{
+    private $tablePrefixes = [];
+
+    protected function setTablePrefix(ObjectManager $manager, string $prefix)
+    {
+        $key = spl_object_hash($manager);
+        $this->tablePrefixes[$key] = Str::ensureEndsWith($prefix, '_');
+
+        return $this;
+    }
+
+    protected function setTablePrefixes($tablePrefixes, ManagerRegistry $managerRegistry)
+    {
+        $prefixes = (array)$tablePrefixes;
+        $this->registry = $managerRegistry;
+
+        foreach ($prefixes as $em => $tablePrefix) {
+            $manager = $managerRegistry->getManager(is_int($em) ? 'default' : $em);
+            if ($manager) {
+                $this->setTablePrefix($manager, $tablePrefix);
+            }
+        }
+
+        return $this;
+    }
+
+    protected function getTablePrefix(ObjectManager $manager)
+    {
+        $key = spl_object_hash($manager);
+
+        return $this->tablePrefixes[$key] ?? '';
+    }
+}
diff --git a/src/Event/Subscriber/TimedPublishSubscriber.php b/src/Event/Subscriber/TimedPublishSubscriber.php
index 6be94605..455740df 100644
--- a/src/Event/Subscriber/TimedPublishSubscriber.php
+++ b/src/Event/Subscriber/TimedPublishSubscriber.php
@@ -4,14 +4,17 @@ declare(strict_types=1);
 
 namespace Bolt\Event\Subscriber;
 
-use Bolt\Common\Str;
+use Bolt\Doctrine\TablePrefixTrait;
 use Carbon\Carbon;
 use Doctrine\ORM\EntityManagerInterface;
+use Symfony\Bridge\Doctrine\ManagerRegistry;
 use Symfony\Component\EventDispatcher\EventSubscriberInterface;
 use Symfony\Component\HttpKernel\KernelEvents;
 
 class TimedPublishSubscriber implements EventSubscriberInterface
 {
+    use TablePrefixTrait;
+
     public const PRIORITY = 30;
 
     /** @var EntityManagerInterface */
@@ -20,10 +23,22 @@ class TimedPublishSubscriber implements EventSubscriberInterface
     /** @var string */
     private $tablePrefix;
 
-    public function __construct(string $tablePrefix, EntityManagerInterface $entityManager)
+    public function __construct($tablePrefixes, ManagerRegistry $managerRegistry)
     {
-        $this->entityManager = $entityManager;
-        $this->tablePrefix = Str::ensureEndsWith($tablePrefix, '_');
+        $this->entityManager = $managerRegistry->getManager('default');
+        $this->tablePrefix = $this
+            ->setTablePrefixes($tablePrefixes, $managerRegistry)
+            ->getTablePrefix($this->entityManager);
+    }
+
+    /**
+     * Return the events to subscribe to.
+     */
+    public static function getSubscribedEvents(): array
+    {
+        return [
+            KernelEvents::REQUEST => [['onKernelRequest', self::PRIORITY]], // Right after route is matched
+        ];
     }
 
     /**
@@ -36,8 +51,14 @@ class TimedPublishSubscriber implements EventSubscriberInterface
 
         // Publish timed Content records when 'publish_at' has passed and Depublish published Content
         // records when 'depublish_at' has passed. Note: Placeholders in DBAL don't work for tablenames.
-        $queryPublish = sprintf('update %scontent SET status = "published", published_at = :now  WHERE status = "timed" AND published_at < :now', $this->tablePrefix);
-        $queryDepublish = sprintf('update %scontent SET status = "held", depublished_at = :now WHERE status = "published" AND depublished_at < :now', $this->tablePrefix);
+        $queryPublish = sprintf(
+            'update %scontent SET status = "published", published_at = :now  WHERE status = "timed" AND published_at < :now',
+            $this->tablePrefix
+        );
+        $queryDepublish = sprintf(
+            'update %scontent SET status = "held", depublished_at = :now WHERE status = "published" AND depublished_at < :now',
+            $this->tablePrefix
+        );
 
         try {
             $conn->executeUpdate($queryPublish, [':now' => $now]);
@@ -46,14 +67,4 @@ class TimedPublishSubscriber implements EventSubscriberInterface
             // Fail silently, output user-friendly exception elsewhere.
         }
     }
-
-    /**
-     * Return the events to subscribe to.
-     */
-    public static function getSubscribedEvents(): array
-    {
-        return [
-            KernelEvents::REQUEST => [['onKernelRequest', self::PRIORITY]], // Right after route is matched
-        ];
-    }
 }