mediagoblin/gmg_commands/import_export.py

   1 # GNU MediaGoblin -- federated, autonomous media hosting
   2 # Copyright (C) 2011 Free Software Foundation, Inc
   3 #
   4 # This program is free software: you can redistribute it and/or modify
   5 # it under the terms of the GNU Affero General Public License as published by
   6 # the Free Software Foundation, either version 3 of the License, or
   7 # (at your option) any later version.
   8 #
   9 # This program is distributed in the hope that it will be useful,
  10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 # GNU Affero General Public License for more details.
  13 #
  14 # You should have received a copy of the GNU Affero General Public License
  15 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16
  17 from mediagoblin.gmg_commands import util as commands_util
  18 from mediagoblin import mg_globals
  19 from mediagoblin.db import util as db_util
  20 from mediagoblin.db.open import setup_connection_and_db_from_config
  21 from mediagoblin.init.config import read_mediagoblin_config
  22 from mediagoblin import util as mg_util
  23 from mediagoblin.storage import BasicFileStorage
  24 from mediagoblin.init import setup_storage, setup_global_and_app_config
  25
  26 import shlex
  27 import shutil
  28 import tarfile
  29 import subprocess
  30 import os.path
  31 import os
  32 import re
  33 import sys
  34
  35
  36 def import_export_parse_setup(subparser):
  37     # TODO: Add default
  38     subparser.add_argument(
  39         'tar_file')
  40     subparser.add_argument(
  41         '-cf', '--conf_file', default='mediagoblin.ini',
  42         help='Config file used to set up environment')
  43     subparser.add_argument(
  44         '--mongodump_path', default='mongodump',
  45         help='mongodump binary')
  46     subparser.add_argument(
  47         '--mongorestore_path', default='mongorestore',
  48         help='mongorestore binary')
  49     subparser.add_argument(
  50         '--cache_path', default='/tmp/mediagoblin/',
  51         help='')
  52
  53
  54 def _export_database(db, args):
  55     print "\n== Exporting database ==\n"
  56
  57     command = '{mongodump_path} -d {database} -o {mongodump_cache}'.format(
  58         mongodump_path=args.mongodump_path,
  59         database=db.name,
  60         mongodump_cache=args._cache_path['database'])
  61
  62     p = subprocess.Popen(
  63         shlex.split(command))
  64
  65     p.wait()
  66
  67     print "\n== Database exported ==\n"
  68
  69
  70 def _export_media(db, args):
  71     print "\n== Exporting media ==\n"
  72
  73     media_cache = BasicFileStorage(
  74         args._cache_path['media'])
  75
  76     # TODO: Add export of queue files
  77     queue_cache = BasicFileStorage(
  78         args._cache_path['queue'])
  79
  80     for entry in db.media_entries.find():
  81         for name, path in entry['media_files'].items():
  82             mc_file = media_cache.get_file(path, mode='wb')
  83             mc_file.write(
  84                 mg_globals.public_store.get_file(path, mode='rb').read())
  85
  86             print(mc_file)
  87         print(entry)
  88
  89     print "\n== Media exported ==\n"
  90
  91
  92 def _import_media(db, args):
  93     """
  94     Import media files
  95
  96     Must be called after _import_database()
  97     """
  98     print "\n== Importing media ==\n"
  99
 100     media_cache = BasicFileStorage(
 101         args._cache_path['media'])
 102
 103     # TODO: Add import of queue files
 104     queue_cache = BasicFileStorage(
 105         args._cache_path['queue'])
 106
 107     for entry in db.media_entries.find():
 108         for name, path in entry['media_files'].items():
 109             media_file = mg_globals.public_store.get_file(path, mode='wb')
 110             media_file.write(
 111                 media_cache.get_file(path, mode='rb').read())
 112
 113             print(media_file)
 114         print(entry)
 115
 116     print "\n== Media imported ==\n"
 117
 118
 119 def _import_database(db, args):
 120     print "\n== Importing database ==\n"
 121     command = '{mongorestore_path} -d {database}'
 122     '{backup_dir}/{database}'.format(
 123         mongorestore_path=args.mongorestore_path,
 124         database=db.name,
 125         backup_dir=args._cache_path['database'])
 126
 127     print command
 128
 129     p = subprocess.Popen(
 130         shlex.split(command))
 131
 132     p.wait()
 133
 134
 135 def env_import(args):
 136     args.cache_path += 'mediagoblin-data'
 137     args = _setup_paths(args)
 138
 139     setup_global_and_app_config(args.conf_file)
 140     setup_storage()
 141
 142     config, validation_result = read_mediagoblin_config(args.conf_file)
 143     connection, db = setup_connection_and_db_from_config(
 144         config['mediagoblin'], use_pymongo=True)
 145
 146     tf = tarfile.open(
 147         args.tar_file,
 148         mode='r|gz')
 149
 150     tf.extractall(args.cache_path)
 151
 152     # Import database from extracted data
 153     _import_database(db, args)
 154
 155     _import_media(db, args)
 156
 157
 158 def _setup_paths(args):
 159     args._cache_path = dict()
 160     PATH_MAP = {
 161         'media': 'media',
 162         'queue': 'queue',
 163         'database': 'database'}
 164
 165     for key, val in PATH_MAP.items():
 166         args._cache_path[key] = os.path.join(args.cache_path, val)
 167
 168     return args
 169
 170
 171 def _create_archive(args):
 172     print "\n== Compressing to archive ==\n"
 173
 174     tf = tarfile.open(
 175         args.tar_file,
 176         mode='w|gz')
 177
 178     with tf:
 179         tf.add(args.cache_path, 'mediagoblin-data/')
 180
 181     print "\n== Archiving done ==\n"
 182
 183
 184 def _clean(args):
 185     shutil.rmtree(args.cache_path)
 186
 187
 188 def _check(args):
 189     if os.path.exists(args.tar_file):
 190         overwrite = raw_input(
 191             'The output file already exists. '
 192             'Are you **SURE** you want to overwrite it? '
 193             '(yes/no)> ')
 194         if not overwrite == 'yes':
 195             print "Aborting."
 196
 197             return False
 198
 199     if os.path.exists(args.cache_path):
 200         print 'The cache directory must not exist before you run this script'
 201         print 'Cache directory: ', args.cache_path
 202
 203         return False
 204
 205     return True
 206
 207
 208 def env_export(args):
 209     args = _setup_paths(args)
 210
 211     if not _check(args):
 212         print "\n== Checks did not pass, exiting ==\n"
 213         sys.exit(0)
 214
 215     setup_global_and_app_config(args.conf_file)
 216     setup_storage()
 217
 218     config, validation_result = read_mediagoblin_config(args.conf_file)
 219     connection, db = setup_connection_and_db_from_config(
 220         config['mediagoblin'], use_pymongo=True)
 221
 222     _export_database(db, args)
 223
 224     _export_media(db, args)
 225
 226     _create_archive(args)
 227
 228     _clean(args)