99
1010from airbyte import cloud , get_destination , get_source
1111from airbyte ._util import api_util
12- from airbyte .cloud .auth import (
13- resolve_cloud_api_url ,
14- resolve_cloud_client_id ,
15- resolve_cloud_client_secret ,
16- )
1712from airbyte .cloud .connectors import CustomCloudSourceDefinition
1813from airbyte .cloud .constants import FAILED_STATUSES
1914from airbyte .cloud .workspaces import CloudWorkspace
@@ -1249,17 +1244,19 @@ def _resolve_organization(
12491244 organization_name : str | None ,
12501245 * ,
12511246 api_root : str ,
1252- client_id : SecretString ,
1253- client_secret : SecretString ,
1247+ client_id : SecretString | None ,
1248+ client_secret : SecretString | None ,
1249+ bearer_token : SecretString | None = None ,
12541250) -> api_util .models .OrganizationResponse :
12551251 """Resolve organization from either ID or exact name match.
12561252
12571253 Args:
12581254 organization_id: The organization ID (if provided directly)
12591255 organization_name: The organization name (exact match required)
12601256 api_root: The API root URL
1261- client_id: OAuth client ID
1262- client_secret: OAuth client secret
1257+ client_id: OAuth client ID (optional if bearer_token is provided)
1258+ client_secret: OAuth client secret (optional if bearer_token is provided)
1259+ bearer_token: Bearer token for authentication (optional if client credentials provided)
12631260
12641261 Returns:
12651262 The resolved OrganizationResponse object
@@ -1283,7 +1280,7 @@ def _resolve_organization(
12831280 api_root = api_root ,
12841281 client_id = client_id ,
12851282 client_secret = client_secret ,
1286- bearer_token = None , # Organization listing requires client credentials
1283+ bearer_token = bearer_token ,
12871284 )
12881285
12891286 if organization_id :
@@ -1327,8 +1324,9 @@ def _resolve_organization_id(
13271324 organization_name : str | None ,
13281325 * ,
13291326 api_root : str ,
1330- client_id : SecretString ,
1331- client_secret : SecretString ,
1327+ client_id : SecretString | None ,
1328+ client_secret : SecretString | None ,
1329+ bearer_token : SecretString | None = None ,
13321330) -> str :
13331331 """Resolve organization ID from either ID or exact name match.
13341332
@@ -1340,6 +1338,7 @@ def _resolve_organization_id(
13401338 api_root = api_root ,
13411339 client_id = client_id ,
13421340 client_secret = client_secret ,
1341+ bearer_token = bearer_token ,
13431342 )
13441343 return org .organization_id
13451344
@@ -1390,24 +1389,23 @@ def list_cloud_workspaces(
13901389 This tool will NOT list workspaces across all organizations - you must specify
13911390 which organization to list workspaces from.
13921391 """
1393- api_root = resolve_cloud_api_url ()
1394- client_id = resolve_cloud_client_id ()
1395- client_secret = resolve_cloud_client_secret ()
1392+ credentials = resolve_cloud_credentials ()
13961393
13971394 resolved_org_id = _resolve_organization_id (
13981395 organization_id = organization_id ,
13991396 organization_name = organization_name ,
1400- api_root = api_root ,
1401- client_id = client_id ,
1402- client_secret = client_secret ,
1397+ api_root = credentials .api_root ,
1398+ client_id = credentials .client_id ,
1399+ client_secret = credentials .client_secret ,
1400+ bearer_token = credentials .bearer_token ,
14031401 )
14041402
14051403 workspaces = api_util .list_workspaces_in_organization (
14061404 organization_id = resolved_org_id ,
1407- api_root = api_root ,
1408- client_id = client_id ,
1409- client_secret = client_secret ,
1410- bearer_token = None , # Workspace listing requires client credentials
1405+ api_root = credentials . api_root ,
1406+ client_id = credentials . client_id ,
1407+ client_secret = credentials . client_secret ,
1408+ bearer_token = credentials . bearer_token ,
14111409 name_contains = name_contains ,
14121410 max_items_limit = max_items_limit ,
14131411 )
@@ -1453,16 +1451,15 @@ def describe_cloud_organization(
14531451 Requires either organization_id OR organization_name (exact match) to be provided.
14541452 This tool is useful for looking up an organization's ID from its name, or vice versa.
14551453 """
1456- api_root = resolve_cloud_api_url ()
1457- client_id = resolve_cloud_client_id ()
1458- client_secret = resolve_cloud_client_secret ()
1454+ credentials = resolve_cloud_credentials ()
14591455
14601456 org = _resolve_organization (
14611457 organization_id = organization_id ,
14621458 organization_name = organization_name ,
1463- api_root = api_root ,
1464- client_id = client_id ,
1465- client_secret = client_secret ,
1459+ api_root = credentials .api_root ,
1460+ client_id = credentials .client_id ,
1461+ client_secret = credentials .client_secret ,
1462+ bearer_token = credentials .bearer_token ,
14661463 )
14671464
14681465 return CloudOrganizationResult (
0 commit comments