Merge pull request #74 from manics/pgdump

Parse config.xml for db parameters, add omero db dump

.travis.yml

@@ -33,6 +33,7 @@ before_script:
     - psql -c "select 1;" -U omero -h localhost omero
     - mkdir $HOME/OMERO
     - echo "config set omero.data.dir $HOME/OMERO" > $HOME/config.omero
+    - echo "config set omero.db.name omero" >> $HOME/config.omero
 script:
     - sh travis-build
 

omego/db.py

@@ -10,7 +10,7 @@
 import fileutils
 from external import External, RunException
 from yaclifw.framework import Command, Stop
-from env import EnvDefault, DbParser
+from env import DbParser
 
 log = logging.getLogger("omego.db")
 
@@ -30,17 +30,15 @@ def __init__(self, dir, command, args, external):
         if not os.path.exists(dir):
             raise Exception("%s does not exist!" % dir)
 
+        self.external = external
+
         psqlv = self.psql('--version')
         log.info('psql version: %s', psqlv)
 
-        self.external = external
-
         self.check_connection()
 
-        if command == 'init':
-            self.initialise()
-        elif command == 'upgrade':
-            self.upgrade()
+        if command in ('init', 'upgrade', 'dump'):
+            getattr(self, command)()
         else:
             raise Stop('Invalid db command: %s', command)
 
@@ -51,7 +49,7 @@ def check_connection(self):
             log.error(e)
             raise Stop(30, 'Database connection check failed')
 
-    def initialise(self):
+    def init(self):
         omerosql = self.args.omerosql
         autoupgrade = False
         if not omerosql:
@@ -161,23 +159,84 @@ def get_current_db_version(self):
         log.info('Current omero db version: %s', v)
         return v
 
-    def psql(self, *psqlargs):
+    def dump(self):
         """
-        Run a psql command
+        Dump the database using the postgres custom format
         """
-        if not self.args.dbname:
+        dumpfile = self.args.dumpfile
+        if not dumpfile:
+            db, env = self.get_db_args_env()
+            dumpfile = fileutils.timestamp_filename(
+                'omero-database-%s' % db['name'], 'pgdump')
+
+        log.info('Dumping database to %s', dumpfile)
+        if not self.args.dry_run:
+            self.pgdump('-Fc', '-f', dumpfile)
+
+    def get_db_args_env(self):
+        """
+        Get a dictionary of database connection parameters, and create an
+        environment for running postgres commands.
+        Falls back to omego defaults.
+        """
+        db = {
+            'name': self.args.dbname,
+            'host': self.args.dbhost,
+            'user': self.args.dbuser,
+            'pass': self.args.dbpass
+            }
+
+        if not self.args.no_db_config:
+            try:
+                c = self.external.get_config(force=True)
+            except Exception as e:
+                log.warn('config.xml not found: %s', e)
+                c = {}
+
+            for k in db:
+                try:
+                    db[k] = c['omero.db.%s' % k]
+                except KeyError:
+                    log.info(
+                        'Failed to lookup parameter omero.db.%s, using %s',
+                        k, db[k])
+
+        if not db['name']:
             raise Exception('Database name required')
 
         env = os.environ.copy()
-        env['PGPASSWORD'] = self.args.dbpass
-        args = ['-d', self.args.dbname, '-h', self.args.dbhost, '-U',
-                self.args.dbuser, '-w', '-A', '-t'] + list(psqlargs)
+        env['PGPASSWORD'] = db['pass']
+        return db, env
+
+    def psql(self, *psqlargs):
+        """
+        Run a psql command
+        """
+        db, env = self.get_db_args_env()
+
+        args = ['-d', db['name'], '-h', db['host'], '-U', db['user'],
+                '-w', '-A', '-t'] + list(psqlargs)
         stdout, stderr = External.run('psql', args, capturestd=True, env=env)
         if stderr:
             log.warn('stderr: %s', stderr)
         log.debug('stdout: %s', stdout)
         return stdout
 
