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(args._cache_path['queue'])
  65
  66     for entry in db.MediaEntry.query.filter_by():
  67         for name, path in entry.media_files.items():
  68             _log.info('Importing: {0} - {1}'.format(
  69                     entry.title.encode('ascii', 'replace'),
  70                     name))
  71
  72             media_file = mg_globals.public_store.get_file(path, mode='wb')
  73             media_file.write(
  74                 media_cache.get_file(path, mode='rb').read())
  75
  76     _log.info('...Media imported')
  77
  78
  79 def _import_database(db, args):
  80     '''
  81     Restore mongo database from ___.bson files
  82     '''
  83     _log.info('-> Importing database...')
  84
  85     p = subprocess.Popen([
  86             args.mongorestore_path,
  87             '-d', db.name,
  88             os.path.join(args._cache_path['database'], db.name)])
  89
  90     p.wait()
  91
  92     _log.info('...Database imported')
  93
  94
  95 def env_import(args):
  96     '''
  97     Restore mongo database and media files from a tar archive
  98     '''
  99     if not args.cache_path:
 100         args.cache_path = tempfile.mkdtemp()
 101
 102     setup_global_and_app_config(args.conf_file)
 103
 104     # Creates mg_globals.public_store and mg_globals.queue_store
 105     setup_storage()
 106
 107     global_config, app_config = setup_global_and_app_config(args.conf_file)
 108     db = setup_connection_and_db_from_config(
 109         app_config)
 110
 111     tf = tarfile.open(
 112         args.tar_file,
 113         mode='r|gz')
 114
 115     tf.extractall(args.cache_path)
 116
 117     args.cache_path = os.path.join(
 118         args.cache_path, 'mediagoblin-data')
 119     args = _setup_paths(args)
 120
 121     # Import database from extracted data
 122     _import_database(db, args)
 123
 124     _import_media(db, args)
 125
 126     _clean(args)
 127
 128
 129 def _setup_paths(args):
 130     '''
 131     Populate ``args`` variable with cache subpaths
 132     '''
 133     args._cache_path = dict()
 134     PATH_MAP = {
 135         'media': 'media',
 136         'queue': 'queue',
 137         'database': 'database'}
 138
 139     for key, val in PATH_MAP.items():
 140         args._cache_path[key] = os.path.join(args.cache_path, val)
 141
 142     return args
 143
 144
 145 def _create_archive(args):
 146     '''
 147     Create the tar archive
 148     '''
 149     _log.info('-> Compressing to archive')
 150
 151     tf = tarfile.open(
 152         args.tar_file,
 153         mode='w|gz')
 154
 155     with closing(tf):
 156         tf.add(args.cache_path, 'mediagoblin-data/')
 157
 158     _log.info('...Archiving done')
 159
 160
 161 def _clean(args):
 162     '''
 163     Remove cache directory
 164     '''
 165     shutil.rmtree(args.cache_path)
 166
 167
 168 def _export_check(args):
 169     '''
 170     Run security checks for export command
 171     '''
 172     if os.path.exists(args.tar_file):
 173         overwrite = raw_input(
 174             'The output file already exists. '
 175             'Are you **SURE** you want to overwrite it? '
 176             '(yes/no)> ')
 177         if not overwrite == 'yes':
 178             print 'Aborting.'
 179
 180             return False
 181
 182     return True
 183
 184
 185 def _export_database(db, args):
 186     _log.info('-> Exporting database...')
 187
 188     p = subprocess.Popen([
 189             args.mongodump_path,
 190             '-d', db.name,
 191             '-o', args._cache_path['database']])
 192
 193     p.wait()
 194
 195     _log.info('...Database exported')
 196
 197
 198 def _export_media(db, args):
 199     _log.info('-> Exporting media...')
 200
 201     media_cache = BasicFileStorage(
 202         args._cache_path['media'])
 203
 204     # TODO: Add export of queue files
 205     queue_cache = BasicFileStorage(args._cache_path['queue'])
 206
 207     for entry in db.MediaEntry.query.filter_by():
 208         for name, path in entry.media_files.items():
 209             _log.info(u'Exporting {0} - {1}'.format(
 210                     entry.title,
 211                     name))
 212             try:
 213                 mc_file = media_cache.get_file(path, mode='wb')
 214                 mc_file.write(
 215                     mg_globals.public_store.get_file(path, mode='rb').read())
 216             except Exception as e:
 217                 _log.error('Failed: {0}'.format(e))
 218
 219     _log.info('...Media exported')
 220
 221
 222 def env_export(args):
 223     '''
 224     Export database and media files to a tar archive
 225     '''
 226     if args.cache_path:
 227         if os.path.exists(args.cache_path):
 228             _log.error('The cache directory must not exist '
 229                        'before you run this script')
 230             _log.error('Cache directory: {0}'.format(args.cache_path))
 231
 232             return False
 233     else:
 234         args.cache_path = tempfile.mkdtemp()
 235
 236     args = _setup_paths(args)
 237
 238     if not _export_check(args):
 239         _log.error('Checks did not pass, exiting')
 240         sys.exit(0)
 241
 242     globa_config, app_config = setup_global_and_app_config(args.conf_file)
 243
 244     setup_storage()
 245
 246     db = setup_connection_and_db_from_config(app_config)
 247
 248     _export_database(db, args)
 249
 250     _export_media(db, args)
 251
 252     _create_archive(args)
 253
 254     _clean(args)