Cleanup error message when database does not exist

Fixes GH-5

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)

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()

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

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)