diff --git a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py index c689324f97dca..d86d0453fa592 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py @@ -179,7 +179,6 @@ def sqlalchemy_url(self) -> URL: :return: the extracted sqlalchemy.engine.URL object. """ - conn = self.get_conn() url_query = { "http_path": self._http_path, "catalog": self.catalog, @@ -189,9 +188,8 @@ def sqlalchemy_url(self) -> URL: return URL.create( drivername="databricks", username="token", - password=conn.password, - host=conn.host, - port=conn.port, + password=self._get_token(raise_error=True), + host=self.host, query=url_query, ) diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py index 9e2ee58d60069..38fd5e26bdd4e 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py @@ -39,7 +39,6 @@ DEFAULT_CONN_ID = "databricks_default" HOST = "xx.cloud.databricks.com" HOST_WITH_SCHEME = "https://xx.cloud.databricks.com" -PORT = 443 TOKEN = "token" HTTP_PATH = "sql/protocolv1/o/1234567890123456/0123-456789-abcd123" SCHEMA = "test_schema" @@ -112,38 +111,36 @@ def mock_timer(): yield mock_timer -def make_mock_connection(): - return Connection( - conn_id=DEFAULT_CONN_ID, - conn_type="databricks", - host=HOST, - port=PORT, - login="token", - password=TOKEN, - ) - - -def test_sqlachemy_url_property(mock_get_conn): - mock_get_conn.return_value = make_mock_connection() +def test_sqlachemy_url_property(): hook = DatabricksSqlHook( databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA ) url = hook.sqlalchemy_url.render_as_string(hide_password=False) expected_url = ( - f"databricks://token:{TOKEN}@{HOST}:{PORT}?" + f"databricks://token:{TOKEN}@{HOST}?" f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}" ) assert url == expected_url -def test_get_uri(mock_get_conn): - mock_get_conn.return_value = make_mock_connection() +def test_get_sqlalchemy_engine(): + hook = DatabricksSqlHook( + databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA + ) + engine = hook.get_sqlalchemy_engine() + assert engine.url.render_as_string(hide_password=False) == ( + f"databricks://token:{TOKEN}@{HOST}?" + f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}" + ) + + +def test_get_uri(): hook = DatabricksSqlHook( databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA ) uri = hook.get_uri() expected_uri = ( - f"databricks://token:{TOKEN}@{HOST}:{PORT}?" + f"databricks://token:{TOKEN}@{HOST}?" f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}" ) assert uri == expected_uri