diff --git a/airflow/hooks/dbapi_hook.py b/airflow/hooks/dbapi_hook.py index cf02385f61c8db..70a30c7865ded9 100644 --- a/airflow/hooks/dbapi_hook.py +++ b/airflow/hooks/dbapi_hook.py @@ -74,8 +74,11 @@ def get_uri(self): host = conn.host if conn.port is not None: host += ':{port}'.format(port=conn.port) - return '{conn.conn_type}://{login}{host}/{conn.schema}'.format( + uri = '{conn.conn_type}://{login}{host}/'.format( conn=conn, login=login, host=host) + if conn.schema: + uri += conn.schema + return uri def get_sqlalchemy_engine(self, engine_kwargs=None): if engine_kwargs is None: diff --git a/tests/hooks/test_dbapi_hook.py b/tests/hooks/test_dbapi_hook.py index 911d4f86c7b32b..afbce87be9f073 100644 --- a/tests/hooks/test_dbapi_hook.py +++ b/tests/hooks/test_dbapi_hook.py @@ -22,6 +22,7 @@ from unittest import mock from airflow.hooks.dbapi_hook import DbApiHook +from airflow.models import Connection class TestDbApiHook(unittest.TestCase): @@ -149,3 +150,25 @@ def test_insert_rows_commit_every(self): sql = "INSERT INTO {} VALUES (%s)".format(table) for row in rows: self.cur.execute.assert_any_call(sql, row) + + def test_get_uri_schema_not_none(self): + self.db_hook.get_connection = mock.MagicMock(return_value=Connection( + conn_type="conn_type", + host="host", + login="login", + password="password", + schema="schema", + port=1 + )) + self.assertEqual("conn_type://login:password@host:1/schema", self.db_hook.get_uri()) + + def test_get_uri_schema_none(self): + self.db_hook.get_connection = mock.MagicMock(return_value=Connection( + conn_type="conn_type", + host="host", + login="login", + password="password", + schema=None, + port=1 + )) + self.assertEqual("conn_type://login:password@host:1/", self.db_hook.get_uri())