Expand source code
def main() -> None:
args = parser.parse_args()
# Manual validation and integration with env. variables.
# We uniformly apply (1) parameters over env vars, (2) astra+local,
# and then check for conflicts and re-parse endpoint to find out the target env.
_endpoint = (
args.endpoint
or os.getenv("ASTRA_DB_API_ENDPOINT")
or os.getenv("LOCAL_DATA_API_ENDPOINT")
)
_keyspace = (
args.keyspace
or os.getenv("ASTRA_DB_KEYSPACE")
or os.getenv("LOCAL_DATA_API_KEYSPACE")
or None
)
_environment = args.environment or None
#
_token_provider: TokenProvider
_token = args.token or os.getenv("ASTRA_DB_APPLICATION_TOKEN")
_username = args.username or os.getenv("LOCAL_DATA_API_USERNAME")
_password = args.password or os.getenv("LOCAL_DATA_API_PASSWORD")
if _token and (_username or _password):
parser.error("Specify either a token or username/password pair, not both.")
if (_username and not _password) or (_password and not _username):
parser.error("Username and password must must be provided together.")
if _username and _password:
_token_provider = UsernamePasswordTokenProvider(_username, _password)
else:
_token_provider = StaticTokenProvider(_token)
if not _endpoint:
raise SystemExit(
"The endpoint is required, either via env. var or command-line argument. "
"Hint: '-h' for help."
)
# instantiate client and database
if not _environment:
parsed_endpoint = parse_api_endpoint(_endpoint)
if parsed_endpoint:
_environment = parsed_endpoint.environment
else:
_environment = Environment.HCD
# inject the token at client-level for repl users to easily spawn admins and such
client = DataAPIClient(
environment=_environment,
token=_token_provider,
)
database = client.get_database( # noqa: F841
_endpoint,
keyspace=_keyspace,
)
# pick the target REPL mode
ipython_repl: bool
if args.repl == "ipython":
if ipython_available:
ipython_repl = True
else:
raise SystemExit(
"'IPython' package not found: cannot start the required REPL. "
"Hint: '-h' for help."
)
elif args.repl == "stdlib":
ipython_repl = False
else:
# unspecified repl choice, pick best default
ipython_repl = ipython_available
banner = BANNER_TEMPLATE.format(
endpoint=_endpoint,
keyspace=_keyspace or "(default)",
environment=_environment,
)
# spawn the repl
namespace: dict[str, Any] = {
**{k: v for k, v in globals().items() if k not in PRIVATE_GLOBALS},
**{
"client": client,
"database": database,
},
}
if args.loglevel:
logging.basicConfig(level=LOGGING_LEVELS[args.loglevel])
if ipython_repl:
embed(
banner1=banner,
header=" ",
user_ns=namespace,
colors="Neutral",
) # type: ignore[no-untyped-call]
else:
stdbanner = banner + ("\n" * 3)
if not sys.flags.interactive:
interact(
banner=stdbanner,
local=namespace,
exitmsg="Closing astrapy-repl.",
)