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 _import_media(db, args):
  55     """
  56     Import media files
  57
  58     Must be called after _import_database()
  59     """
  60     print "\n== Importing media ==\n"
  61
  62     media_cache = BasicFileStorage(
  63         args._cache_path['media'])
  64
  65     # TODO: Add import of queue files
  66     queue_cache = BasicFileStorage(
  67         args._cache_path['queue'])
  68
  69     for entry in db.media_entries.find():
  70         for name, path in entry['media_files'].items():
  71             media_file = mg_globals.public_store.get_file(path, mode='wb')
  72             media_file.write(
  73                 media_cache.get_file(path, mode='rb').read())
  74
  75             print(media_file)
  76         print(entry)
  77
  78     print "\n== Media imported ==\n"
  79
  80
  81 def _import_database(db, args):
  82     """
  83     Restore mongo database from ___.bson files
  84     """
  85     print "\n== Importing database ==\n"
  86
  87     p = subprocess.Popen([
  88             args.mongorestore_path,
  89             '-d', db.name,
  90             os.path.join(args._cache_path['database'], db.name)])
  91
  92     print p
  93
  94     p.wait()
  95
  96     print "\n== Database imported ==\n"
  97
  98
  99 def env_import(args):
 100     """
 101     Restore mongo database and media files from a tar archive
 102     """
 103     # args.cache_path += 'mediagoblin-data'
 104     setup_global_and_app_config(args.conf_file)
 105
 106     # Creates mg_globals.public_store and mg_globals.queue_store
 107     setup_storage()
 108
 109     config, validation_result = read_mediagoblin_config(args.conf_file)
 110     connection, db = setup_connection_and_db_from_config(
 111         config['mediagoblin'], use_pymongo=True)
 112
 113     tf = tarfile.open(
 114         args.tar_file,
 115         mode='r|gz')
 116
 117     tf.extractall(args.cache_path)
 118
 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 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     print "\n== Compressing to archive ==\n"
 150
 151     tf = tarfile.open(
 152         args.tar_file,
 153         mode='w|gz')
 154
 155     with tf:
 156         tf.add(args.cache_path, 'mediagoblin-data/')
 157
 158     print "\n== Archiving done ==\n"
 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     if os.path.exists(args.cache_path):
 183         print 'The cache directory must not exist before you run this script'
 184         print 'Cache directory: ', args.cache_path
 185
 186         return False
 187
 188     return True
 189
 190
 191 def _export_database(db, args):
 192     print "\n== Exporting database ==\n"
 193
 194     command = '{mongodump_path} -d {database} -o {mongodump_cache}'.format(
 195         mongodump_path=args.mongodump_path,
 196         database=db.name,
 197         mongodump_cache=args._cache_path['database'])
 198
 199     p = subprocess.Popen([
 200             args.mongodump_path,
 201             '-d', db.name,
 202             '-o', args._cache_path['database']])
 203
 204     p.wait()
 205
 206     print "\n== Database exported ==\n"
 207
 208
 209 def _export_media(db, args):
 210     print "\n== Exporting media ==\n"
 211
 212     media_cache = BasicFileStorage(
 213         args._cache_path['media'])
 214
 215     # TODO: Add export of queue files
 216     queue_cache = BasicFileStorage(
 217         args._cache_path['queue'])
 218
 219     for entry in db.media_entries.find():
 220         for name, path in entry['media_files'].items():
 221             mc_file = media_cache.get_file(path, mode='wb')
 222             mc_file.write(
 223                 mg_globals.public_store.get_file(path, mode='rb').read())
 224
 225             print(mc_file)
 226         print(entry)
 227
 228     print "\n== Media exported ==\n"
 229
 230
 231 def env_export(args):
 232     """
 233     Export database and media files to a tar archive
 234     """
 235     args = _setup_paths(args)
 236
 237     if not _export_check(args):
 238         print "\n== Checks did not pass, exiting ==\n"
 239         sys.exit(0)
 240
 241     setup_global_and_app_config(args.conf_file)
 242     setup_storage()
 243
 244     config, validation_result = read_mediagoblin_config(args.conf_file)
 245     connection, db = setup_connection_and_db_from_config(
 246         config['mediagoblin'], use_pymongo=True)
 247
 248     _export_database(db, args)
 249
 250     _export_media(db, args)
 251
 252     _create_archive(args)
 253
 254     _clean(args)