From 9ac170435829e603b2b0727f48d90f36d2f0e85a Mon Sep 17 00:00:00 2001 From: Jonathan Ellis Date: Sat, 14 Sep 2024 15:43:11 -0500 Subject: [PATCH 1/6] add convenience method for authenticating with astra via db_id and token --- cassandra/cluster.py | 46 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index d5f80290a9..382e66c574 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -37,6 +37,9 @@ import time from threading import Lock, RLock, Thread, Event import uuid +import os +import urllib.request +import json import weakref from weakref import WeakValueDictionary @@ -1169,6 +1172,26 @@ def __init__(self, uses_twisted = TwistedConnection and issubclass(self.connection_class, TwistedConnection) 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: + # download SCB if necessary + if 'secure_connect_bundle' not in cloud: + bundle_path = f'astra-secure-connect-{cloud["db_id"]}.zip' + if not os.path.exists(bundle_path): + print('Downloading Secure Cloud Bundle') + url = self._get_astra_bundle_url(cloud['db_id'], cloud['token']) + try: + with urllib.request.urlopen(url) as r: + with open(bundle_path, 'wb') as f: + f.write(r.read()) + except urllib.error.URLError as e: + raise Exception(f"Error downloading secure connect bundle: {str(e)}") + cloud['secure_connect_bundle'] = bundle_path + # Set up auth_provider if not provided + if auth_provider is None: + auth_provider = PlainTextAuthProvider('token', cloud['token']) + cloud_config = dscloud.get_cloud_config(cloud, create_pyopenssl_context=uses_twisted or uses_eventlet) ssl_context = cloud_config.ssl_context @@ -2184,6 +2207,29 @@ def get_control_connection_host(self): endpoint = connection.endpoint if connection else None return self.metadata.get_host(endpoint) if endpoint else None + @staticmethod + def _get_astra_bundle_url(db_id, token): + # set up the request + url = f"https://api.astra.datastax.com/v2/databases/{db_id}/secureBundleURL" + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + req = urllib.request.Request(url, method="POST", 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'] + # handle errors + if 'errors' in response_data: + raise Exception(response_data['errors'][0]['message']) + raise Exception('Unknown error in ' + str(response_data)) + except urllib.error.URLError as e: + raise Exception(f"Error connecting to Astra API: {str(e)}") + def refresh_schema_metadata(self, max_schema_agreement_wait=None): """ Synchronously refresh all schema metadata. From 13aed0d4003212dc681244902d0d14ddba5573fb Mon Sep 17 00:00:00 2001 From: Madhavan Sridharan Date: Thu, 19 Sep 2024 16:54:23 -0400 Subject: [PATCH 2/6] 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 3/6] 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: From ee0b190cfcd3bd541403a3cd78ef3c56539650e9 Mon Sep 17 00:00:00 2001 From: Jonathan Ellis Date: Sat, 21 Sep 2024 09:02:27 -0500 Subject: [PATCH 4/6] fix KeyError when db_region is unspecified --- cassandra/cluster.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 2caa63d54d..c672897d26 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -1198,7 +1198,7 @@ def __init__(self, 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['token'], cloud["db_region"]) + url = self._get_astra_bundle_url(cloud['db_id'], cloud['token'], cloud.get("db_region")) try: with urllib.request.urlopen(url) as r: with open(bundle_path, 'wb') as f: From f9922bec16fc0975309248054339877a8a3d30e8 Mon Sep 17 00:00:00 2001 From: Jonathan Ellis Date: Sat, 21 Sep 2024 09:04:43 -0500 Subject: [PATCH 5/6] cleanup --- cassandra/cluster.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index c672897d26..8bb1149ac6 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -1046,7 +1046,7 @@ def default_retry_policy(self, policy): (or) { - # Astra DB cluster UUID. Only if secure_connect_bundle is not provided + # Astra DB cluster UUID, used if secure_connect_bundle is not provided 'db_id': 'db_id', # required with db_id. Astra DB token @@ -1188,7 +1188,7 @@ 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', '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"]}' @@ -2233,7 +2233,7 @@ def _get_astra_bundle_url(db_id, token, db_region: Optional[str] = None): Args: db_id (str): The Astra DB cluster UUID. - token (str): The Astra token. + token (str): The Astra security token. db_region (optional str): The Astra DB cluster region. Returns: From 6f5a4933847ba784a8cdb69cb29bdc5e65b7ebdc Mon Sep 17 00:00:00 2001 From: Jonathan Ellis Date: Sat, 21 Sep 2024 09:11:52 -0500 Subject: [PATCH 6/6] simplify to a single path for retrieving secureBundleUrl, and improve error handling when region not found --- cassandra/cluster.py | 40 ++++++++++++++++++---------------------- 1 file changed, 18 insertions(+), 22 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 8bb1149ac6..4fd99a0042 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -2240,9 +2240,9 @@ def _get_astra_bundle_url(db_id, token, db_region: Optional[str] = None): str: The secure connect bundle URL for the given inputs. Raises: - URLError: If the request to the Astra DB API fails. + ValueError: If the Astra DB API response is missing the download url or the specified db_region is not found. + Exception: If there's an error connecting to the Astra API or processing the response. """ - # set up the request url = f"https://api.astra.datastax.com/v2/databases/{db_id}/datacenters" headers = { "Authorization": f"Bearer {token}", @@ -2253,28 +2253,24 @@ def _get_astra_bundle_url(db_id, token, db_region: Optional[str] = None): try: with urllib.request.urlopen(req) as response: response_data = json.loads(response.read().decode()) - - 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: - raise ValueError("'secureBundleUrl' is missing from the Astra DB API response") + + # Convert list of responses to a dict keyed by region + datacenter_dict = {dc['region']: dc for dc in response_data if 'region' in dc} + + # Pull out the specified region, or the first one if not specified + if not datacenter_dict: + raise ValueError("No valid datacenter information found in the Astra DB API response") + if db_region: + if db_region not in datacenter_dict: + raise ValueError(f"Astra DB region '{db_region}' not found in list of regions") + datacenter = datacenter_dict[db_region] 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: - raise ValueError("'secureBundleUrl' is missing from the Astra DB API response for the primary region") + # Use the first datacenter as the primary region + datacenter = next(iter(datacenter_dict.values())) - # handle errors - if 'errors' in response_data: - raise Exception(response_data['errors'][0]['message']) - raise Exception('Unknown error in ' + str(response_data)) + if 'secureBundleUrl' not in datacenter or not datacenter['secureBundleUrl']: + raise ValueError("'secureBundleUrl' is missing or empty in the Astra DB API response") + return datacenter['secureBundleUrl'] except urllib.error.URLError as e: raise Exception(f"Error connecting to Astra API: {str(e)}")