forked from databricks/databricks-sql-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathauth.py
More file actions
executable file
·210 lines (178 loc) · 8.13 KB
/
Copy pathauth.py
File metadata and controls
executable file
·210 lines (178 loc) · 8.13 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
from enum import Enum
from typing import Optional, List
from databricks.sql.auth.authenticators import (
AuthProvider,
AccessTokenAuthProvider,
ExternalAuthProvider,
CredentialsProvider,
DatabricksOAuthProvider,
)
class AuthType(Enum):
DATABRICKS_OAUTH = "databricks-oauth"
AZURE_OAUTH = "azure-oauth"
# TODO: Token federation should be a feature that works with different auth types,
# not an auth type itself. This will be refactored in a future change.
# We will add a use_token_federation flag that can be used with any auth type.
TOKEN_FEDERATION = "token-federation"
# other supported types (access_token) can be inferred
# we can add more types as needed later
class ClientContext:
def __init__(
self,
hostname: str,
access_token: Optional[str] = None,
auth_type: Optional[str] = None,
oauth_scopes: Optional[List[str]] = None,
oauth_client_id: Optional[str] = None,
oauth_redirect_port_range: Optional[List[int]] = None,
use_cert_as_auth: Optional[str] = None,
tls_client_cert_file: Optional[str] = None,
oauth_persistence=None,
credentials_provider=None,
identity_federation_client_id: Optional[str] = None,
):
self.hostname = hostname
self.access_token = access_token
self.auth_type = auth_type
self.oauth_scopes = oauth_scopes
self.oauth_client_id = oauth_client_id
self.oauth_redirect_port_range = oauth_redirect_port_range
self.use_cert_as_auth = use_cert_as_auth
self.tls_client_cert_file = tls_client_cert_file
self.oauth_persistence = oauth_persistence
self.credentials_provider = credentials_provider
self.identity_federation_client_id = identity_federation_client_id
def get_auth_provider(cfg: ClientContext):
"""
Get an appropriate auth provider based on the provided configuration.
Token Federation Support:
-----------------------
Currently, token federation is implemented as a separate auth type, but the goal is to
refactor it as a feature that can work with any auth type. The current implementation
is maintained for backward compatibility while the refactoring is planned.
Future refactoring will introduce a `use_token_federation` flag that can be combined
with any auth type to enable token federation.
Args:
cfg: The client context containing configuration parameters
Returns:
An appropriate AuthProvider instance
Raises:
RuntimeError: If no valid authentication settings are provided
"""
# If credentials_provider is explicitly provided
if cfg.credentials_provider:
# If token federation is enabled and credentials provider is provided,
# wrap the credentials provider with DatabricksTokenFederationProvider
if cfg.auth_type == AuthType.TOKEN_FEDERATION.value:
from databricks.sql.auth.token_federation import (
DatabricksTokenFederationProvider,
)
federation_provider = DatabricksTokenFederationProvider(
cfg.credentials_provider,
cfg.hostname,
cfg.identity_federation_client_id,
)
return ExternalAuthProvider(federation_provider)
# If not token federation, just use the credentials provider directly
return ExternalAuthProvider(cfg.credentials_provider)
# If we don't have a credentials provider but have token federation auth type with access token
if cfg.auth_type == AuthType.TOKEN_FEDERATION.value and cfg.access_token:
# Create a simple credentials provider and wrap it with token federation provider
from databricks.sql.auth.token_federation import (
DatabricksTokenFederationProvider,
SimpleCredentialsProvider,
)
simple_provider = SimpleCredentialsProvider(cfg.access_token)
federation_provider = DatabricksTokenFederationProvider(
simple_provider, cfg.hostname, cfg.identity_federation_client_id
)
return ExternalAuthProvider(federation_provider)
if cfg.auth_type in [AuthType.DATABRICKS_OAUTH.value, AuthType.AZURE_OAUTH.value]:
assert cfg.oauth_redirect_port_range is not None
assert cfg.oauth_client_id is not None
assert cfg.oauth_scopes is not None
return DatabricksOAuthProvider(
cfg.hostname,
cfg.oauth_persistence,
cfg.oauth_redirect_port_range,
cfg.oauth_client_id,
cfg.oauth_scopes,
cfg.auth_type,
)
elif cfg.access_token is not None:
return AccessTokenAuthProvider(cfg.access_token)
elif cfg.use_cert_as_auth and cfg.tls_client_cert_file:
# no op authenticator. authentication is performed using ssl certificate outside of headers
return AuthProvider()
else:
if (
cfg.oauth_redirect_port_range is not None
and cfg.oauth_client_id is not None
and cfg.oauth_scopes is not None
):
return DatabricksOAuthProvider(
cfg.hostname,
cfg.oauth_persistence,
cfg.oauth_redirect_port_range,
cfg.oauth_client_id,
cfg.oauth_scopes,
)
else:
raise RuntimeError("No valid authentication settings!")
PYSQL_OAUTH_SCOPES = ["sql", "offline_access"]
PYSQL_OAUTH_CLIENT_ID = "databricks-sql-python"
PYSQL_OAUTH_AZURE_CLIENT_ID = "96eecda7-19ea-49cc-abb5-240097d554f5"
PYSQL_OAUTH_REDIRECT_PORT_RANGE = list(range(8020, 8025))
PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE = [8030]
def normalize_host_name(hostname: str):
maybe_scheme = "https://" if not hostname.startswith("https://") else ""
maybe_trailing_slash = "/" if not hostname.endswith("/") else ""
return f"{maybe_scheme}{hostname}{maybe_trailing_slash}"
def get_client_id_and_redirect_port(use_azure_auth: bool):
return (
(PYSQL_OAUTH_CLIENT_ID, PYSQL_OAUTH_REDIRECT_PORT_RANGE)
if not use_azure_auth
else (PYSQL_OAUTH_AZURE_CLIENT_ID, PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE)
)
def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
"""
Get an auth provider for the Python SQL connector.
This function is the main entry point for authentication in the SQL connector.
It processes the parameters and creates an appropriate auth provider.
TODO: Future refactoring needed:
1. Add a use_token_federation flag that can be combined with any auth type
2. Remove TOKEN_FEDERATION as an auth_type while maintaining backward compatibility
3. Create a token federation wrapper that can wrap any existing auth provider
Args:
hostname: The Databricks server hostname
**kwargs: Additional configuration parameters
Returns:
An appropriate AuthProvider instance
Raises:
ValueError: If username/password authentication is attempted (no longer supported)
"""
auth_type = kwargs.get("auth_type")
(client_id, redirect_port_range) = get_client_id_and_redirect_port(
auth_type == AuthType.AZURE_OAUTH.value
)
if kwargs.get("username") or kwargs.get("password"):
raise ValueError(
"Username/password authentication is no longer supported. "
"Please use OAuth or access token instead."
)
cfg = ClientContext(
hostname=normalize_host_name(hostname),
auth_type=auth_type,
access_token=kwargs.get("access_token"),
use_cert_as_auth=kwargs.get("_use_cert_as_auth"),
tls_client_cert_file=kwargs.get("_tls_client_cert_file"),
oauth_scopes=PYSQL_OAUTH_SCOPES,
oauth_client_id=kwargs.get("oauth_client_id") or client_id,
oauth_redirect_port_range=[kwargs["oauth_redirect_port"]]
if kwargs.get("oauth_client_id") and kwargs.get("oauth_redirect_port")
else redirect_port_range,
oauth_persistence=kwargs.get("experimental_oauth_persistence"),
credentials_provider=kwargs.get("credentials_provider"),
identity_federation_client_id=kwargs.get("identity_federation_client_id"),
)
return get_auth_provider(cfg)