Add a "--database <name> ..." flag to "bin/storage dump"

Summary: Ref T13671. Allow "bin/storage dump" to dump a subset of databases, primarily to support merging previously-partitioned databases.

Test Plan: Ran `bin/storage dump` with and without `--database ...` flags. Ran `--database invalid`, `--database a --database a` to hit error cases.

Maniphest Tasks: T13671

Differential Revision: https://secure.phabricator.com/D21745

Author: epriestley

src/infrastructure/storage/management/PhabricatorStorageManagementAPI.php

@@ -89,6 +89,17 @@ public function getDatabaseName($fragment) {
     return $this->namespace.'_'.$fragment;
   }
 
+  public function getInternalDatabaseName($name) {
+    $namespace = $this->getNamespace();
+
+    $prefix = $namespace.'_';
+    if (strncmp($name, $prefix, strlen($prefix))) {
+      return null;
+    }
+
+    return substr($name, strlen($prefix));
+  }
+
   public function getDisplayName() {
     return $this->getRef()->getDisplayName();
   }

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementDumpWorkflow.php

@@ -44,6 +44,16 @@ protected function didConstruct() {
               'With __--output__, overwrite the output file if it already '.
               'exists.'),
           ),
+          array(
+            'name' => 'database',
+            'param' => 'database-name',
+            'help' => pht(
+              'Dump only tables in the named database (or databases, if '.
+              'the flag is repeated). Specify database names without the '.
+              'namespace prefix (that is: use "differential", not '.
+              '"phabricator_differential").'),
+            'repeat' => true,
+          ),
         ));
   }
 
@@ -58,6 +68,8 @@ public function didExecute(PhutilArgumentParser $args) {
     $is_noindex = $args->getArg('no-indexes');
     $is_replica = $args->getArg('for-replica');
 
+    $database_filter = $args->getArg('database');
+
     if ($is_compress) {
       if ($output_file === null) {
         throw new PhutilArgumentUsageException(
@@ -128,11 +140,60 @@ public function didExecute(PhutilArgumentParser $args) {
     $schemata = $actual_map[$ref_key];
     $expect = $expect_map[$ref_key];
 
+    if ($database_filter) {
+      $internal_names = array();
+
+      $expect_databases = $expect->getDatabases();
+      foreach ($expect_databases as $expect_database) {
+        $database_name = $expect_database->getName();
+
+        $internal_name = $api->getInternalDatabaseName($database_name);
+        if ($internal_name !== null) {
+          $internal_names[$internal_name] = $database_name;
+        }
+      }
+
+      ksort($internal_names);
+
+      $seen = array();
+      foreach ($database_filter as $filter) {
+        if (!isset($internal_names[$filter])) {
+          throw new PhutilArgumentUsageException(
+            pht(
+              'Database "%s" is unknown. This script can only dump '.
+              'databases known to the current version of Phabricator. '.
+              'Valid databases are: %s.',
+              $filter,
+              implode(', ', array_keys($internal_names))));
+        }
+
+        if (isset($seen[$filter])) {
+          throw new PhutilArgumentUsageException(
+            pht(
+              'Database "%s" is specified more than once. Specify each '.
+              'database at most once.',
+              $filter));
+        }
+
+        $seen[$filter] = true;
+      }
+
+      $dump_databases = array_select_keys($internal_names, $database_filter);
+      $dump_databases = array_fuse($dump_databases);
+    } else {
+      $dump_databases = array_keys($schemata->getDatabases());
+      $dump_databases = array_fuse($dump_databases);
+    }
+
     $with_caches = $is_replica;
     $with_indexes = !$is_noindex;
 
     $targets = array();
     foreach ($schemata->getDatabases() as $database_name => $database) {
+      if (!isset($dump_databases[$database_name])) {
+        continue;
+      }
+
       $expect_database = $expect->getDatabase($database_name);
       foreach ($database->getTables() as $table_name => $table) {