Fix user check for jira ticketing

Author: chi-yelp

service_configuration_lib/spark_config.py

@@ -77,6 +77,9 @@
 
 SUPPORTED_CLUSTER_MANAGERS = ['kubernetes', 'local']
 DEFAULT_SPARK_RUN_CONFIG = '/nail/srv/configs/spark.yaml'
+USER_BATCH = 'batch'  # used by batch servers
+USER_TRON = 'TRON'  # used by Tron jobs, or other paasta CLI commands such as `paasta validate/mark-for-deployment`
+USER_UNSPECIFIED = 'UNSPECIFIED'
 
 log = logging.Logger(__name__)
 log.setLevel(logging.WARN)
@@ -305,7 +308,7 @@ def _get_k8s_spark_env(
     service_account_name: Optional[str] = None,
     include_self_managed_configs: bool = True,
     k8s_server_address: Optional[str] = None,
-    user: Optional[str] = None,
+    user: Optional[str] = USER_UNSPECIFIED,
     jira_ticket: Optional[str] = None,
 ) -> Dict[str, str]:
     # RFC 1123: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#dns-label-names
@@ -314,8 +317,6 @@ def _get_k8s_spark_env(
     _paasta_cluster = utils.get_k8s_resource_name_limit_size_with_hash(paasta_cluster)
     _paasta_service = utils.get_k8s_resource_name_limit_size_with_hash(paasta_service)
     _paasta_instance = utils.get_k8s_resource_name_limit_size_with_hash(paasta_instance)
-    if not user:
-        user = os.environ.get('USER', 'UNSPECIFIED')
 
     spark_env = {
         'spark.master': f'k8s://https://k8s.{paasta_cluster}.paasta:6443',
@@ -1040,6 +1041,7 @@ def get_spark_conf(
         :param service_account_name: The k8s service account to use for spark k8s authentication.
         :param force_spark_resource_configs: skip the resource/instances recalculation.
             This is strongly not recommended.
+        :param user: the user who is running the spark job.
         :returns: spark opts in a dict.
         """
         # Mesos deprecation
@@ -1051,8 +1053,11 @@ def get_spark_conf(
         # is str type.
         user_spark_opts = _convert_user_spark_opts_value_to_str(user_spark_opts)
 
+        # Get user from environment variables if it's not set
+        user = user or os.environ.get('USER', None)
+
         if self.mandatory_default_spark_srv_conf.get('spark.yelp.jira_ticket.enabled') == 'true':
-            needs_jira_check = os.environ.get('USER', '') not in ['batch', 'TRON', '']
+            needs_jira_check = user not in [USER_BATCH, USER_TRON, None]
             if needs_jira_check:
                 valid_ticket = self._get_valid_jira_ticket(jira_ticket)
                 if valid_ticket is None:

tests/spark_config_test.py

@@ -623,6 +623,20 @@ def gpu_pool(self, tmpdir, monkeypatch):
                 },
                 False,
             ),
+            (
+                'Local spark cluster',
+                'local',
+                {
+                    'spark.executor.cores': '2',
+                    'spark.executor.instances': '600',
+                },
+                {
+                    'spark.executor.memory': '28g',
+                    'spark.executor.cores': '4',
+                    'spark.executor.instances': '600',
+                },
+                False,
+            ),
         ],
     )
     def test_adjust_spark_requested_resources(
@@ -1887,35 +1901,22 @@ def test_get_spark_conf_with_jira_validation_disabled(self, mock_spark_srv_conf_
         assert 'spark.kubernetes.executor.label.spark.yelp.com/jira_ticket' not in result
 
     @pytest.mark.parametrize(
-        'user_env,should_check', [
+        'user,should_check', [
             ('regular_user', True),
             ('batch', False),
             ('TRON', False),
-            ('', False),
+            (None, False),
         ],
     )
     def test_jira_ticket_check_for_different_users(
-        self, user_env, should_check, mock_spark_srv_conf_file_with_jira_enabled, mock_log,
+        self, user, should_check, mock_spark_srv_conf_file_with_jira_enabled, mock_log,
     ):
         """Test that Jira ticket validation is skipped for certain users."""
-        with mock.patch.dict(os.environ, {'USER': user_env}):
-            spark_conf_builder = spark_config.SparkConfBuilder()
+        spark_conf_builder = spark_config.SparkConfBuilder()
 
-            if should_check:
-                # For regular users, validation should be enforced
-                with pytest.raises(RuntimeError):
-                    spark_conf_builder.get_spark_conf(
-                        cluster_manager='kubernetes',
-                        spark_app_base_name='test-app',
-                        user_spark_opts={},
-                        paasta_cluster='test-cluster',
-                        paasta_pool='test-pool',
-                        paasta_service='test-service',
-                        paasta_instance='test-instance',
-                        docker_img='test-image',
-                    )
-            else:
-                # For special users, validation should be skipped
+        if should_check:
+            # For regular users, validation should be enforced
+            with pytest.raises(RuntimeError):
                 spark_conf_builder.get_spark_conf(
                     cluster_manager='kubernetes',
                     spark_app_base_name='test-app',
@@ -1925,5 +1926,19 @@ def test_jira_ticket_check_for_different_users(
                     paasta_service='test-service',
                     paasta_instance='test-instance',
                     docker_img='test-image',
+                    user=user,
                 )
-                mock_log.debug.assert_called_with('Jira ticket check not required for this job configuration.')
+        else:
+            # For special users, validation should be skipped
+            spark_conf_builder.get_spark_conf(
+                cluster_manager='kubernetes',
+                spark_app_base_name='test-app',
+                user_spark_opts={},
+                paasta_cluster='test-cluster',
+                paasta_pool='test-pool',
+                paasta_service='test-service',
+                paasta_instance='test-instance',
+                docker_img='test-image',
+                user=user,
+            )
+            mock_log.debug.assert_called_with('Jira ticket check not required for this job configuration.')