123 lines
3.9 KiB
Python
123 lines
3.9 KiB
Python
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Optional
|
|
|
|
|
|
@dataclass
|
|
class AzureEnvironment:
|
|
name: str
|
|
service_management_endpoint: str
|
|
resource_manager_endpoint: str
|
|
active_directory_endpoint: str
|
|
|
|
|
|
ARM_DATABRICKS_RESOURCE_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
|
|
|
|
ENVIRONMENTS = dict(
|
|
PUBLIC=AzureEnvironment(
|
|
name="PUBLIC",
|
|
service_management_endpoint="https://management.core.windows.net/",
|
|
resource_manager_endpoint="https://management.azure.com/",
|
|
active_directory_endpoint="https://login.microsoftonline.com/",
|
|
),
|
|
USGOVERNMENT=AzureEnvironment(
|
|
name="USGOVERNMENT",
|
|
service_management_endpoint="https://management.core.usgovcloudapi.net/",
|
|
resource_manager_endpoint="https://management.usgovcloudapi.net/",
|
|
active_directory_endpoint="https://login.microsoftonline.us/",
|
|
),
|
|
CHINA=AzureEnvironment(
|
|
name="CHINA",
|
|
service_management_endpoint="https://management.core.chinacloudapi.cn/",
|
|
resource_manager_endpoint="https://management.chinacloudapi.cn/",
|
|
active_directory_endpoint="https://login.chinacloudapi.cn/",
|
|
),
|
|
)
|
|
|
|
|
|
class Cloud(Enum):
|
|
AWS = "AWS"
|
|
AZURE = "AZURE"
|
|
GCP = "GCP"
|
|
|
|
|
|
@dataclass
|
|
class DatabricksEnvironment:
|
|
cloud: Cloud
|
|
dns_zone: str
|
|
azure_application_id: Optional[str] = None
|
|
azure_environment: Optional[AzureEnvironment] = None
|
|
|
|
def deployment_url(self, name: str) -> str:
|
|
return f"https://{name}{self.dns_zone}"
|
|
|
|
@property
|
|
def azure_service_management_endpoint(self) -> Optional[str]:
|
|
if self.azure_environment is None:
|
|
return None
|
|
return self.azure_environment.service_management_endpoint
|
|
|
|
@property
|
|
def azure_resource_manager_endpoint(self) -> Optional[str]:
|
|
if self.azure_environment is None:
|
|
return None
|
|
return self.azure_environment.resource_manager_endpoint
|
|
|
|
@property
|
|
def azure_active_directory_endpoint(self) -> Optional[str]:
|
|
if self.azure_environment is None:
|
|
return None
|
|
return self.azure_environment.active_directory_endpoint
|
|
|
|
|
|
DEFAULT_ENVIRONMENT = DatabricksEnvironment(Cloud.AWS, ".cloud.databricks.com")
|
|
|
|
ALL_ENVS = [
|
|
DatabricksEnvironment(Cloud.AWS, ".dev.databricks.com"),
|
|
DatabricksEnvironment(Cloud.AWS, ".staging.cloud.databricks.com"),
|
|
DatabricksEnvironment(Cloud.AWS, ".cloud.databricks.us"),
|
|
DEFAULT_ENVIRONMENT,
|
|
DatabricksEnvironment(
|
|
Cloud.AZURE,
|
|
".dev.azuredatabricks.net",
|
|
azure_application_id="62a912ac-b58e-4c1d-89ea-b2dbfc7358fc",
|
|
azure_environment=ENVIRONMENTS["PUBLIC"],
|
|
),
|
|
DatabricksEnvironment(
|
|
Cloud.AZURE,
|
|
".staging.azuredatabricks.net",
|
|
azure_application_id="4a67d088-db5c-48f1-9ff2-0aace800ae68",
|
|
azure_environment=ENVIRONMENTS["PUBLIC"],
|
|
),
|
|
DatabricksEnvironment(
|
|
Cloud.AZURE,
|
|
".azuredatabricks.net",
|
|
azure_application_id=ARM_DATABRICKS_RESOURCE_ID,
|
|
azure_environment=ENVIRONMENTS["PUBLIC"],
|
|
),
|
|
DatabricksEnvironment(
|
|
Cloud.AZURE,
|
|
".databricks.azure.us",
|
|
azure_application_id=ARM_DATABRICKS_RESOURCE_ID,
|
|
azure_environment=ENVIRONMENTS["USGOVERNMENT"],
|
|
),
|
|
DatabricksEnvironment(
|
|
Cloud.AZURE,
|
|
".databricks.azure.cn",
|
|
azure_application_id=ARM_DATABRICKS_RESOURCE_ID,
|
|
azure_environment=ENVIRONMENTS["CHINA"],
|
|
),
|
|
DatabricksEnvironment(Cloud.GCP, ".dev.gcp.databricks.com"),
|
|
DatabricksEnvironment(Cloud.GCP, ".staging.gcp.databricks.com"),
|
|
DatabricksEnvironment(Cloud.GCP, ".gcp.databricks.com"),
|
|
]
|
|
|
|
|
|
def get_environment_for_hostname(hostname: Optional[str]) -> DatabricksEnvironment:
|
|
if not hostname:
|
|
return DEFAULT_ENVIRONMENT
|
|
for env in ALL_ENVS:
|
|
if hostname.endswith(env.dns_zone):
|
|
return env
|
|
return DEFAULT_ENVIRONMENT
|