From 2c5b1dc20a6bd0ba593237dc430585e9a5827b8a Mon Sep 17 00:00:00 2001
From: epriestley <git@epriestley.com>
Date: Wed, 17 Aug 2016 07:59:59 -0700
Subject: [PATCH] Provide "--output" flags for "bin/storage renamespace"

Summary:
Ref T6996. Depends on D16407. This does the same stuff as D16407, but for `bin/storage renamespace`. In particular:

  - Support writing directly to a file (so we can get good errors on failure).
  - Support in-process compression.

Also add support for reading out of a `storage dump` subprocess, so we don't have to do a dump-to-disk + renamespace + compress dance and can just stream out of MySQL directly to a compressed file on disk.

This is used in the second stage of instance exports (see T7148).

It would be nice to share more code with `bin/storage dump`, and possibly to just make this a flag for it, although we still do need to do the file-based version when importing (vs exporting). I figured that was better left for another time.

Test Plan:
Ran `bin/storage renamespace --live --output x --from A --to B --compress --overwrite` and similar commands.

Verified that a compressed, renamespaced dump came out of the other end.

Reviewers: chad

Reviewed By: chad

Maniphest Tasks: T6996

Differential Revision: https://secure.phabricator.com/D16410
---
 ...orStorageManagementRenamespaceWorkflow.php | 158 +++++++++++++++---
 1 file changed, 137 insertions(+), 21 deletions(-)

diff --git a/src/infrastructure/storage/management/workflow/PhabricatorStorageManagementRenamespaceWorkflow.php b/src/infrastructure/storage/management/workflow/PhabricatorStorageManagementRenamespaceWorkflow.php
index e64242344b..ad5afac7fb 100644
--- a/src/infrastructure/storage/management/workflow/PhabricatorStorageManagementRenamespaceWorkflow.php
+++ b/src/infrastructure/storage/management/workflow/PhabricatorStorageManagementRenamespaceWorkflow.php
@@ -8,15 +8,20 @@ final class PhabricatorStorageManagementRenamespaceWorkflow
       ->setName('renamespace')
       ->setExamples(
         '**renamespace** [__options__] '.
-        '--in __dump.sql__ --from __old__ --to __new__ > __out.sql__')
+        '--input __dump.sql__ --from __old__ --to __new__ > __out.sql__')
       ->setSynopsis(pht('Change the database namespace of a .sql dump file.'))
       ->setArguments(
         array(
           array(
-            'name' => 'in',
+            'name' => 'input',
             'param' => 'file',
             'help' => pht('SQL dumpfile to process.'),
           ),
+          array(
+            'name' => 'live',
+            'help' => pht(
+              'Generate a live dump instead of processing a file on disk.'),
+          ),
           array(
             'name' => 'from',
             'param' => 'namespace',
@@ -27,6 +32,21 @@ final class PhabricatorStorageManagementRenamespaceWorkflow
             'param' => 'namespace',
             'help' => pht('Desired database namespace for output.'),
           ),
+          array(
+            'name' => 'output',
+            'param' => 'file',
+            'help' => pht('Write output directly to a file on disk.'),
+          ),
+          array(
+            'name' => 'compress',
+            'help' => pht('Emit gzipped output instead of plain text.'),
+          ),
+          array(
+            'name' => 'overwrite',
+            'help' => pht(
+              'With __--output__, write to disk even if the file already '.
+              'exists.'),
+          ),
         ));
   }
 
@@ -37,12 +57,13 @@ final class PhabricatorStorageManagementRenamespaceWorkflow
   public function didExecute(PhutilArgumentParser $args) {
     $console = PhutilConsole::getConsole();
 
-    $in = $args->getArg('in');
-    if (!strlen($in)) {
+    $input = $args->getArg('input');
+    $is_live = $args->getArg('live');
+    if (!strlen($input) && !$is_live) {
       throw new PhutilArgumentUsageException(
         pht(
-          'Specify the dumpfile to read with %s.',
-          '--in'));
+          'Specify the dumpfile to read with "--in", or use "--live" to '.
+          'generate one automatically.'));
     }
 
     $from = $args->getArg('from');
