From 13aed0d4003212dc681244902d0d14ddba5573fb Mon Sep 17 00:00:00 2001 From: Madhavan Sridharan Date: Thu, 19 Sep 2024 16:54:23 -0400 Subject: [PATCH 1/2] Improvized Astra DB connection based on the given region --- cassandra/cluster.py | 55 ++++++++++++++++++++++++++++++++++++-------- 1 file changed, 45 insertions(+), 10 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 382e66c574..1203be4d44 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -1042,6 +1042,19 @@ def default_retry_policy(self, policy): 'use_default_tempdir': True # use the system temp dir for the zip extraction } + (or) + + { + # Astra DB cluster UUID. Only if secure_connect_bundle is not provided + 'db_id': 'db_id', + + # required with db_id. Astra DB region + 'db_region': 'us-east1', + + # required with db_id. Astra DB token + 'token': 'AstraCS:change_me:change_me' + } + The zip file will be temporarily extracted in the same directory to load the configuration and certificates. """ @@ -1174,13 +1187,13 @@ def __init__(self, uses_eventlet = EventletConnection and issubclass(self.connection_class, EventletConnection) # Check if we need to download the secure connect bundle - if 'db_id' in cloud and 'token' in cloud: + if all(akey in cloud for akey in ('db_id', 'db_region', 'token')): # download SCB if necessary if 'secure_connect_bundle' not in cloud: - bundle_path = f'astra-secure-connect-{cloud["db_id"]}.zip' + bundle_path = f'astradb-scb-{cloud["db_id"]}-{cloud["db_region"]}.zip' if not os.path.exists(bundle_path): - print('Downloading Secure Cloud Bundle') - url = self._get_astra_bundle_url(cloud['db_id'], cloud['token']) + log.info('Downloading Secure Cloud Bundle...') + url = self._get_astra_bundle_url(cloud['db_id'], cloud["db_region"], cloud['token']) try: with urllib.request.urlopen(url) as r: with open(bundle_path, 'wb') as f: @@ -2208,21 +2221,43 @@ def get_control_connection_host(self): return self.metadata.get_host(endpoint) if endpoint else None @staticmethod - def _get_astra_bundle_url(db_id, token): + def _get_astra_bundle_url(db_id, db_region, token): + """ + Retrieves the secure connect bundle URL for an Astra DB cluster based on the provided 'db_id', + 'db_region' and 'token'. + + Args: + db_id (str): The Astra DB cluster UUID. + db_region (str): The Astra DB cluster region. + token (str): The Astra token. + + Returns: + str: The secure connect bundle URL for the given inputs. + + Raises: + URLError: If the request to the Astra DB API fails. + """ # set up the request - url = f"https://api.astra.datastax.com/v2/databases/{db_id}/secureBundleURL" + url = f"https://api.astra.datastax.com/v2/databases/{db_id}/datacenters" headers = { "Authorization": f"Bearer {token}", "Content-Type": "application/json" } - req = urllib.request.Request(url, method="POST", headers=headers, data=b"") + req = urllib.request.Request(url, method="GET", headers=headers, data=b"") try: with urllib.request.urlopen(req) as response: response_data = json.loads(response.read().decode()) - # happy path - if 'downloadURL' in response_data: - return response_data['downloadURL'] + for datacenter in response_data: + if 'secureBundleUrl' in datacenter: + # happy path + if datacenter['region'] == db_region: + return datacenter['secureBundleUrl'] + else: + log.warning("Astra DB cluster region [%s] does not match input [%s]", datacenter['region'], db_region) + else: + log.warning("'secureBundleUrl' is missing from the Astra DB API response") + # handle errors if 'errors' in response_data: raise Exception(response_data['errors'][0]['message']) From f764adeac6a759bce182ae7aebf647cd578f4af0 Mon Sep 17 00:00:00 2001 From: Madhavan Sridharan Date: Thu, 19 Sep 2024 19:37:57 -0400 Subject: [PATCH 2/2] Make providing region as optional --- cassandra/cluster.py | 45 ++++++++++++++++++++++++++++---------------- 1 file changed, 29 insertions(+), 16 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 1203be4d44..2caa63d54d 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -40,6 +40,7 @@ import os import urllib.request import json +from typing import Optional import weakref from weakref import WeakValueDictionary @@ -1048,11 +1049,11 @@ def default_retry_policy(self, policy): # Astra DB cluster UUID. Only if secure_connect_bundle is not provided 'db_id': 'db_id', - # required with db_id. Astra DB region - 'db_region': 'us-east1', - # required with db_id. Astra DB token 'token': 'AstraCS:change_me:change_me' + + # optional with db_id & token. Astra DB region + 'db_region': 'us-east1', } The zip file will be temporarily extracted in the same directory to @@ -1187,13 +1188,17 @@ def __init__(self, uses_eventlet = EventletConnection and issubclass(self.connection_class, EventletConnection) # Check if we need to download the secure connect bundle - if all(akey in cloud for akey in ('db_id', 'db_region', 'token')): + if all(akey in cloud for akey in ('db_id', 'token')): # download SCB if necessary if 'secure_connect_bundle' not in cloud: - bundle_path = f'astradb-scb-{cloud["db_id"]}-{cloud["db_region"]}.zip' + bundle_path = f'astradb-scb-{cloud["db_id"]}' + if 'db_region' in cloud: + bundle_path += f'-{cloud["db_region"]}.zip' + else: + bundle_path += '.zip' if not os.path.exists(bundle_path): log.info('Downloading Secure Cloud Bundle...') - url = self._get_astra_bundle_url(cloud['db_id'], cloud["db_region"], cloud['token']) + url = self._get_astra_bundle_url(cloud['db_id'], cloud['token'], cloud["db_region"]) try: with urllib.request.urlopen(url) as r: with open(bundle_path, 'wb') as f: @@ -2221,15 +2226,15 @@ def get_control_connection_host(self): return self.metadata.get_host(endpoint) if endpoint else None @staticmethod - def _get_astra_bundle_url(db_id, db_region, token): + def _get_astra_bundle_url(db_id, token, db_region: Optional[str] = None): """ Retrieves the secure connect bundle URL for an Astra DB cluster based on the provided 'db_id', - 'db_region' and 'token'. + 'db_region' (optional) and 'token'. Args: db_id (str): The Astra DB cluster UUID. - db_region (str): The Astra DB cluster region. token (str): The Astra token. + db_region (optional str): The Astra DB cluster region. Returns: str: The secure connect bundle URL for the given inputs. @@ -2248,15 +2253,23 @@ def _get_astra_bundle_url(db_id, db_region, token): try: with urllib.request.urlopen(req) as response: response_data = json.loads(response.read().decode()) - for datacenter in response_data: - if 'secureBundleUrl' in datacenter: - # happy path - if datacenter['region'] == db_region: - return datacenter['secureBundleUrl'] + + if db_region is not None and len(db_region) > 0: + for datacenter in response_data: + if 'secureBundleUrl' in datacenter and datacenter['secureBundleUrl']: + # happy path + if db_region == datacenter['region']: + return datacenter['secureBundleUrl'] + else: + log.warning("Astra DB cluster region [%s] does not match input [%s]", datacenter['region'], db_region) else: - log.warning("Astra DB cluster region [%s] does not match input [%s]", datacenter['region'], db_region) + raise ValueError("'secureBundleUrl' is missing from the Astra DB API response") + else: + # Return just the primary region SCB URL + if 'secureBundleUrl' in response_data[0] and response_data[0]['secureBundleUrl']: + return response_data[0]['secureBundleUrl'] else: - log.warning("'secureBundleUrl' is missing from the Astra DB API response") + raise ValueError("'secureBundleUrl' is missing from the Astra DB API response for the primary region") # handle errors if 'errors' in response_data: