Fix dropping tables and handle the prefix automatically

Signed-off-by: Joas Schilling <coding@schilljs.com>

2 files changed, 143 insertions, 3 deletions

diff --git a/lib/private/DB/MigrationService.php b/lib/private/DB/MigrationService.php
index 912166988c0..e4a001c6d97 100644
--- a/lib/private/DB/MigrationService.php
+++ b/lib/private/DB/MigrationService.php
@@ -395,11 +395,12 @@ class MigrationService {
 		}, ['tablePrefix' => $this->connection->getPrefix()]);
 
 		$toSchema = $instance->changeSchema($this->output, function() {
-			return $this->connection->createSchema();
+			return new SchemaWrapper($this->connection);
 		}, ['tablePrefix' => $this->connection->getPrefix()]);
 
-		if ($toSchema instanceof Schema) {
-			$this->connection->migrateToSchema($toSchema);
+		if ($toSchema instanceof SchemaWrapper) {
+			$this->connection->migrateToSchema($toSchema->getWrappedSchema());
+			$toSchema->performDropTableCalls();
 		}
 
 		$instance->postSchemaChange($this->output, function() {
diff --git a/lib/private/DB/SchemaWrapper.php b/lib/private/DB/SchemaWrapper.php
new file mode 100644
index 00000000000..0be45d8d3f7
--- /dev/null
+++ b/lib/private/DB/SchemaWrapper.php
@@ -0,0 +1,139 @@
+<?php
+/**
+ * @copyright Copyright (c) 2017 Joas Schilling <coding@schilljs.com>
+ *
+ * @license GNU AGPL version 3 or any later version
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+namespace OC\DB;
+
+use Doctrine\DBAL\DBALException;
+use Doctrine\DBAL\Schema\Schema;
+use OCP\IDBConnection;
+
+class SchemaWrapper {
+
+	/** @var IDBConnection|Connection */
+	protected $connection;
+
+	/** @var Schema */
+	protected $schema;
+
+	/** @var array */
+	protected $tablesToDelete;
+
+	/**
+	 * @param IDBConnection $connection
+	 */
+	public function __construct(IDBConnection $connection) {
+		$this->connection = $connection;
+		$this->schema = $this->connection->createSchema();
+	}
+
+	public function getWrappedSchema() {
+		return $this->schema;
+	}
+
+	public function performDropTableCalls() {
+		foreach ($this->tablesToDelete as $tableName => $true) {
+			$this->connection->dropTable($tableName);
+			unset($this->tablesToDelete[$tableName]);
+		}
+	}
+
+	/**
+	 * Gets all table names
+	 *
+	 * @return array
+	 */
+	public function getTableNamesWithoutPrefix() {
+		$tableNames = $this->schema->getTableNames();
+		return array_map(function($tableName) {
+			if (strpos($tableName, $this->connection->getPrefix()) === 0) {
+				return substr($tableName, strlen($this->connection->getPrefix()));
+			}
+
+			return $tableName;
+		}, $tableNames);
+	}
+
+	// Overwritten methods
+
+	/**
+	 * @param string $tableName
+	 *
+	 * @return \Doctrine\DBAL\Schema\Table
+	 * @throws \Doctrine\DBAL\Schema\SchemaException
+	 */
+	public function getTable($tableName) {
+		return $this->schema->getTable($this->connection->getPrefix() . $tableName);
+	}
+
+	/**
+	 * Does this schema have a table with the given name?
+	 *
+	 * @param string $tableName
+	 *
+	 * @return boolean
+	 */
+	public function hasTable($tableName) {
+		return $this->schema->hasTable($this->connection->getPrefix() . $tableName);
+	}
+
+	/**
+	 * Creates a new table.
+	 *
+	 * @param string $tableName
+	 * @return \Doctrine\DBAL\Schema\Table
+	 */
+	public function createTable($tableName) {
+		return $this->schema->createTable($this->connection->getPrefix() . $tableName);
+	}
+
+	/**
+	 * Renames a table.
+	 *
+	 * @param string $oldTableName
+	 * @param string $newTableName
+	 *
+	 * @return \Doctrine\DBAL\Schema\Schema
+	 * @throws DBALException
+	 */
+	public function renameTable($oldTableName, $newTableName) {
+		throw new DBALException('Renaming tables is not supported. Please create and drop the tables manually.');
+	}
+
+	/**
+	 * Drops a table from the schema.
+	 *
+	 * @param string $tableName
+	 * @return \Doctrine\DBAL\Schema\Schema
+	 */
+	public function dropTable($tableName) {
+		$this->tablesToDelete[$tableName] = true;
+		return $this->schema->dropTable($this->connection->getPrefix() . $tableName);
+	}
+
+	/**
+	 * @param string $name
+	 * @param array $arguments
+	 * @return mixed
+	 */
+	public function __call($name, $arguments) {
+		return call_user_func_array([$this->schema, $name], $arguments);
+	}
+}