@@ -61,6 +82,62 @@ final class PhabricatorStorageManagementRenamespaceWorkflow
           '--to'));
     }
 
+
+    $output_file = $args->getArg('output');
+    $is_overwrite = $args->getArg('overwrite');
+    $is_compress = $args->getArg('compress');
+
+    if ($is_overwrite) {
+      if ($output_file === null) {
+        throw new PhutilArgumentUsageException(
+          pht(
+            'The "--overwrite" flag can only be used alongside "--output".'));
+      }
+    }
+
+    if ($output_file !== null) {
+      if (Filesystem::pathExists($output_file)) {
+        if (!$is_overwrite) {
+          throw new PhutilArgumentUsageException(
+            pht(
+              'Output file "%s" already exists. Use "--overwrite" '.
+              'to overwrite.',
+              $output_file));
+        }
+      }
+    }
+
+    if ($is_live) {
+      $root = dirname(phutil_get_library_root('phabricator'));
+
+      $future = new ExecFuture(
+        '%R dump',
+        $root.'/bin/storage');
+
+      $lines = new LinesOfALargeExecFuture($future);
+    } else {
+      $lines = new LinesOfALargeFile($input);
+    }
+
+    if ($output_file === null) {
+      $file = fopen('php://stdout', 'wb');
+      $output_name = pht('stdout');
+    } else {
+      if ($is_compress) {
+        $file = gzopen($output_file, 'wb');
+      } else {
+        $file = fopen($output_file, 'wb');
+      }
+      $output_name = $output_file;
+    }
+
+    if (!$file) {
+      throw new Exception(
+        pht(
+          'Failed to open output file "%s" for writing.',
+          $output_name));
+    }
+
     $patterns = array(
       'use' => '@^(USE `)([^_]+)(_.*)$@',
       'create' => '@^(CREATE DATABASE /\*.*?\*/ `)([^_]+)(_.*)$@',
@@ -68,27 +145,66 @@ final class PhabricatorStorageManagementRenamespaceWorkflow
 
     $found = array_fill_keys(array_keys($patterns), 0);
 
-    $matches = null;
-    foreach (new LinesOfALargeFile($in) as $line) {
+    try {
+      $matches = null;
+      foreach ($lines as $line) {
 
-      foreach ($patterns as $key => $pattern) {
-        if (preg_match($pattern, $line, $matches)) {
-          $namespace = $matches[2];
-          if ($namespace != $from) {
-            throw new Exception(
-              pht(
-                'Expected namespace "%s", found "%s": %s.',
-                $from,
-                $namespace,
-                $line));
+        foreach ($patterns as $key => $pattern) {
+          if (preg_match($pattern, $line, $matches)) {
+            $namespace = $matches[2];
+            if ($namespace != $from) {
+              throw new Exception(
+                pht(
+                  'Expected namespace "%s", found "%s": %s.',
+                  $from,
+                  $namespace,
+                  $line));
+            }
+
+            $line = $matches[1].$to.$matches[3];
+            $found[$key]++;
           }
+        }
 
-          $line = $matches[1].$to.$matches[3];
-          $found[$key]++;
+        $data = $line."\n";
+
+        if ($is_compress) {
+          $bytes = gzwrite($file, $data);
+        } else {
+          $bytes = fwrite($file, $data);
+        }
+
+        if ($bytes !== strlen($data)) {
+          throw new Exception(
+            pht(
+              'Failed to write %d byte(s) to "%s".',
+              new PhutilNumber(strlen($data)),
+              $output_name));
         }
       }
 
-      echo $line."\n";
+      if ($is_compress) {
+        $ok = gzclose($file);
+      } else {
+        $ok = fclose($file);
+      }
+
+      if ($ok !== true) {
+        throw new Exception(
+          pht(
+            'Failed to close file "%s".',
+            $output_file));
+      }
+    } catch (Exception $ex) {
+      try {
+        if ($output_file !== null) {
+          Filesystem::remove($output_file);
+        }
+      } catch (Exception $ex) {
+        // Ignore any exception.
+      }
+
+      throw $ex;
     }
 
     // Give the user a chance to catch things if the results are crazy.