+    def pgdump(self, *pgdumpargs):
+        """
+        Run a pg_dump command
+        """
+        db, env = self.get_db_args_env()
+
+        args = ['-d', db['name'], '-h', db['host'], '-U', db['user'], '-w'
+                ] + list(pgdumpargs)
+        stdout, stderr = External.run(
+            'pg_dump', args, capturestd=True, env=env)
+        if stderr:
+            log.warn('stderr: %s', stderr)
+        log.debug('stdout: %s', stdout)
+        return stdout
+
 
 class DbCommand(Command):
     """
@@ -192,13 +251,14 @@ def __init__(self, sub_parsers):
         self.parser = DbParser(self.parser)
         self.parser.add_argument("-n", "--dry-run", action="store_true")
 
-        Add = EnvDefault.add
         # TODO: Kind of duplicates Upgrade args.sym/args.server
-        Add(self.parser, 'serverdir', 'Root directory of the server')
+        self.parser.add_argument(
+            '--serverdir', help='Root directory of the server')
         self.parser.add_argument(
             "dbcommand",
-            choices=['init', 'upgrade'],
+            choices=['init', 'upgrade', 'dump'],
             help='Initialise or upgrade a database')
+        self.parser.add_argument('--dumpfile', help='Database dump file')
 
     def __call__(self, args):
         super(DbCommand, self).__call__(args)

omego/env.py

@@ -86,6 +86,9 @@ def __init__(self, parser):
             help="Username for connecting to the OMERO database")
         Add(group, "dbpass", "omero",
             help="Password for connecting to the OMERO database")
+        group.add_argument(
+            "--no-db-config", action="store_true",
+            help="Ignore the database settings in omero config")
         # TODO Admin credentials: dbauser, dbapass
 
         Add(group, "omerosql", None,

omego/external.py

@@ -49,6 +49,8 @@ def __init__(self, dir=None):
         if dir:
             self.set_server_dir(dir)
 
+        self._omero = None
+
     def set_server_dir(self, dir):
         """
         Set the directory of the server to be controlled
@@ -68,13 +70,43 @@ def has_config(self):
             raise Exception('No server directory set')
         return self.configured
 
+    def get_config(self, force=False):
+        """
+        Returns a dictionary of all config.xml properties
+
+        If `force = True` then ignore any cached state and read config.xml
+        if possible
+
+        setup_omero_cli() must be called before this method to import the
+        correct omero module to minimise the possibility of version conflicts
+        """
+        if not force and not self.has_config():
+            raise Exception('No config file')
+
+        configxml = os.path.join(self.dir, 'etc', 'grid', 'config.xml')
+        if not os.path.exists(configxml):
+            raise Exception('No config file')
+
+        try:
+            # Attempt to open config.xml read-only, though this flag is not
+            # present in early versions of OMERO 5.0
+            c = self._omero.config.ConfigXml(
+                configxml, exclusive=False, read_only=True)
+        except TypeError:
+            c = self._omero.config.ConfigXml(configxml, exclusive=False)
+
+        try:
+            return c.as_map()
+        finally:
+            c.close()
+
     def setup_omero_cli(self):
         """
         Imports the omero CLI module so that commands can be run directly.
         Note Python does not allow a module to be imported multiple times,
         so this will only work with a single omero instance.
 
-        This can have several surprisingly effects, so setup_omero_cli()
+        This can have several surprising effects, so setup_omero_cli()
         must be explcitly called.
         """
         if not self.dir:
@@ -96,6 +128,7 @@ def setup_omero_cli(self):
 
         self.cli = omero.cli.CLI()
         self.cli.loadplugins()
