Single user notebooks (#62)

* Add admin option to set personal authentication for all clusters spawned from a single Hub
* Option also available at the property level for users if custom clusters are allowed
* Option also available at the template level (YAML) and adapt the key to match to the Hub user identity

Author: m-mayran

Diff for: dataprocspawner/spawner.py

@@ -333,6 +333,13 @@ class DataprocSpawner(Spawner):
       Whether to show spawned single-user clusters in the list of Dataproc Notebooks.
       """)
 
+  force_single_user = Bool(
+    False,
+    config=True,
+    help="""Whether a notebook on a cluster can only be accessed by the user who
+    spawned it.
+    """)
+
   def __init__(self, *args, **kwargs):
     mock = kwargs.pop('_mock', False)
     super().__init__(*args, **kwargs)
@@ -950,7 +957,8 @@ async def get_cluster(self, clustername):
 ################################################################################
 
   def get_username(self, raw=False):
-    return self.user.name if raw else re.sub(r'[^a-zA-Z0-9=]', '-', str(self.user.name))
+    username, _ = (self.user.name.split('@') + [None])[:2]
+    return username if raw else re.sub(r'[^a-zA-Z0-9=]', '-', str(username))
 
   def clustername(self, cluster_name=None):
     """ JupyterHub provides a notebook per user, so the username is used to
@@ -1301,7 +1309,7 @@ def _apply_users_configs(self, cluster_data):
         int(self.user_options.get('sec_worker_node_amount'))
 
     autoscaling_policy = self.user_options.get('autoscaling_policy', '')
-    if autoscaling_policy != 'None':
+    if autoscaling_policy and autoscaling_policy != 'None':
       cluster_data['config']['autoscaling_config'] = {
           'policy_uri': (
               f"""https://www.googleapis.com/compute/v1/projects/"""
@@ -1343,6 +1351,7 @@ def _build_cluster_config(self, cluster_data=None):
     # but must be set in case there is no form.
     cluster_data = cluster_data or {}
     cluster_zone = self.zone
+    personal_auth_property = 'dataproc:dataproc.personal-auth.user'
 
     # Sets the cluster definition with form data.
     if self.user_options:
@@ -1463,6 +1472,15 @@ def _build_cluster_config(self, cluster_data=None):
       else:
         cluster_data['config']['software_config']['image_version'] = '1.4-debian9'
 
+    # A user can only set 'dataproc:dataproc.personal-auth.user' for themselves.
+    # Otherwise the CG Url is not accessible by either identities. So, if set,
+    # the personal auth property can only have the value of self.user.name. For
+    # that reason, there is no priority between yaml and user.
+    if (self.force_single_user or
+        personal_auth_property in cluster_data['config']['software_config']['properties']):
+      (cluster_data['config']['software_config']['properties']
+                   [personal_auth_property]) = self.user.name
+
     # Forces Component Gateway
     cluster_data['config'].setdefault('endpoint_config', {})
     cluster_data['config']['endpoint_config']['enable_http_port_access'] = True

Diff for: docker/jupyterhub.sh

@@ -104,6 +104,7 @@ function set-environment-from-metadata {
   local machine_types_list
   local notebooks_examples_location
   local hub_allow_named_servers
+  local force_single_user
 
   notebooks_location=$( curl -s -f \
     ${metadata_base_url}/instance/attributes/notebooks-location -H "Metadata-Flavor: Google" || echo )
@@ -123,6 +124,8 @@ function set-environment-from-metadata {
     ${metadata_base_url}/instance/attributes/notebooks-examples-location -H "Metadata-Flavor: Google" || echo )
   hub_allow_named_servers=$( curl -s -f \
     ${metadata_base_url}/instance/attributes/hub-allow-named-servers -H "Metadata-Flavor: Google" || echo )
+  force_single_user=$( curl -s -f \
+    ${metadata_base_url}/instance/attributes/force-single-user -H "Metadata-Flavor: Google" || echo )
 
   if [ -n "${notebooks_location}" ]; then
     export NOTEBOOKS_LOCATION="${notebooks_location}"
@@ -151,6 +154,10 @@ function set-environment-from-metadata {
   if [ -n "${hub_allow_named_servers}" ]; then
     export HUB_ALLOW_NAMED_SERVERS="${hub_allow_named_servers}"
   fi
+  if [ -n "${force_single_user}" ]; then
+    export FORCE_SINGLE_USER="${force_single_user}"
+  fi
+
 }
 
 # 'set-region-and-zone-from-metadata'

Diff for: docker/jupyterhub_config.py

@@ -67,6 +67,7 @@ def is_true(boolstring: str):
 c.DataprocSpawner.allow_custom_clusters = is_true(os.environ.get('DATAPROC_ALLOW_CUSTOM_CLUSTERS', ''))
 c.DataprocSpawner.allow_random_cluster_names = is_true(os.environ.get('ALLOW_RANDOM_CLUSTER_NAMES', ''))
 c.DataprocSpawner.show_spawned_clusters_in_notebooks_list = is_true(os.environ.get('SHOW_SPAWNED_CLUSTERS', ''))
+c.DataprocSpawner.force_single_user = is_true(os.environ.get('FORCE_SINGLE_USER', ''))
 c.DataprocSpawner.gcs_notebooks = os.environ.get('GCS_NOTEBOOKS', '')
 if not c.DataprocSpawner.gcs_notebooks:
   c.DataprocSpawner.gcs_notebooks = os.environ.get('NOTEBOOKS_LOCATION', '')

Diff for: examples/deploy_local.sh

@@ -40,4 +40,6 @@ docker run -it \
 -e PROJECT="${PROJECT}" \
 -e DATAPROC_CONFIGS="${CONFIGS_LOCATION}" \
 -e JUPYTERHUB_REGION="us-west1" \
+-e DATAPROC_ALLOW_CUSTOM_CLUSTERS=true \
+-e FORCE_SINGLE_USER=true \
 "${DOCKER_IMAGE}"

Diff for: tests/test_data/perso.yaml

@@ -0,0 +1,5 @@
+clusterName: 'personal'
+config:
+  softwareConfig:
+    properties:
+      dataproc:dataproc.personal-auth.user: '[email protected]'

Diff for: tests/test_spawner.py

@@ -61,7 +61,7 @@ def set_delay_done(self):
     self.op_done = True
 
 class MockUser(mock.Mock):
-  name = 'fake'
+  name = 'fake@example.com'
   base_url = '/user/fake'
   server = Server()
 
@@ -112,7 +112,7 @@ async def test_get_cluster_notebook_endpoint(*args, **kwargs):
 
     mock_client.create_cluster.assert_called_once()
 
-    assert spawner.cluster_definition['cluster_name'] == 'dataprochub-fake'
+    assert spawner.cluster_definition['cluster_name'] == f'dataprochub-{spawner.get_username()}'
     assert (spawner.cluster_definition['config']['gce_cluster_config']['zone_uri']) == (
         f'https://www.googleapis.com/compute/v1/projects/{spawner.project}/zones/{spawner.zone}')
 
@@ -150,7 +150,7 @@ async def test_get_cluster_notebook_endpoint(*args, **kwargs):
     url = await spawner.start()
     mock_client.create_cluster.assert_called_once()
 
-    assert spawner.cluster_definition['cluster_name'] == 'dataprochub-fake-server1'
+    assert spawner.cluster_definition['cluster_name'] == f'dataprochub-{spawner.get_username()}-server1'
 
   @pytest.mark.asyncio
   async def test_start_existing_clustername(self, monkeypatch):
@@ -195,7 +195,7 @@ async def test_stop_normal(self):
     mock_client.delete_cluster.assert_called_once_with(
         project_id='test-stop',
         region=self.region,
-        cluster_name='dataprochub-fake')
+        cluster_name=f'dataprochub-{spawner.get_username()}')
 
   @pytest.mark.asyncio
   async def test_stop_no_cluster(self):
@@ -374,13 +374,11 @@ def test_minimium_cluster_definition(self, monkeypatch):
 
     def test_read_file(*args, **kwargs):
       config_string = open('./tests/test_data/minimum.yaml', 'r').read()
-      print(config_string)
       return config_string
 
     def test_read_file_preview(*args, **kwargs):
       config_string = open('./tests/test_data/minimum.yaml', 'r').read()
       config_string = config_string.replace('1.4.16', 'preview')
-      print(config_string)
       return config_string
 
     def test_read_file_2_0(*args, **kwargs):
@@ -435,7 +433,6 @@ def test_clustername(*args, **kwargs):
     monkeypatch.setattr(spawner, "read_gcs_file", test_read_file_preview)
 
     config_built = spawner._build_cluster_config()
-    print(config_built)
     assert Component['JUPYTER'].value in config_built['config']['software_config']['optional_components']
     assert Component['ANACONDA'].value not in config_built['config']['software_config']['optional_components']
 
@@ -534,6 +531,7 @@ def test_clustername(*args, **kwargs):
     assert config_built['config']['software_config']['properties']['dataproc:jupyter.hub.env'] == 'test-env-str'
     assert config_built['config']['software_config']['properties']['dataproc:jupyter.hub.menu.enabled'] == 'true'
     assert 'dataproc:jupyter.hub.enabled' not in config_built['config']['software_config']['properties']
+    assert 'dataproc:dataproc.personal-auth.user' not in config_built['config']['software_config']['properties']
 
   def test_cluster_definition_overrides(self, monkeypatch):
     """Check that config settings incompatible with JupyterHub are overwritten correctly."""
@@ -665,7 +663,7 @@ def test_clustername(*args, **kwargs):
             'KeyCamelCase': 'UlowUlow',
             'key_with_underscore': 'https://downloads.io/protected/files/enterprise-trial.tar.gz',
             'key_with_underscore_too': 'some_UPPER_and_UlowerU:1234',
-            'session-user': MockUser.name
+            'session-user': spawner.get_username()
           },
           'zone_uri': 'https://www.googleapis.com/compute/v1/projects/test-project/zones/test-form1-a'
         },
@@ -689,7 +687,7 @@ def test_clustername(*args, **kwargs):
             'dataproc:jupyter.hub.env': 'test-env-str',
             'dataproc:jupyter.hub.menu.enabled': 'true',
             'dataproc:jupyter.instance-tag.enabled': 'false',
-            'dataproc:jupyter.notebook.gcs.dir': 'gs://users-notebooks/fake',
+            'dataproc:jupyter.notebook.gcs.dir': f'gs://users-notebooks/{spawner.get_username()}',
             'key-with-dash:UPPER_UPPER': '4000',
             'key-with-dash-too:UlowUlowUlow': '85196m',
             'key:and.multiple.dots.lowUlowUlow': '13312m'
@@ -775,7 +773,7 @@ def test_clustername(*args, **kwargs):
     assert config_built['config']['gce_cluster_config']['metadata'] == {
       'm1': 'v1',
       'm2': 'v2',
-      'session-user': MockUser.name
+      'session-user': spawner.get_username()
     }
 
   def test_uris(self, monkeypatch):
@@ -998,7 +996,6 @@ def test_read_file(*args, **kwargs):
 
     # Now check that the config with resolved fields is correct as well
     config_built = spawner._build_cluster_config()
-    print(config_built)
 
     assert 'unknown_field_top_level' not in config_built
     assert 'unknown_field_config_level' not in config_built['config']
@@ -1193,3 +1190,77 @@ def test_image_version(*args, **kwargs):
 
     assert config_built['config']['software_config']['image_version'] == '1.5-debian10'
     assert config_built['config']['master_config']['image_uri'] == 'projects/test-project/global/images/custom-image'
+
+  def test_unified_auth_flag(self, monkeypatch):
+    fake_creds = AnonymousCredentials()
+    mock_dataproc_client = mock.create_autospec(ClusterControllerClient(credentials=fake_creds))
+    mock_gcs_client = mock.create_autospec(storage.Client(credentials=fake_creds, project='project'))
+    mock_compute_client = mock.create_autospec(discovery.build('compute', 'v1',
+                                               credentials=fake_creds, cache_discovery=False))
+    spawner = DataprocSpawner(hub=Hub(), dataproc=mock_dataproc_client, gcs=mock_gcs_client,
+                              user=MockUser(), _mock=True, gcs_notebooks=self.gcs_notebooks,
+                              compute=mock_compute_client, project='test-project')
+
+    spawner.force_single_user = True
+    spawner.env_str = "test-env-str"
+    spawner.args_str = "test-args-str"
+    config_built = spawner._build_cluster_config()
+    assert (config_built['config']['software_config']['properties']
+        ['dataproc:dataproc.personal-auth.user']) == spawner.user.name
+
+  def test_unified_auth_yaml(self, monkeypatch):
+    fake_creds = AnonymousCredentials()
+    mock_dataproc_client = mock.create_autospec(ClusterControllerClient(credentials=fake_creds))
+    mock_gcs_client = mock.create_autospec(storage.Client(credentials=fake_creds, project='project'))
+    mock_compute_client = mock.create_autospec(discovery.build('compute', 'v1',
+                                               credentials=fake_creds, cache_discovery=False))
+    spawner = DataprocSpawner(hub=Hub(), dataproc=mock_dataproc_client, gcs=mock_gcs_client,
+                              user=MockUser(), _mock=True, gcs_notebooks=self.gcs_notebooks,
+                              compute=mock_compute_client, project='test-project')
+
+    def test_read_file(*args, **kwargs):
+      config_string = open('./tests/test_data/perso.yaml', 'r').read()
+      return config_string
+
+    monkeypatch.setattr(spawner, "read_gcs_file", test_read_file)
+    spawner.env_str = "test-env-str"
+    spawner.args_str = "test-args-str"
+    spawner.user_options = {
+      'cluster_type': 'perso.yaml',
+      'cluster_zone': 'us-central-1'
+    }
+
+    config_built = spawner._build_cluster_config()
+
+    assert (config_built['config']['software_config']['properties']
+        ['dataproc:dataproc.personal-auth.user']) == spawner.user.name
+
+  def test_unified_auth_user(self, monkeypatch):
+    fake_creds = AnonymousCredentials()
+    mock_dataproc_client = mock.create_autospec(ClusterControllerClient(credentials=fake_creds))
+    mock_gcs_client = mock.create_autospec(storage.Client(credentials=fake_creds, project='project'))
+    mock_compute_client = mock.create_autospec(discovery.build('compute', 'v1',
+                                               credentials=fake_creds, cache_discovery=False))
+    spawner = DataprocSpawner(hub=Hub(), dataproc=mock_dataproc_client, gcs=mock_gcs_client,
+                              user=MockUser(), _mock=True, gcs_notebooks=self.gcs_notebooks,
+                              compute=mock_compute_client, project='test-project')
+
+    def test_read_file(*args, **kwargs):
+      config_string = open('./tests/test_data/perso.yaml', 'r').read()
+      return config_string
+
+    monkeypatch.setattr(spawner, "read_gcs_file", test_read_file)
+    spawner.env_str = "test-env-str"
+    spawner.args_str = "test-args-str"
+    spawner.allow_custom_clusters = True
+    spawner.user_options = {
+      'cluster_type': 'perso.yaml',
+      'cluster_zone': 'us-central-1',
+      "cluster_props_prefix_0": "dataproc",
+      "cluster_props_key_0": "dataproc.personal-auth.user",
+      "cluster_props_val_0": "[email protected]"
+    }
+    config_built = spawner._build_cluster_config()
+    assert (config_built['config']['software_config']['properties']
+        ['dataproc:dataproc.personal-auth.user']) == spawner.user.name
+