Skip to content

Commit

Permalink
Add method to get session and keyspace without setting them globally (#…
Browse files Browse the repository at this point in the history 
…168)

* Add method to get session and keyspace without setting them globally

* add 0.1.9 to changesfile

---------

Co-authored-by: Stefano Lottini <[email protected]>
  • Loading branch information
@cbornet @hemidactylus
cbornet and hemidactylus authored Oct 1, 2024
1 parent ba19858 commit 4455889
Show file tree
Hide file tree
Showing 3 changed files with 142 additions and 110 deletions.
9 changes: 9 additions & 0 deletions CHANGES.txt
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,12 @@
v 0.1.10
========
Add method to get session and keyspace from parameters.
Calling `cassio.init()` with insufficient arguments now raises an exception.

v 0.1.9
=======
introduce default proto version 4 for Astra (w/ fallback); alias ASTRA_DB_ID (#164)

v 0.1.8
=======
Add support for arbitrary CQL type in the 'body_blob' column (body_type init parameter)
Expand Down
236 changes: 131 additions & 105 deletions src/cassio/config/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
default_keyspace: Optional[str] = None


def init(
def get_session_and_keyspace(
auto: bool = False,
session: Optional[Session] = None,
secure_connect_bundle: Optional[str] = None,
Expand All @@ -38,88 +38,10 @@ def init(
tempfile_basedir: Optional[str] = None,
bundle_url_template: Optional[str] = None,
cloud_kwargs: Optional[Dict[str, Any]] = None,
) -> None:
"""
Globally set the default Cassandra connection (/keyspace) for CassIO.
This default will be used by all other db-requiring CassIO instantiations,
unless passed to the respective classes' __init__.
There are various ways to achieve this, depending on the passed parameters.
Broadly speaking, this method allows to pass one's own ready Session,
or to have it created in the method. For this second case, both Astra DB
and a regular Cassandra cluster can be targeted.
One can also simply call cassio.init(auto=True) and let it figure out
what to do based on standard environment variables.
CASSANDRA
If one passes `contact_points`, it is assumed that this is Cassandra.
In that case, only the following arguments will be used:
`contact_points`, `keyspace`, `username`, `password`, `cluster_kwargs`
Note that when passing a `session` all other parameters are ignored.
ASTRA DB:
If `contact_points` is not passed, one of several methods to connect to
Astra should be supplied for the connection to Astra. These overlap:
see below for their precedence.
Note that when passing a `session` all other parameters are ignored.
AUTO:
If passing auto=True, no other parameter is allowed except for
`tempfile_basedir`. The rest, including choosing Astra DB / Cassandra,
is autofilled by inspecting the following environment variables:
CASSANDRA_CONTACT_POINTS
CASSANDRA_USERNAME
CASSANDRA_PASSWORD
CASSANDRA_KEYSPACE
ASTRA_DB_APPLICATION_TOKEN
ASTRA_DB_INIT_STRING
ASTRA_DB_SECURE_BUNDLE_PATH
ASTRA_DB_KEYSPACE
ASTRA_DB_DATABASE_ID / ASTRA_DB_ID (equivalently)
PARAMETERS:
`auto`: (bool = False), whether to auto-guess all connection params.
`session` (optional cassandra.cluster.Session), an established connection.
`secure_connect_bundle` (optional str), full path to a Secure Bundle.
`init_string` (optional str), a stand-alone "db init string" credential
(which can optionally contain keyspace and/or token).
`token` (optional str), the Astra DB "AstraCS:..." token.
`database_id` (optional str), the Astra DB ID. Used only for Astra
connections with no `secure_connect_bundle` parameter passed.
`keyspace` (optional str), the keyspace to work.
`contact_points` (optional List[str]), for Cassandra connection
`username` (optional str), username for Cassandra connection
`password` (optional str), password for Cassandra connection
`cluster_kwargs` (optional dict), additional arguments to `Cluster(...)`.
`tempfile_basedir` (optional str), where to create temporary work directories.
`bundle_url_template` (optional str), url template for getting the database
secure bundle. The "databaseId" variable is resolved with the actual value.
Default (for Astra DB):
"https://api.astra.datastax.com/v2/databases/{database_id}/secureBundleURL"
`cloud_kwargs` (optional dict), additional arguments to `Cluster(cloud={...})`
(i.e. additional key-value pairs within the passed `cloud` dict).
ASTRA DB:
The Astra-related parameters are arranged in a chain of fallbacks.
In case redundant information is supplied, these are the precedences:
session > secure_connect_bundle > init_string
token > (from init_string if any)
keyspace > (from init_string if any) > (from bundle if any)
If a secure-connect-bundle is needed and not passed, it will be downloaded:
this requires `database_id` to be passed, suitable token permissions.
Constraints and caveats:
`secure_connect_bundle` requires `token`.
`session` does not make use of `cluster_kwargs` and will ignore it.
The Session is created at `init` time and kept around, available to any
subsequent table creation. If calling `init` a second time, a new Session
will be made available replacing the previous one.
"""
global default_session
global default_keyspace
) -> tuple[Session, Optional[str]]:
result_session: Optional[Session]
temp_dir_created: bool = False
temp_dir: Optional[str] = None
direct_session: Optional[Session] = None
bundle_from_is: Optional[str] = None
bundle_from_arg: Optional[str] = None
bundle_from_download: Optional[str] = None
Expand Down Expand Up @@ -185,9 +107,6 @@ def init(
)
keyspace_from_is = options_from_is.get("keyspace")
token_from_is = options_from_is.get("token")
# for the session
if session:
direct_session = session
if secure_connect_bundle:
if not token:
raise ValueError(
Expand All @@ -197,16 +116,9 @@ def init(
bundle_from_arg = secure_connect_bundle
token_from_arg = token
keyspace_from_arg = keyspace
can_be_astra = any(
[
secure_connect_bundle is not None,
init_string is not None,
token is not None,
]
)
# resolution of priority among args
if direct_session:
default_session = direct_session
if session:
result_session = session
else:
# first determine if Cassandra or Astra
is_cassandra = all(
Expand All @@ -217,12 +129,13 @@ def init(
contact_points is not None,
]
)
if is_cassandra:
is_astra_db = False
else:
# determine if Astra DB
is_astra_db = can_be_astra
#
can_be_astra = any(
[
secure_connect_bundle is not None,
init_string is not None,
token is not None,
]
)
if is_cassandra:
# need to take care of:
# contact_points, username, password, cluster_kwargs
Expand Down Expand Up @@ -263,8 +176,8 @@ def init(
auth_provider=chosen_auth_provider,
**(cluster_kwargs if cluster_kwargs is not None else {}),
)
default_session = cluster.connect()
elif is_astra_db:
result_session = cluster.connect()
elif can_be_astra:
# Astra DB
chosen_token = _first_valid(token_from_arg, token_from_is)
if chosen_token is None:
Expand Down Expand Up @@ -323,7 +236,7 @@ def init(
}
),
)
default_session = cluster.connect()
result_session = cluster.connect()
except NoHostAvailable:
# this is how a mismatched explicitly-set protocol version
# manifests itself. If the V4 guess fails, retry w/out version
Expand All @@ -338,19 +251,132 @@ def init(
),
**(cluster_kwargs if cluster_kwargs is not None else {}),
)
default_session = cluster.connect()
result_session = cluster.connect()
else:
raise ValueError("No secure-connect-bundle was available.")
else:
raise ValueError("Insufficient parameters to connect")
# keyspace to be resolved in any case
chosen_keyspace = _first_valid(
keyspace_from_arg, keyspace_from_is, keyspace_from_bundle
)
default_keyspace = chosen_keyspace
return result_session, chosen_keyspace
finally:
if temp_dir_created and temp_dir is not None:
shutil.rmtree(temp_dir)


def init(
auto: bool = False,
session: Optional[Session] = None,
secure_connect_bundle: Optional[str] = None,
init_string: Optional[str] = None,
token: Optional[str] = None,
database_id: Optional[str] = None,
keyspace: Optional[str] = None,
contact_points: Optional[Union[str, List[str]]] = None,
username: Optional[str] = None,
password: Optional[str] = None,
cluster_kwargs: Optional[Dict[str, Any]] = None,
tempfile_basedir: Optional[str] = None,
bundle_url_template: Optional[str] = None,
cloud_kwargs: Optional[Dict[str, Any]] = None,
) -> None:
"""
Globally set the default Cassandra connection (/keyspace) for CassIO.
This default will be used by all other db-requiring CassIO instantiations,
unless passed to the respective classes' __init__.
There are various ways to achieve this, depending on the passed parameters.
Broadly speaking, this method allows to pass one's own ready Session,
or to have it created in the method. For this second case, both Astra DB
and a regular Cassandra cluster can be targeted.
One can also simply call cassio.init(auto=True) and let it figure out
what to do based on standard environment variables.
CASSANDRA
If one passes `contact_points`, it is assumed that this is Cassandra.
In that case, only the following arguments will be used:
`contact_points`, `keyspace`, `username`, `password`, `cluster_kwargs`
Note that when passing a `session` all other parameters are ignored.
ASTRA DB:
If `contact_points` is not passed, one of several methods to connect to
Astra should be supplied for the connection to Astra. These overlap:
see below for their precedence.
Note that when passing a `session` all other parameters are ignored.
AUTO:
If passing auto=True, no other parameter is allowed except for
`tempfile_basedir`. The rest, including choosing Astra DB / Cassandra,
is autofilled by inspecting the following environment variables:
CASSANDRA_CONTACT_POINTS
CASSANDRA_USERNAME
CASSANDRA_PASSWORD
CASSANDRA_KEYSPACE
ASTRA_DB_APPLICATION_TOKEN
ASTRA_DB_INIT_STRING
ASTRA_DB_SECURE_BUNDLE_PATH
ASTRA_DB_KEYSPACE
ASTRA_DB_DATABASE_ID / ASTRA_DB_ID (equivalently)
PARAMETERS:
`auto`: (bool = False), whether to auto-guess all connection params.
`session` (optional cassandra.cluster.Session), an established connection.
`secure_connect_bundle` (optional str), full path to a Secure Bundle.
`init_string` (optional str), a stand-alone "db init string" credential
(which can optionally contain keyspace and/or token).
`token` (optional str), the Astra DB "AstraCS:..." token.
`database_id` (optional str), the Astra DB ID. Used only for Astra
connections with no `secure_connect_bundle` parameter passed.
`keyspace` (optional str), the keyspace to work.
`contact_points` (optional List[str]), for Cassandra connection
`username` (optional str), username for Cassandra connection
`password` (optional str), password for Cassandra connection
`cluster_kwargs` (optional dict), additional arguments to `Cluster(...)`.
`tempfile_basedir` (optional str), where to create temporary work directories.
`bundle_url_template` (optional str), url template for getting the database
secure bundle. The "databaseId" variable is resolved with the actual value.
Default (for Astra DB):
"https://api.astra.datastax.com/v2/databases/{database_id}/secureBundleURL"
`cloud_kwargs` (optional dict), additional arguments to `Cluster(cloud={...})`
(i.e. additional key-value pairs within the passed `cloud` dict).
ASTRA DB:
The Astra-related parameters are arranged in a chain of fallbacks.
In case redundant information is supplied, these are the precedences:
session > secure_connect_bundle > init_string
token > (from init_string if any)
keyspace > (from init_string if any) > (from bundle if any)
If a secure-connect-bundle is needed and not passed, it will be downloaded:
this requires `database_id` to be passed, suitable token permissions.
Constraints and caveats:
`secure_connect_bundle` requires `token`.
`session` does not make use of `cluster_kwargs` and will ignore it.
The Session is created at `init` time and kept around, available to any
subsequent table creation. If calling `init` a second time, a new Session
will be made available replacing the previous one.
"""
global default_session, default_keyspace
default_session, default_keyspace = get_session_and_keyspace(
auto=auto,
session=session,
secure_connect_bundle=secure_connect_bundle,
init_string=init_string,
token=token,
database_id=database_id,
keyspace=keyspace,
contact_points=contact_points,
username=username,
password=password,
cluster_kwargs=cluster_kwargs,
tempfile_basedir=tempfile_basedir,
bundle_url_template=bundle_url_template,
cloud_kwargs=cloud_kwargs,
)


def resolve_session(arg_session: Optional[Session] = None) -> Optional[Session]:
"""Utility to fall back to the global defaults when null args are passed."""
if arg_session is not None:
Expand Down Expand Up @@ -383,7 +409,7 @@ def check_resolve_keyspace(arg_keyspace: Optional[str] = None) -> str:
return s


def _first_valid(*pargs: Any) -> Union[Any, None]:
def _first_valid(*pargs: Optional[str]) -> Optional[str]:
for entry in pargs:
if entry is not None:
return entry
Expand Down
7 changes: 2 additions & 5 deletions tests/integration/test_init_astra.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,11 +66,8 @@ def test_init_noop(self) -> None:
assert resolve_keyspace() is None
assert resolve_session("s") == "s"
assert resolve_keyspace("k") == "k"
cassio.init()
assert resolve_session() is None
assert resolve_keyspace() is None
assert resolve_session("s") == "s"
assert resolve_keyspace("k") == "k"
with pytest.raises(ValueError, match="Insufficient parameters to connect"):
cassio.init()

def test_init_session(self) -> None:
_reset_cassio_globals()
Expand Down

0 comments on commit 4455889

Please sign in to comment.