diff --git a/worek/cli.py b/worek/cli.py index 70d4e70..0532a1b 100644 --- a/worek/cli.py +++ b/worek/cli.py @@ -47,5 +47,16 @@ def restore(host, port, user, dbname, engine, schema, restore_file, file_format) ' automatically determine the file format when using a pipe.', ) - core.restore(file_name, schema=schema, host=host, port=port, user=user, dbname=dbname, - file_format=file_format) + __import__('pdb').set_trace() + try: + core.restore( + file_name, + schema=schema, + host=host, + port=port, + user=user, + dbname=dbname, + file_format=file_format, + ) + except core.WorekOperationException as e: + click.echo(str(e), err=True) diff --git a/worek/core.py b/worek/core.py index 8067c82..02993ec 100644 --- a/worek/core.py +++ b/worek/core.py @@ -27,7 +27,9 @@ def backup(backup_file, backup_type='full', **params): schemas=params.get('schemas')) if not PG.engine_can_connect: - raise WorekOperationException('Can\'t connect to the database.') + raise WorekOperationException( + "Can't connect to the database. Trying to connect to {}.".format(repr(PG.engine.url)) + ) if backup_type == 'full': PG.backup_binary(backup_file) @@ -58,7 +60,9 @@ def restore(restore_file, file_format=None, clean_existing_database=True, **para schemas=params.get('schemas')) if not PG.engine_can_connect: - raise WorekOperationException('Can\'t connect to the database.') + raise WorekOperationException( + "Can't connect to the database. Trying to connect to {}.".format(repr(PG.engine.url)) + ) if clean_existing_database: PG.clean_existing_database() diff --git a/worek/dialects/postgres.py b/worek/dialects/postgres.py index ce1e5f2..b6aec33 100644 --- a/worek/dialects/postgres.py +++ b/worek/dialects/postgres.py @@ -72,7 +72,7 @@ def engine_can_connect(self): try: self.engine.execute("SELECT 1") return True - except psycopg2.Error: + except (psycopg2.Error, sa.exc.OperationalError): return False @classmethod diff --git a/worek/tests/test_core.py b/worek/tests/test_core.py index 7608267..929eea9 100644 --- a/worek/tests/test_core.py +++ b/worek/tests/test_core.py @@ -1,6 +1,9 @@ +import io + +import pytest import sqlalchemy as sa -import worek +import worek from worek.tests.helpers import PostgresDialectTestBase @@ -25,3 +28,17 @@ def test_backup_creates_full_restoreable_backup(self, tmpdir, pg_clean_engine): pg_clean_engine.execute('SELECT * FROM should_be_removed;') except sa.exc.ProgrammingError as e: assert 'relation "should_be_removed" does not exist' in str(e) + + def test_exception_with_no_database_backup(self): + with pytest.raises(worek.core.WorekOperationException) as e: + worek.core.backup('back.txt') + + assert "Can't connect to the database. Trying to connect to " in str(e.value) + + +class TestCorePGRestore(PostgresDialectTestBase): + def test_exception_with_no_database_restore(self,): + with pytest.raises(worek.core.WorekOperationException) as e: + worek.core.restore(io.StringIO('test')) + + assert "Can't connect to the database. Trying to connect to " in str(e.value)