feat: Add idp_partition connection option

Author: Jonathan Glaser

redshift_connector/__init__.py

@@ -153,6 +153,7 @@ def connect(
  client_secret: typing.Optional[str] = None,
  partner_sp_id: typing.Optional[str] = None,
  idp_response_timeout: typing.Optional[int] = None,
+ idp_partition: typing.Optional[str] = None,
  listen_port: typing.Optional[int] = None,
  login_to_rp: typing.Optional[str] = None,
  login_url: typing.Optional[str] = None,
@@ -334,6 +335,7 @@ def connect(
  info.put("idc_region", idc_region)
  info.put("identity_namespace", identity_namespace)
  info.put("idp_host", idp_host)
+ info.put("idp_partition", idp_partition)
  info.put("idp_response_timeout", idp_response_timeout)
  info.put("idp_tenant", idp_tenant)
  info.put("issuer_url", issuer_url)

redshift_connector/plugin/azure_credentials_provider.py

@@ -3,7 +3,9 @@
 import typing
 
 from redshift_connector.error import InterfaceError
+from redshift_connector.plugin.azure_utils import validate_idp_partition
 from redshift_connector.plugin.credential_provider_constants import azure_headers
+from redshift_connector.plugin.plugin_utils import get_microsoft_idp_host
 from redshift_connector.plugin.saml_credentials_provider import SamlCredentialsProvider
 from redshift_connector.redshift_property import RedshiftProperty
 
@@ -23,6 +25,7 @@ def __init__(self: "AzureCredentialsProvider") -> None:
  self.idp_tenant: typing.Optional[str] = None
  self.client_secret: typing.Optional[str] = None
  self.client_id: typing.Optional[str] = None
+ self.idp_partition: typing.Optional[str] = None
 
  # method to grab the field parameters specified by end user.
  # This method adds to it Azure specific parameters.
@@ -34,6 +37,9 @@ def add_parameter(self: "AzureCredentialsProvider", info: RedshiftProperty) -> N
  self.client_secret = info.client_secret
  # The value of parameter client_id.
  self.client_id = info.client_id
+ 
+ # Validate and set idp_partition
+ self.idp_partition = validate_idp_partition(info.idp_partition)
 
  # Required method to grab the SAML Response. Used in base class to refresh temporary credentials.
  def get_saml_assertion(self: "AzureCredentialsProvider") -> str:
@@ -63,7 +69,10 @@ def azure_oauth_based_authentication(self: "AzureCredentialsProvider") -> str:
  import requests
 
  # endpoint to connect with Microsoft Azure to get SAML Assertion token
- url: str = "https://login.microsoftonline.com/{tenant}/oauth2/token".format(tenant=self.idp_tenant)
+ url: str = "https://{host}/{tenant}/oauth2/token".format(
+ host=get_microsoft_idp_host(self.idp_partition), 
+ tenant=self.idp_tenant
+ )
  _logger.debug("Uri: %s", url)
  self.validate_url(url)
 

redshift_connector/plugin/azure_utils.py

@@ -0,0 +1,16 @@
+"""Shared utilities for Azure credential providers."""
+from redshift_connector.error import InterfaceError
+import typing
+
+
+def validate_idp_partition(idp_partition: typing.Optional[str]) -> typing.Optional[str]:
+ """Validate idp_partition parameter and return normalized value."""
+ if idp_partition is not None:
+ if not isinstance(idp_partition, str):
+ raise InterfaceError("idp_partition must be a string")
+ # Validate against allowed values
+ valid_partitions = ["", "us-gov", "cn"]
+ normalized = idp_partition.strip().lower()
+ if normalized not in valid_partitions:
+ raise InterfaceError(f"idp_partition must be one of: {', '.join(repr(p) if p else 'empty string' for p in valid_partitions)}")
+ return idp_partition

redshift_connector/plugin/browser_azure_credentials_provider.py

@@ -6,7 +6,9 @@
 import typing
 
 from redshift_connector.error import InterfaceError
+from redshift_connector.plugin.azure_utils import validate_idp_partition
 from redshift_connector.plugin.credential_provider_constants import azure_headers
+from redshift_connector.plugin.plugin_utils import get_microsoft_idp_host
 from redshift_connector.plugin.saml_credentials_provider import SamlCredentialsProvider
 from redshift_connector.redshift_property import RedshiftProperty
 
@@ -28,6 +30,7 @@ def __init__(self: "BrowserAzureCredentialsProvider") -> None:
 
  self.idp_response_timeout: int = 120
  self.listen_port: int = 0
+ self.idp_partition: typing.Optional[str] = None
 
  self.redirectUri: typing.Optional[str] = None
 
@@ -48,6 +51,9 @@ def add_parameter(self: "BrowserAzureCredentialsProvider", info: RedshiftPropert
  self.idp_tenant = info.idp_tenant
  # The value of parameter client_id.
  self.client_id = info.client_id
+ 
+ # Validate and set idp_partition
+ self.idp_partition = validate_idp_partition(info.idp_partition)
 
  self.idp_response_timeout = info.idp_response_timeout
 
@@ -110,7 +116,10 @@ def fetch_saml_response(self: "BrowserAzureCredentialsProvider", token):
  _logger.debug("BrowserAzureCredentialsProvider.fetch_saml_response")
  import requests
 
- url: str = "https://login.microsoftonline.com/{tenant}/oauth2/token".format(tenant=self.idp_tenant)
+ url: str = "https://{host}/{tenant}/oauth2/token".format(
+ host=get_microsoft_idp_host(self.idp_partition), 
+ tenant=self.idp_tenant
+ )
  # headers to pass with POST request
  headers: typing.Dict[str, str] = azure_headers
  self.validate_url(url)
@@ -237,16 +246,21 @@ def run_server(
  def open_browser(self: "BrowserAzureCredentialsProvider", state: str) -> None:
  _logger.debug("BrowserAzureCredentialsProvider.open_browser")
  import webbrowser
-
- url: str = (
- "https://login.microsoftonline.com/{tenant}"
- "/oauth2/authorize"
- "?scope=openid"
- "&response_type=code"
- "&response_mode=form_post"
- "&client_id={id}"
- "&redirect_uri={uri}"
- "&state={state}".format(tenant=self.idp_tenant, id=self.client_id, uri=self.redirectUri, state=state)
+ from urllib.parse import quote, urlencode
+
+ # For query parameters, use urlencode for the entire query string
+ query_params = {
+ 'scope': 'openid',
+ 'response_type': 'code',
+ 'response_mode': 'form_post',
+ 'client_id': self.client_id,
+ 'redirect_uri': self.redirectUri,
+ 'state': state
+ }
+ url: str = "https://{host}/{tenant}/oauth2/authorize?{query}".format(
+ host=get_microsoft_idp_host(self.idp_partition), 
+ tenant=self.idp_tenant,
+ query=urlencode(query_params)
  )
  self.validate_url(url)
  webbrowser.open(url)

redshift_connector/plugin/browser_azure_oauth2_credentials_provider.py

@@ -4,8 +4,10 @@
 from enum import Enum
 
 from redshift_connector.error import InterfaceError
+from redshift_connector.plugin.azure_utils import validate_idp_partition
 from redshift_connector.plugin.credential_provider_constants import azure_headers
 from redshift_connector.plugin.jwt_credentials_provider import JwtCredentialsProvider
+from redshift_connector.plugin.plugin_utils import get_microsoft_idp_host
 from redshift_connector.redshift_property import RedshiftProperty
 
 if typing.TYPE_CHECKING:
@@ -35,15 +37,16 @@ class OAuthParamNames(Enum):
  SCOPE: str = "scope"
  RESPONSE_MODE: str = "response_mode"
  RESOURCE: str = "resource"
+ IDP_PARTITION: str = "idp_partition"
 
- MICROSOFT_IDP_HOST: str = "login.microsoftonline.com"
  CURRENT_INTERACTION_SCHEMA: str = "https"
 
  def __init__(self: "BrowserAzureOAuth2CredentialsProvider") -> None:
  super().__init__()
  self.idp_tenant: typing.Optional[str] = None
  self.client_id: typing.Optional[str] = None
  self.scope: str = ""
+ self.idp_partition: typing.Optional[str] = None
  self.idp_response_timeout: int = 120
  self.listen_port: int = 0
 
@@ -55,6 +58,9 @@ def add_parameter(
  self.idp_tenant = info.idp_tenant
  self.client_id = info.client_id
  self.scope = info.scope
+ 
+ # Validate and set idp_partition
+ self.idp_partition = validate_idp_partition(info.idp_partition)
 
  if info.idp_response_timeout:
  self.idp_response_timeout = info.idp_response_timeout
@@ -207,7 +213,7 @@ def get_authorization_token_url(self, state: str) -> str:
  return urlunsplit(
  (
  BrowserAzureOAuth2CredentialsProvider.CURRENT_INTERACTION_SCHEMA,
- BrowserAzureOAuth2CredentialsProvider.MICROSOFT_IDP_HOST,
+ get_microsoft_idp_host(self.idp_partition),
  "/{}/oauth2/v2.0/authorize".format(self.idp_tenant),
  encoded_params,
  "",
@@ -251,7 +257,7 @@ def get_jwt_post_request_url(self: "BrowserAzureOAuth2CredentialsProvider") -> s
  """
  return "{}://{}{}".format(
  BrowserAzureOAuth2CredentialsProvider.CURRENT_INTERACTION_SCHEMA,
- BrowserAzureOAuth2CredentialsProvider.MICROSOFT_IDP_HOST,
+ get_microsoft_idp_host(self.idp_partition),
  "/{}/oauth2/v2.0/token".format(self.idp_tenant),
  )
 

redshift_connector/plugin/plugin_utils.py

@@ -0,0 +1,27 @@
+from redshift_connector.error import InterfaceError
+import typing
+
+# Microsoft IdP host constants
+MICROSOFT_COMMERCIAL_HOST = "login.microsoftonline.com"
+MICROSOFT_IDP_HOSTS = {
+ "us-gov": "login.microsoftonline.us",
+ "cn": "login.chinacloudapi.cn"
+}
+
+
+def get_microsoft_idp_host(idp_partition: typing.Optional[str] = None) -> str:
+ """Returns the appropriate Microsoft IDP host based on the idp_partition value."""
+ from redshift_connector.plugin.azure_utils import validate_idp_partition
+ 
+ validated_partition = validate_idp_partition(idp_partition)
+ if not validated_partition or not validated_partition.strip():
+ return MICROSOFT_COMMERCIAL_HOST
+ 
+ partition = validated_partition.strip().lower()
+ if partition in MICROSOFT_IDP_HOSTS:
+ return MICROSOFT_IDP_HOSTS[partition]
+ else:
+ supported_values = list(MICROSOFT_IDP_HOSTS.keys())
+ raise InterfaceError("Invalid IdP partition: '{}' (normalized: '{}'). Supported values are: {}, or empty/None for commercial cloud.".format(
+ idp_partition, partition, ", ".join("'{}'".format(v) for v in supported_values)
+ ))

redshift_connector/redshift_property.py

@@ -67,6 +67,8 @@ def __init__(self: "RedshiftProperty", **kwargs):
  self.idp_host: typing.Optional[str] = None
  # timeout for authentication via Browser IDP
  self.idp_response_timeout: int = 120
+ # The IdP partition for multi-tenant IdP configurations (Azure AD: us-gov, cn)
+ self.idp_partition: typing.Optional[str] = None
  # The Azure AD tenant ID for your Redshift application.Only used for Azure AD.
  self.idp_tenant: typing.Optional[str] = None
  # The port used by an IdP (identity provider).

test/conftest.py

@@ -360,6 +360,25 @@ def redshift_browser_idc() -> typing.Dict[str, typing.Union[str, typing.Optional
  return db_connect
 
 
+@pytest.fixture(scope="class")
+def redshift_browser_azure_oauth2() -> typing.Dict[str, typing.Union[str, typing.Optional[bool], int]]:
+ db_connect = {
+ "host": conf.get("redshift-browser-azure-oauth2", "host", fallback=None),
+ "port": conf.getint("redshift-browser-azure-oauth2", "port", fallback=5439),
+ "database": conf.get("redshift-browser-azure-oauth2", "database", fallback="dev"),
+ "iam": conf.getboolean("redshift-browser-azure-oauth2", "iam", fallback=True),
+ "credentials_provider": conf.get(
+ "redshift-browser-azure-oauth2", "credentials_provider", fallback="BrowserAzureOAuth2CredentialsProvider"
+ ),
+ "idp_tenant": conf.get("redshift-browser-azure-oauth2", "idp_tenant", fallback=None),
+ "client_id": conf.get("redshift-browser-azure-oauth2", "client_id", fallback=None),
+ "scope": conf.get("redshift-browser-azure-oauth2", "scope", fallback=None),
+ "idp_partition": conf.get("redshift-browser-azure-oauth2", "idp_partition", fallback=None),
+ "idp_response_timeout": conf.getint("redshift-browser-azure-oauth2", "idp_response_timeout", fallback=120),
+ }
+ return db_connect
+
+
 @pytest.fixture
 def con(request, db_kwargs) -> redshift_connector.Connection:
  conn: redshift_connector.Connection = redshift_connector.connect(**db_kwargs)

test/unit/plugin/test_azure_common.py

@@ -0,0 +1,48 @@
+"""Common test utilities for Azure credential providers."""
+import pytest
+from redshift_connector import InterfaceError
+from redshift_connector.plugin.azure_utils import validate_idp_partition
+from redshift_connector.plugin.plugin_utils import get_microsoft_idp_host
+
+
+@pytest.mark.parametrize("partition", ["", " ", None])
+def test_get_microsoft_idp_host_empty_partition_returns_commercial_host(partition) -> None:
+ """Test that empty/None partitions return commercial host."""
+ assert get_microsoft_idp_host(partition) == "login.microsoftonline.com"
+
+
+@pytest.mark.parametrize("partition", ["us-gov", "US-GOV", "Us-gov "])
+def test_get_microsoft_idp_host_us_gov_partition(partition) -> None:
+ """Test that us-gov partitions return US government host."""
+ assert get_microsoft_idp_host(partition) == "login.microsoftonline.us"
+
+
+@pytest.mark.parametrize("partition", ["cn", "CN", "Cn "])
+def test_get_microsoft_idp_host_china_partition(partition) -> None:
+ """Test that cn partitions return China host."""
+ assert get_microsoft_idp_host(partition) == "login.chinacloudapi.cn"
+
+
+def test_get_microsoft_idp_host_invalid_partition_throws_error() -> None:
+ """Test that invalid partitions raise InterfaceError."""
+ with pytest.raises(InterfaceError, match="idp_partition must be one of"):
+ get_microsoft_idp_host("random_partition")
+
+
+@pytest.mark.parametrize("partition", ["", "us-gov", "cn", None])
+def test_validate_idp_partition_valid_values(partition) -> None:
+ """Test that valid partition values are accepted."""
+ result = validate_idp_partition(partition)
+ assert result == partition
+
+
+def test_validate_idp_partition_invalid_type() -> None:
+ """Test that non-string partition values raise InterfaceError."""
+ with pytest.raises(InterfaceError, match="idp_partition must be a string"):
+ validate_idp_partition(123)
+
+
+def test_validate_idp_partition_invalid_value() -> None:
+ """Test that invalid partition values raise InterfaceError."""
+ with pytest.raises(InterfaceError, match="idp_partition must be one of"):
+ validate_idp_partition("invalid_partition")

test/unit/plugin/test_azure_credentials_provider.py

@@ -7,6 +7,15 @@
 from redshift_connector import InterfaceError, RedshiftProperty
 from redshift_connector.plugin import AzureCredentialsProvider
 from redshift_connector.plugin.credential_provider_constants import azure_headers
+from redshift_connector.plugin.plugin_utils import get_microsoft_idp_host
+
+# Import common Azure tests
+from test.unit.plugin.test_azure_common import (
+ test_get_microsoft_idp_host_empty_partition_returns_commercial_host,
+ test_get_microsoft_idp_host_us_gov_partition,
+ test_get_microsoft_idp_host_china_partition,
+ test_get_microsoft_idp_host_invalid_partition_throws_error
+)
 
 
 def make_valid_azure_credentials_provider() -> typing.Tuple[AzureCredentialsProvider, RedshiftProperty]: