def test_do_db_reset_initiate_without_iam_credentials(self):
with self.assertRaises(HTTPError):
initiate_database_reset(self.host, self.port, self.ssl)
def test_do_db_reset_initiate_with_iam_credentials(self):
request_generator = create_request_generator(
AuthModeEnum.IAM, IAMAuthCredentialsProvider.ENV)
result = initiate_database_reset(self.host, self.port, self.ssl,
request_generator)
self.assertNotEqual(result['payload']['token'], '')
def db_reset(self, line):
host = self.graph_notebook_config.host
port = self.graph_notebook_config.port
ssl = self.graph_notebook_config.ssl
logger.info(f'calling system endpoint {host}')
parser = argparse.ArgumentParser()
parser.add_argument('-g',
'--generate-token',
action='store_true',
help='generate token for database reset')
parser.add_argument('-t',
'--token',
nargs=1,
default='',
help='perform database reset with given token')
parser.add_argument('-y',
'--yes',
action='store_true',
help='skip the prompt and perform database reset')
args = parser.parse_args(line.split())
generate_token = args.generate_token
skip_prompt = args.yes
request_generator = create_request_generator(
self.graph_notebook_config.auth_mode,
self.graph_notebook_config.iam_credentials_provider_type)
logger.info(
f'used credentials_provider_mode={self.graph_notebook_config.iam_credentials_provider_type.name} and auth_mode={self.graph_notebook_config.auth_mode.name} to make system request'
)
if generate_token is False and args.token == '':
if skip_prompt:
res = initiate_database_reset(host, port, ssl,
request_generator)
token = res['payload']['token']
res = perform_database_reset(token, host, port, ssl,
request_generator)
logger.info(f'got the response {res}')
return res
output = widgets.Output()
source = 'Are you sure you want to delete all the data in your cluster?'
label = widgets.Label(source)
text_hbox = widgets.HBox([label])
check_box = widgets.Checkbox(
value=False,
disabled=False,
indent=False,
description=
'I acknowledge that upon deletion the cluster data will no longer be available.',
layout=widgets.Layout(width='600px', margin='5px 5px 5px 5px'))
button_delete = widgets.Button(description="Delete")
button_cancel = widgets.Button(description="Cancel")
button_hbox = widgets.HBox([button_delete, button_cancel])
display(text_hbox, check_box, button_hbox, output)
def on_button_delete_clicked(b):
result = initiate_database_reset(host, port, ssl,
request_generator)
text_hbox.close()
check_box.close()
button_delete.close()
button_cancel.close()
button_hbox.close()
if not check_box.value:
with output:
print('Checkbox is not checked.')
return
token = result['payload']['token']
if token == "":
with output:
print('Failed to get token.')
print(result)
return
result = perform_database_reset(token, host, port, ssl,
request_generator)
if 'status' not in result or result['status'] != '200 OK':
with output:
print(
'Database reset failed, please try the operation again or reboot the cluster.'
)
print(result)
logger.error(result)
return
retry = 10
poll_interval = 5
interval_output = widgets.Output()
job_status_output = widgets.Output()
status_hbox = widgets.HBox([interval_output])
vbox = widgets.VBox([status_hbox, job_status_output])
display(vbox)
last_poll_time = time.time()
while retry > 0:
time_elapsed = int(time.time() - last_poll_time)
time_remaining = poll_interval - time_elapsed
interval_output.clear_output()
if time_elapsed > poll_interval:
with interval_output:
print('checking status...')
job_status_output.clear_output()
with job_status_output:
display_html(HTML(loading_wheel_html))
try:
retry -= 1
interval_check_response = get_status(
host, port, ssl, request_generator)
except Exception as e:
# Exception is expected when database is resetting, continue waiting
with job_status_output:
last_poll_time = time.time()
time.sleep(1)
continue
job_status_output.clear_output()
with job_status_output:
if interval_check_response["status"] == 'healthy':
interval_output.close()
print('Database has been reset.')
return
last_poll_time = time.time()
else:
with interval_output:
print(
f'checking status in {time_remaining} seconds')
time.sleep(1)
with output:
print(result)
if interval_check_response["status"] != 'healthy':
print(
"Could not retrieve the status of the reset operation within the allotted time. "
"If the database is not healthy after 1 min, please try the operation again or "
"reboot the cluster.")
def on_button_cancel_clicked(b):
text_hbox.close()
check_box.close()
button_delete.close()
button_cancel.close()
button_hbox.close()
with output:
print('Database reset operation has been canceled.')
button_delete.on_click(on_button_delete_clicked)
button_cancel.on_click(on_button_cancel_clicked)
return
elif generate_token:
res = initiate_database_reset(host, port, ssl, request_generator)
else:
# args.token is an array of a single string, e.g., args.token=['ade-23-c23'], use index 0 to take the string
res = perform_database_reset(args.token[0], host, port, ssl,
request_generator)
logger.info(f'got the response {res}')
return res
def on_button_delete_clicked(b):
result = initiate_database_reset(host, port, ssl,
request_generator)
text_hbox.close()
check_box.close()
button_delete.close()
button_cancel.close()
button_hbox.close()
if not check_box.value:
with output:
print('Checkbox is not checked.')
return
token = result['payload']['token']
if token == "":
with output:
print('Failed to get token.')
print(result)
return
result = perform_database_reset(token, host, port, ssl,
request_generator)
if 'status' not in result or result['status'] != '200 OK':
with output:
print(
'Database reset failed, please try the operation again or reboot the cluster.'
)
print(result)
logger.error(result)
return
retry = 10
poll_interval = 5
interval_output = widgets.Output()
job_status_output = widgets.Output()
status_hbox = widgets.HBox([interval_output])
vbox = widgets.VBox([status_hbox, job_status_output])
display(vbox)
last_poll_time = time.time()
while retry > 0:
time_elapsed = int(time.time() - last_poll_time)
time_remaining = poll_interval - time_elapsed
interval_output.clear_output()
if time_elapsed > poll_interval:
with interval_output:
print('checking status...')
job_status_output.clear_output()
with job_status_output:
display_html(HTML(loading_wheel_html))
try:
retry -= 1
interval_check_response = get_status(
host, port, ssl, request_generator)
except Exception as e:
# Exception is expected when database is resetting, continue waiting
with job_status_output:
last_poll_time = time.time()
time.sleep(1)
continue
job_status_output.clear_output()
with job_status_output:
if interval_check_response["status"] == 'healthy':
interval_output.close()
print('Database has been reset.')
return
last_poll_time = time.time()
else:
with interval_output:
print(
f'checking status in {time_remaining} seconds')
time.sleep(1)
with output:
print(result)
if interval_check_response["status"] != 'healthy':
print(
"Could not retrieve the status of the reset operation within the allotted time. "
"If the database is not healthy after 1 min, please try the operation again or "
"reboot the cluster.")
def test_do_database_reset_initiate(self):
result = initiate_database_reset(self.host, self.port, self.ssl)
self.assertNotEqual(result['payload']['token'], '')