|
1 | 1 | import os |
2 | 2 | import tempfile |
3 | 3 | from os.path import join |
| 4 | +from unittest import mock |
4 | 5 |
|
5 | 6 | import keyring.backend |
6 | 7 | import pytest |
| 8 | +from databricks.sdk import config as _sdk_config |
7 | 9 |
|
8 | 10 | from dbt.adapters.databricks.credentials import DatabricksCredentials |
9 | 11 |
|
| 12 | +# databricks-sdk >=0.103 added a network probe to Config(); on older SDKs |
| 13 | +# this symbol doesn't exist and the lazy `_ensure_config` path is unnecessary. |
| 14 | +_REQUIRES_LAZY_CONFIG = pytest.mark.skipif( |
| 15 | + not hasattr(_sdk_config, "get_host_metadata"), |
| 16 | + reason="lazy _ensure_config is only required when Config() does a network probe", |
| 17 | +) |
| 18 | + |
| 19 | + |
| 20 | +_COMMON_KWARGS = { |
| 21 | + "host": "yourorg.databricks.com", |
| 22 | + "database": "andre", |
| 23 | + "http_path": "sql/protocolv1/o/1234567890123456/1234-cluster", |
| 24 | + "schema": "dbt", |
| 25 | +} |
| 26 | + |
| 27 | + |
| 28 | +@_REQUIRES_LAZY_CONFIG |
| 29 | +class TestParseTimeIsOffline: |
| 30 | + """`dbt parse/list/compile` must stay fully offline. Building |
| 31 | + `DatabricksCredentials` is on that path, so for every supported auth |
| 32 | + method we verify that constructing the credentials never calls |
| 33 | + `Config()` (which is what triggers the SDK's network I/O).""" |
| 34 | + |
| 35 | + def test_pat_credentials_init_does_not_call_config(self): |
| 36 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 37 | + DatabricksCredentials(token="foo", **_COMMON_KWARGS) |
| 38 | + mock_config.assert_not_called() |
| 39 | + |
| 40 | + def test_oauth_m2m_credentials_init_does_not_call_config(self): |
| 41 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 42 | + DatabricksCredentials( |
| 43 | + client_id="cid", |
| 44 | + client_secret="dose-secret", |
| 45 | + **_COMMON_KWARGS, |
| 46 | + ) |
| 47 | + mock_config.assert_not_called() |
| 48 | + |
| 49 | + def test_external_browser_credentials_init_does_not_call_config(self): |
| 50 | + # client_id with no client_secret triggers external-browser auth |
| 51 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 52 | + DatabricksCredentials(client_id="cid", **_COMMON_KWARGS) |
| 53 | + mock_config.assert_not_called() |
| 54 | + |
| 55 | + def test_azure_client_secret_credentials_init_does_not_call_config(self): |
| 56 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 57 | + DatabricksCredentials( |
| 58 | + azure_client_id="acid", |
| 59 | + azure_client_secret="asecret", |
| 60 | + **_COMMON_KWARGS, |
| 61 | + ) |
| 62 | + mock_config.assert_not_called() |
| 63 | + |
| 64 | + |
| 65 | +@_REQUIRES_LAZY_CONFIG |
| 66 | +class TestEnsureConfigTriggersTheRightAuth: |
| 67 | + """Connect-time counterpart to TestParseTimeIsOffline: when something |
| 68 | + actually does need the config (e.g. opening a connection), `_ensure_config` |
| 69 | + must build it via the correct auth method per credential shape. This also |
| 70 | + exercises every branch of the selection logic in `_ensure_config`.""" |
| 71 | + |
| 72 | + def test_pat_uses_pat_auth(self): |
| 73 | + creds = DatabricksCredentials(token="foo", **_COMMON_KWARGS) |
| 74 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 75 | + creds.authenticate().config |
| 76 | + mock_config.assert_called_once_with(host=_COMMON_KWARGS["host"], token="foo") |
| 77 | + |
| 78 | + def test_explicit_azure_uses_azure_client_secret(self): |
| 79 | + creds = DatabricksCredentials( |
| 80 | + azure_client_id="acid", |
| 81 | + azure_client_secret="asecret", |
| 82 | + auth_type="oauth", |
| 83 | + **_COMMON_KWARGS, |
| 84 | + ) |
| 85 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 86 | + creds.authenticate().config |
| 87 | + mock_config.assert_called_once_with( |
| 88 | + host=_COMMON_KWARGS["host"], |
| 89 | + azure_client_id="acid", |
| 90 | + azure_client_secret="asecret", |
| 91 | + auth_type="azure-client-secret", |
| 92 | + ) |
| 93 | + |
| 94 | + def test_client_id_only_uses_external_browser(self): |
| 95 | + creds = DatabricksCredentials(client_id="cid", auth_type="oauth", **_COMMON_KWARGS) |
| 96 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 97 | + creds.authenticate().config |
| 98 | + mock_config.assert_called_once_with( |
| 99 | + host=_COMMON_KWARGS["host"], |
| 100 | + client_id="cid", |
| 101 | + client_secret="", |
| 102 | + auth_type="external-browser", |
| 103 | + ) |
| 104 | + |
| 105 | + def test_dose_secret_tries_oauth_m2m_first(self): |
| 106 | + creds = DatabricksCredentials( |
| 107 | + client_id="cid", |
| 108 | + client_secret="dose-secret", |
| 109 | + auth_type="oauth", |
| 110 | + **_COMMON_KWARGS, |
| 111 | + ) |
| 112 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 113 | + creds.authenticate().config |
| 114 | + mock_config.assert_called_once_with( |
| 115 | + host=_COMMON_KWARGS["host"], |
| 116 | + client_id="cid", |
| 117 | + client_secret="dose-secret", |
| 118 | + auth_type="oauth-m2m", |
| 119 | + ) |
| 120 | + |
| 121 | + def test_non_dose_secret_tries_legacy_azure_first(self): |
| 122 | + creds = DatabricksCredentials( |
| 123 | + client_id="cid", |
| 124 | + client_secret="plain-secret", |
| 125 | + auth_type="oauth", |
| 126 | + **_COMMON_KWARGS, |
| 127 | + ) |
| 128 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 129 | + creds.authenticate().config |
| 130 | + mock_config.assert_called_once_with( |
| 131 | + host=_COMMON_KWARGS["host"], |
| 132 | + azure_client_id="cid", |
| 133 | + azure_client_secret="plain-secret", |
| 134 | + auth_type="azure-client-secret", |
| 135 | + ) |
| 136 | + |
| 137 | + def test_falls_back_to_second_method_when_first_raises(self): |
| 138 | + creds = DatabricksCredentials( |
| 139 | + client_id="cid", |
| 140 | + client_secret="plain-secret", |
| 141 | + auth_type="oauth", |
| 142 | + **_COMMON_KWARGS, |
| 143 | + ) |
| 144 | + fake_config = mock.MagicMock() |
| 145 | + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: |
| 146 | + mock_config.side_effect = [RuntimeError("first method fails"), fake_config] |
| 147 | + creds.authenticate().config |
| 148 | + assert mock_config.call_count == 2 |
| 149 | + # First call: legacy-azure (default order for non-dose secret). |
| 150 | + assert mock_config.call_args_list[0].kwargs["auth_type"] == "azure-client-secret" |
| 151 | + # Second call: oauth-m2m fallback. |
| 152 | + assert mock_config.call_args_list[1].kwargs["auth_type"] == "oauth-m2m" |
| 153 | + |
10 | 154 |
|
11 | 155 | @pytest.mark.skip(reason="Need to mock requests to OIDC") |
12 | 156 | class TestM2MAuth: |
|
0 commit comments