+        self._omero = omero
 
     def setup_previous_omero_env(self, olddir, savevarsfile):
         """

test/unit/test_db.py

@@ -74,7 +74,7 @@ def test_check_connection(self, connected):
 
     @pytest.mark.parametrize('sqlfile', ['exists', 'missing', 'notprovided'])
     @pytest.mark.parametrize('dryrun', [True, False])
-    def test_initialise(self, sqlfile, dryrun):
+    def test_init(self, sqlfile, dryrun):
         ext = self.mox.CreateMock(External)
         if sqlfile != 'notprovided':
             omerosql = 'omero.sql'
@@ -109,10 +109,10 @@ def test_initialise(self, sqlfile, dryrun):
 
         if sqlfile == 'missing':
             with pytest.raises(Stop) as excinfo:
-                db.initialise()
+                db.init()
             assert str(excinfo.value) == 'SQL file not found'
         else:
-            db.initialise()
+            db.init()
         self.mox.VerifyAll()
 
     def test_sort_schema(self):
@@ -203,28 +203,111 @@ def test_get_current_db_version(self):
         assert db.get_current_db_version() == ('OMERO4.4', '0')
         self.mox.VerifyAll()
 
+    @pytest.mark.parametrize('dumpfile', ['test.pgdump', None])
+    @pytest.mark.parametrize('dryrun', [True, False])
+    def test_dump(self, dumpfile, dryrun):
+        args = self.Args({'dry_run': dryrun, 'dumpfile': dumpfile})
+        db = self.PartialMockDb(args, None)
+        self.mox.StubOutWithMock(omego.fileutils, 'timestamp_filename')
+        self.mox.StubOutWithMock(db, 'get_db_args_env')
+        self.mox.StubOutWithMock(db, 'pgdump')
+
+        if not dumpfile:
+            db.get_db_args_env().AndReturn(self.create_db_test_params())
+
+            dumpfile = 'omero-database-name-00000000-000000-000000.pgdump'
+            omego.fileutils.timestamp_filename(
+                'omero-database-name', 'pgdump').AndReturn(dumpfile)
+
+        if not dryrun:
+            db.pgdump('-Fc', '-f', dumpfile).AndReturn('')
+
+        self.mox.ReplayAll()
+
+        db.dump()
+        self.mox.VerifyAll()
+
+    def create_db_test_params(self, prefix=''):
+        db = {
+            'name': '%sname' % prefix,
+            'host': '%shost' % prefix,
+            'user': '%suser' % prefix,
+            'pass': '%spass' % prefix,
+        }
+        env = {'PGPASSWORD': '%spass' % prefix}
+        return db, env
+
     @pytest.mark.parametrize('dbname', ['name', ''])
-    def test_psql(self, dbname):
+    @pytest.mark.parametrize('hasconfig', [True, False])
+    @pytest.mark.parametrize('noconfig', [True, False])
+    def test_get_db_args_env(self, dbname, hasconfig, noconfig):
+        ext = self.mox.CreateMock(External)
         args = self.Args({'dbhost': 'host', 'dbname': dbname,
-                          'dbuser': 'user', 'dbpass': 'pass'})
-
+                          'dbuser': 'user', 'dbpass': 'pass',
+                          'no_db_config': noconfig})
+        db = self.PartialMockDb(args, ext)
+        self.mox.StubOutWithMock(db.external, 'has_config')
+        self.mox.StubOutWithMock(db.external, 'get_config')
         self.mox.StubOutWithMock(os.environ, 'copy')
-        self.mox.StubOutWithMock(External, 'run')
 
-        if dbname:
-            os.environ.copy().AndReturn({'PGPASSWORD': 'incorrect'})
-            psqlargs = ['-d', dbname, '-h', 'host', '-U', 'user',
-                        '-w', '-A', '-t', 'arg1', 'arg2']
-            External.run('psql', psqlargs, capturestd=True,
-                         env={'PGPASSWORD': 'pass'}).AndReturn(('', ''))
-        self.mox.ReplayAll()
+        if noconfig or not hasconfig:
+            expecteddb, expectedenv = self.create_db_test_params()
+        else:
+            expecteddb, expectedenv = self.create_db_test_params('ext')
 
-        db = self.PartialMockDb(args, None)
+        if not noconfig:
+            cfg = {}
+            if hasconfig:
+                cfg = {
+                    'omero.db.host': 'exthost',
+                    'omero.db.user': 'extuser',
+                    'omero.db.pass': 'extpass',
+                }
+                if dbname:
+                    cfg['omero.db.name'] = 'extname'
+
+                db.external.get_config(force=True).AndReturn(cfg)
+            else:
+                db.external.get_config().AndRaise(Exception())
+
+        os.environ.copy().AndReturn({'PGPASSWORD': 'incorrect'})
+
+        self.mox.ReplayAll()
         if dbname:
-            db.psql('arg1', 'arg2')
+            rcfg, renv = db.get_db_args_env()
+            assert rcfg == expecteddb
+            assert renv == expectedenv
         else:
             with pytest.raises(Exception) as excinfo:
-                db.psql('arg1', 'arg2')
+                db.get_db_args_env()
             assert str(excinfo.value) == 'Database name required'
 
+    def test_psql(self):
+        db = self.PartialMockDb(None, None)
+        self.mox.StubOutWithMock(db, 'get_db_args_env')
+        self.mox.StubOutWithMock(External, 'run')
+
+        psqlargs = ['-d', 'name', '-h', 'host', '-U', 'user',
+                    '-w', '-A', '-t', 'arg1', 'arg2']
+        db.get_db_args_env().AndReturn(self.create_db_test_params())
+        External.run('psql', psqlargs, capturestd=True,
+                     env={'PGPASSWORD': 'pass'}).AndReturn(('', ''))
+        self.mox.ReplayAll()
+
+        db.psql('arg1', 'arg2')
+        self.mox.VerifyAll()
+
+    def test_pgdump(self):
+        db = self.PartialMockDb(None, None)
+        self.mox.StubOutWithMock(db, 'get_db_args_env')
+        self.mox.StubOutWithMock(External, 'run')
+
+        pgdumpargs = ['-d', 'name', '-h', 'host', '-U', 'user',
+                      '-w', 'arg1', 'arg2']
+        db.get_db_args_env().AndReturn(self.create_db_test_params())
+        External.run('pg_dump', pgdumpargs, capturestd=True,
+                     env={'PGPASSWORD': 'pass'}).AndReturn(('', ''))
+        self.mox.ReplayAll()
+
+        db.pgdump('arg1', 'arg2')
         self.mox.VerifyAll()

test/unit/test_external.py

@@ -85,6 +85,9 @@ def test_set_server_dir_and_has_config(self, tmpdir, configured):
         tmpdir.ensure('etc', 'grid', 'config.xml')
         assert self.ext.has_config() == configured
 
+    # def test_get_config(self):
+    # Not easily testable since it requires the omero module
+
     # def test_setup_omero_cli(self):
     # Not easily testable since it does a direct import