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 import mg_globals
  18 from mediagoblin.db.open import setup_connection_and_db_from_config
  19 from mediagoblin.init.config import read_mediagoblin_config
  20 from mediagoblin.storage import BasicFileStorage
  21 from mediagoblin.init import setup_storage, setup_global_and_app_config
  22
  23 import shutil
  24 import tarfile
  25 import subprocess
  26 import os.path
  27 import os
  28 import sys
  29 from contextlib import closing
  30
  31
  32 def import_export_parse_setup(subparser):
  33     # TODO: Add default
  34     subparser.add_argument(
  35         'tar_file')
  36     subparser.add_argument(
  37         '-cf', '--conf_file', default='mediagoblin.ini',
  38         help='Config file used to set up environment')
  39     subparser.add_argument(
  40         '--mongodump_path', default='mongodump',
  41         help='mongodump binary')
  42     subparser.add_argument(
  43         '--mongorestore_path', default='mongorestore',
  44         help='mongorestore binary')
  45     subparser.add_argument(
  46         '--cache_path', default='/tmp/mediagoblin/',
  47         help='')
  48
  49
  50 def _import_media(db, args):
  51     """
  52     Import media files
  53
  54     Must be called after _import_database()
  55     """
  56     print "\n== Importing media ==\n"
  57
  58     media_cache = BasicFileStorage(
  59         args._cache_path['media'])
  60
  61     # TODO: Add import of queue files
  62     queue_cache = BasicFileStorage(
  63         args._cache_path['queue'])
  64
  65     for entry in db.media_entries.find():
  66         for name, path in entry['media_files'].items():
  67             media_file = mg_globals.public_store.get_file(path, mode='wb')
  68             media_file.write(
  69                 media_cache.get_file(path, mode='rb').read())
  70
  71     print "\n== Media imported ==\n"
  72
  73
  74 def _import_database(db, args):
  75     """
  76     Restore mongo database from ___.bson files
  77     """
  78     print "\n== Importing database ==\n"
  79
  80     p = subprocess.Popen([
  81             args.mongorestore_path,
  82             '-d', db.name,
  83             os.path.join(args._cache_path['database'], db.name)])
  84
  85     p.wait()
  86
  87     print "\n== Database imported ==\n"
  88
  89
  90 def env_import(args):
  91     """
  92     Restore mongo database and media files from a tar archive
  93     """
  94     # args.cache_path += 'mediagoblin-data'
  95     setup_global_and_app_config(args.conf_file)
  96
  97     # Creates mg_globals.public_store and mg_globals.queue_store
  98     setup_storage()
  99
 100     config, validation_result = read_mediagoblin_config(args.conf_file)
 101     connection, db = setup_connection_and_db_from_config(
 102         config['mediagoblin'], use_pymongo=True)
 103
 104     tf = tarfile.open(
 105         args.tar_file,
 106         mode='r|gz')
 107
 108     tf.extractall(args.cache_path)
 109
 110     args.cache_path += 'mediagoblin-data'
 111     args = _setup_paths(args)
 112
 113     # Import database from extracted data
 114     _import_database(db, args)
 115
 116     _import_media(db, args)
 117
 118     # _clean(args)
 119
 120 def _setup_paths(args):
 121     """
 122     Populate ``args`` variable with cache subpaths
 123     """
 124     args._cache_path = dict()
 125     PATH_MAP = {
 126         'media': 'media',
 127         'queue': 'queue',
 128         'database': 'database'}
 129
 130     for key, val in PATH_MAP.items():
 131         args._cache_path[key] = os.path.join(args.cache_path, val)
 132
 133     return args
 134
 135
 136 def _create_archive(args):
 137     """
 138     Create the tar archive
 139     """
 140     print "\n== Compressing to archive ==\n"
 141
 142     tf = tarfile.open(
 143         args.tar_file,
 144         mode='w|gz')
 145
 146     with closing(tf):
 147         tf.add(args.cache_path, 'mediagoblin-data/')
 148
 149     print "\n== Archiving done ==\n"
 150
 151
 152 def _clean(args):
 153     """
 154     Remove cache directory
 155     """
 156     shutil.rmtree(args.cache_path)
 157
 158
 159 def _export_check(args):
 160     """
 161     Run security checks for export command
 162     """
 163     if os.path.exists(args.tar_file):
 164         overwrite = raw_input(
 165             'The output file already exists. '
 166             'Are you **SURE** you want to overwrite it? '
 167             '(yes/no)> ')
 168         if not overwrite == 'yes':
 169             print "Aborting."
 170
 171             return False
 172
 173     if os.path.exists(args.cache_path):
 174         print 'The cache directory must not exist before you run this script'
 175         print 'Cache directory: ', args.cache_path
 176
 177         return False
 178
 179     return True
 180
 181
 182 def _export_database(db, args):
 183     print "\n== Exporting database ==\n"
 184
 185     command = '{mongodump_path} -d {database} -o {mongodump_cache}'.format(
 186         mongodump_path=args.mongodump_path,
 187         database=db.name,
 188         mongodump_cache=args._cache_path['database'])
 189
 190     p = subprocess.Popen([
 191             args.mongodump_path,
 192             '-d', db.name,
 193             '-o', args._cache_path['database']])
 194
 195     p.wait()
 196
 197     print "\n== Database exported ==\n"
 198
 199
 200 def _export_media(db, args):
 201     print "\n== Exporting media ==\n"
 202
 203     media_cache = BasicFileStorage(
 204         args._cache_path['media'])
 205
 206     # TODO: Add export of queue files
 207     queue_cache = BasicFileStorage(
 208         args._cache_path['queue'])
 209
 210     for entry in db.media_entries.find():
 211         for name, path in entry['media_files'].items():
 212             mc_file = media_cache.get_file(path, mode='wb')
 213             mc_file.write(
 214                 mg_globals.public_store.get_file(path, mode='rb').read())
 215
 216     print "\n== Media exported ==\n"
 217
 218
 219 def env_export(args):
 220     """
 221     Export database and media files to a tar archive
 222     """
 223     args = _setup_paths(args)
 224
 225     if not _export_check(args):
 226         print "\n== Checks did not pass, exiting ==\n"
 227         sys.exit(0)
 228
 229     setup_global_and_app_config(args.conf_file)
 230     setup_storage()
 231
 232     config, validation_result = read_mediagoblin_config(args.conf_file)
 233     connection, db = setup_connection_and_db_from_config(
 234         config['mediagoblin'], use_pymongo=True)
 235
 236     _export_database(db, args)
 237
 238     _export_media(db, args)
 239
 240     _create_archive(args)
 241
 242     _clean(args)