def run(args, gcloud_compute, email='', in_cloud_shell=False, **unused_kwargs):
"""Implementation of the `datalab connect` subcommand.
Args:
args: The Namespace instance returned by argparse
gcloud_compute: Function that can be used to invoke `gcloud compute`
email: The user's email address
in_cloud_shell: Whether or not the command is being run in the
Google Cloud Shell
Raises:
subprocess.CalledProcessError: If a nested `gcloud` calls fails
"""
instance = args.instance
status, metadata_items = utils.describe_instance(args, gcloud_compute,
instance)
for_user = metadata_items.get('for-user', '')
sdk_version = metadata_items.get('created-with-sdk-version', 'UNKNOWN')
datalab_version = metadata_items.get('created-with-datalab-version',
'UNKNOWN')
if (not args.no_user_checking) and for_user and (for_user != email):
print(wrong_user_message_template.format(for_user, email))
return
if args.diagnose_me:
print('Instance {} was created with the following '
'Cloud SDK component versions:'
'\n\tCloud SDK: {}'
'\n\tDatalab: {}'.format(instance, sdk_version, datalab_version))
maybe_start(args, gcloud_compute, instance, status)
connect(args, gcloud_compute, email, in_cloud_shell)
return
def run(args, gcloud_compute, email='', in_cloud_shell=False, **unused_kwargs):
"""Implementation of the `datalab connect` subcommand.
Args:
args: The Namespace instance returned by argparse
gcloud_compute: Function that can be used to invoke `gcloud compute`
email: The user's email address
in_cloud_shell: Whether or not the command is being run in the
Google Cloud Shell
Raises:
subprocess.CalledProcessError: If a nested `gcloud` calls fails
"""
instance = args.instance
status, metadata_items = utils.describe_instance(args, gcloud_compute,
instance)
for_user = metadata_items.get('for-user', '')
if (not args.no_user_checking) and for_user and (for_user != email):
print(wrong_user_message_template.format(for_user, email))
return
maybe_start(args, gcloud_compute, instance, status)
connect(args, gcloud_compute, email, in_cloud_shell)
return
def connect(args, gcloud_compute, email, in_cloud_shell):
"""Create a persistent connection to a Datalab instance.
Args:
args: The Namespace object constructed by argparse
gcloud_compute: A function that can be called to invoke `gcloud compute`
email: The user's email address
in_cloud_shell: Whether or not the command is being run in the
Google Cloud Shell
"""
instance = args.instance
connect_msg = ('Connecting to {0}.\n'
'This will create an SSH tunnel '
'and may prompt you to create an rsa key pair. '
'To manage these keys, see https://cloud.google.com/'
'compute/docs/instances/adding-removing-ssh-keys')
print(connect_msg.format(instance))
datalab_port = args.port
datalab_address = 'http://*****:*****@{0}'.format(instance))
gcloud_compute(args, cmd)
return
def maybe_open_browser(address):
"""Try to open a browser if we reasonably can."""
try:
browser_context = webbrowser.get()
browser_name = type(browser_context).__name__
if browser_name in unsupported_browsers:
return
webbrowser.open(datalab_address)
except webbrowser.Error as e:
print('Unable to open the webbrowser: ' + str(e))
def on_ready():
"""Callback that handles a successful connection."""
print('\nThe connection to Datalab is now open and will '
'remain until this command is killed.')
if in_cloud_shell:
print(web_preview_message_template.format(datalab_port))
else:
print('You can connect to Datalab at ' + datalab_address)
if not args.no_launch_browser:
maybe_open_browser(datalab_address)
return
def health_check(cancelled_event, healthy_event):
"""Check if the Datalab instance is reachable via the connection.
After the instance is reachable, the `on_ready` method is called.
This method is meant to be suitable for running in a separate thread,
and takes an event argument to indicate when that thread should exit.
Args:
cancelled_event: A threading.Event instance that indicates we should
give up on the instance becoming reachable.
healthy_event: A threading.Event instance that can be used to
indicate that the instance became reachable.
"""
health_url = '{0}_info/'.format(datalab_address)
healthy = False
print('Waiting for Datalab to be reachable at ' + datalab_address)
while not cancelled_event.is_set():
try:
health_resp = urllib2.urlopen(health_url)
if health_resp.getcode() == 200:
healthy = True
break
except Exception:
continue
if healthy:
healthy_event.set()
on_ready()
return
def connect_and_check(healthy_event):
"""Create a connection to Datalab and notify the user when ready.
This method blocks for as long as the connection is open.
Args:
healthy_event: A threading.Event instance that can be used to
indicate that the instance became reachable.
Returns:
True iff the Datalab instance became reachable.
Raises:
KeyboardInterrupt: If the user kills the connection.
"""
cancelled_event = threading.Event()
health_check_thread = threading.Thread(
target=health_check, args=[cancelled_event, healthy_event])
health_check_thread.start()
try:
create_tunnel()
except subprocess.CalledProcessError:
print('Connection broken')
finally:
cancelled_event.set()
health_check_thread.join()
return healthy_event.is_set()
remaining_reconnects = args.max_reconnects
while True:
healthy_event = threading.Event()
try:
connect_and_check(healthy_event)
except KeyboardInterrupt:
if healthy_event.is_set():
cli_flags = ' '
if args.project:
cli_flags += '--project {} '.format(args.project)
if args.zone:
cli_flags += '--zone {} '.format(args.zone)
cli_flags += '--port {} '.format(args.port)
print(
connection_closed_message_template.format(
instance, cli_flags))
return
if remaining_reconnects == 0:
return
# Before we try to reconnect, check to see if the VM is still running.
status, unused_metadata_items = utils.describe_instance(
args, gcloud_compute, instance)
if status != _STATUS_RUNNING:
print('Instance {0} is no longer running ({1})'.format(
instance, status))
return
print('Attempting to reconnect...')
remaining_reconnects -= 1
# Don't launch the browser on reconnect...
args.no_launch_browser = True
return
