Set AWS creds via Spark Config (#78)

Set AWS creds via Spark Configs

Author: jsleight

service_configuration_lib/spark_config.py

@@ -53,9 +53,6 @@
     'spark.executorEnv.PAASTA_INSTANCE',
     'spark.executorEnv.PAASTA_CLUSTER',
     'spark.executorEnv.SPARK_EXECUTOR_DIRS',
-    'spark.hadoop.fs.s3a.access.key',
-    'spark.hadoop.fs.s3a.secret.key',
-    'spark.hadoop.fs.s3a.session.token',
     'spark.kubernetes.pyspark.pythonVersion',
     'spark.kubernetes.container.image',
     'spark.kubernetes.namespace',
@@ -303,6 +300,31 @@ def _append_event_log_conf(
     return spark_opts
 
 
+def _append_aws_credentials_conf(
+    spark_opts: Dict[str, str],
+    access_key: Optional[str],
+    secret_key: Optional[str],
+    session_token: Optional[str] = None,
+) -> Dict[str, str]:
+    """It is important that we set the aws creds via the spark configs and not via the AWS
+    environment variables.  See HADOOP-18233 for details
+
+    We set both s3a and s3 credentials because s3a is the actual hadoop-aws driver, but our
+    glue-metatore integration will attempt to use s3 path drivers which are monkeypatched
+    to use the the s3a driver.
+    """
+    if access_key:
+        spark_opts['spark.hadoop.fs.s3a.access.key'] = access_key
+        spark_opts['spark.hadoop.fs.s3.access.key'] = access_key
+    if secret_key:
+        spark_opts['spark.hadoop.fs.s3a.secret.key'] = secret_key
+        spark_opts['spark.hadoop.fs.s3.secret.key'] = secret_key
+    if session_token:
+        spark_opts['spark.hadoop.fs.s3a.session.token'] = session_token
+        spark_opts['spark.hadoop.fs.s3.session.token'] = session_token
+    return spark_opts
+
+
 def _adjust_spark_requested_resources(
     user_spark_opts: Dict[str, str],
     cluster_manager: str,
@@ -727,6 +749,8 @@ def get_spark_conf(
 
     # configure spark Console Progress
     spark_conf = _append_spark_config(spark_conf, 'spark.ui.showConsoleProgress', 'true')
+
+    spark_conf = _append_aws_credentials_conf(spark_conf, *aws_creds)
     return spark_conf
 
 

setup.py

@@ -17,7 +17,7 @@
 
 setup(
     name='service-configuration-lib',
-    version='2.10.3',
+    version='2.10.4',
     provides=['service_configuration_lib'],
     description='Start, stop, and inspect Yelp SOA services',
     url='https://github.com/Yelp/service_configuration_lib',

tests/spark_config_test.py

@@ -643,6 +643,17 @@ def test_append_console_progress_conf(
 
         assert output[key] == expected_output
 
+    def test_append_aws_credentials_conf(self):
+        output = spark_config._append_aws_credentials_conf(
+            {},
+            mock.sentinel.access,
+            mock.sentinel.secret,
+            mock.sentinel.token,
+        )
+        assert output['spark.hadoop.fs.s3a.access.key'] == mock.sentinel.access
+        assert output['spark.hadoop.fs.s3a.secret.key'] == mock.sentinel.secret
+        assert output['spark.hadoop.fs.s3a.session.token'] == mock.sentinel.token
+
     @pytest.fixture
     def mock_append_spark_conf_log(self):
         return_value = {'spark.logConf': 'true'}
@@ -682,6 +693,19 @@ def mock_append_event_log_conf(self):
         with MockConfigFunction('_append_event_log_conf', return_value) as m:
             yield m
 
+    @pytest.fixture
+    def mock_append_aws_credentials_conf(self):
+        return_value = {
+            'spark.hadoop.fs.s3a.access.key': 'my_key',
+            'spark.hadoop.fs.s3.access.key': 'my_key',
+            'spark.hadoop.fs.s3a.secret.key': 'your_key',
+            'spark.hadoop.fs.s3.secret.key': 'your_key',
+            'spark.hadoop.fs.s3a.session.token': 'ice_cream',
+            'spark.hadoop.fs.s3.session.token': 'ice_cream',
+        }
+        with MockConfigFunction('_append_aws_credentials_conf', return_value) as m:
+            yield m
+
     @pytest.fixture
     def mock_adjust_spark_requested_resources_mesos(self):
         return_value = {
@@ -909,6 +933,7 @@ def test_get_spark_conf_mesos(
         extra_docker_params,
         mock_get_mesos_docker_volumes_conf,
         mock_append_event_log_conf,
+        mock_append_aws_credentials_conf,
         mock_append_sql_shuffle_partitions_conf,
         mock_adjust_spark_requested_resources_mesos,
         mock_time,
@@ -964,6 +989,7 @@ def test_get_spark_conf_mesos(
             list(mock_get_mesos_docker_volumes_conf.return_value.keys()) +
             list(mock_adjust_spark_requested_resources_mesos.return_value.keys()) +
             list(mock_append_event_log_conf.return_value.keys()) +
+            list(mock_append_aws_credentials_conf.return_value.keys()) +
             list(mock_append_sql_shuffle_partitions_conf.return_value.keys()) +
             list(mock_append_spark_conf_log.return_value.keys()) +
             list(mock_append_console_progress_conf.return_value.keys()),
@@ -978,6 +1004,7 @@ def test_get_spark_conf_mesos(
         mock_append_event_log_conf.mocker.assert_called_once_with(
             mock.ANY, *aws_creds,
         )
+        mock_append_aws_credentials_conf.mocker.assert_called_once_with(mock.ANY, *aws_creds)
         mock_append_sql_shuffle_partitions_conf.mocker.assert_called_once_with(
             mock.ANY,
         )
@@ -1030,6 +1057,7 @@ def tes_leaderst_get_spark_conf_kubernetes(
         spark_opts_from_env,
         ui_port,
         mock_append_event_log_conf,
+        mock_append_aws_credentials_conf,
         mock_append_sql_shuffle_partitions_conf,
         mock_adjust_spark_requested_resources_kubernetes,
         mock_time,
@@ -1067,6 +1095,7 @@ def tes_leaderst_get_spark_conf_kubernetes(
             list(other_spark_opts.keys()) +
             list(mock_adjust_spark_requested_resources_kubernetes.return_value.keys()) +
             list(mock_append_event_log_conf.return_value.keys()) +
+            list(mock_append_aws_credentials_conf.return_value.keys()) +
             list(mock_append_sql_shuffle_partitions_conf.return_value.keys()),
         )
         assert len(set(output.keys()) - verified_keys) == 0
@@ -1076,6 +1105,7 @@ def tes_leaderst_get_spark_conf_kubernetes(
         mock_append_event_log_conf.mocker.assert_called_once_with(
             mock.ANY, *aws_creds,
         )
+        mock_append_aws_credentials_conf.mocker.assert_called_once_with(mock.ANY, *aws_creds)
         mock_append_sql_shuffle_partitions_conf.mocker.assert_called_once_with(
             mock.ANY,
         )