mediagoblin/gmg_commands/import_export.py

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