Add extra tag for env (#1750)

Script tokens can have the same name, add env tag to differentiate.
pinterest · Nov 14, 2024 · 62ad465 · 62ad465
1 parent df4cc0d
commit 62ad465
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/...service/teletraanservice/src/test/java/com/pinterest/teletraan/handler/EnvStagesTest.java b/...service/teletraanservice/src/test/java/com/pinterest/teletraan/handler/EnvStagesTest.java
@@ -67,7 +67,7 @@ public void updateEnvStageOffDev_unsetsAllowPrivateBuild() throws Exception {
         origBean.setStage_type(EnvType.DEV);
         envBean.setStage_type(EnvType.STAGING);
         Mockito.when(environDAO.getByStage(Mockito.anyString(), Mockito.anyString()))
-            .thenReturn(origBean);
+                .thenReturn(origBean);
         SecurityContext mockSC = mock(SecurityContext.class);
         Principal mockPrincipal = mock(Principal.class);
         Mockito.when(mockSC.getUserPrincipal()).thenReturn(mockPrincipal);

diff --git a/...al/src/main/java/com/pinterest/teletraan/universal/security/ScriptTokenAuthenticator.java b/...al/src/main/java/com/pinterest/teletraan/universal/security/ScriptTokenAuthenticator.java
@@ -22,6 +22,7 @@
 import io.dropwizard.auth.Authenticator;
 import io.micrometer.core.instrument.Counter;
 import io.micrometer.core.instrument.Metrics;
+import java.util.Objects;
 import java.util.Optional;
 
 /**
@@ -51,7 +52,11 @@ public Optional<ScriptTokenPrincipal<R>> authenticate(String credentials)
             Optional<ScriptTokenPrincipal<R>> principal = tokenProvider.getPrincipal(credentials);
             if (principal.isPresent()) {
                 scriptTokenSuccessCounterBuilder
-                        .tags("principal", principal.get().getName())
+                        .tags(
+                                "principal",
+                                principal.get().getName(),
+                                "env",
+                                Objects.toString(principal.get().getResource().getEnvName(), "NA"))
                         .register(Metrics.globalRegistry)
                         .increment();
             } else {

diff --git a/...rc/test/java/com/pinterest/teletraan/universal/security/ScriptTokenAuthenticatorTest.java b/...rc/test/java/com/pinterest/teletraan/universal/security/ScriptTokenAuthenticatorTest.java
@@ -23,6 +23,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.pinterest.teletraan.universal.security.bean.AuthZResource;
 import com.pinterest.teletraan.universal.security.bean.ScriptTokenPrincipal;
 import com.pinterest.teletraan.universal.security.bean.ValueBasedRole;
 import io.dropwizard.auth.AuthenticationException;
@@ -50,12 +51,14 @@ void setUp() throws Exception {
         registry = new SimpleMeterRegistry();
         Metrics.addRegistry(registry);
 
+        scriptTokenPrincipal =
+                new ScriptTokenPrincipal<ValueBasedRole>(
+                        PRINCIPAL_NAME, new ValueBasedRole(1), new AuthZResource("rId", "rType"));
+
         scriptTokenProvider = mock(ScriptTokenProvider.class);
-        scriptTokenPrincipal = mock(ScriptTokenPrincipal.class);
         when(scriptTokenProvider.getPrincipal(anyString())).thenReturn(Optional.empty());
         when(scriptTokenProvider.getPrincipal(CREDENTIALS))
                 .thenReturn(Optional.of(scriptTokenPrincipal));
-        when(scriptTokenPrincipal.getName()).thenReturn(PRINCIPAL_NAME);
 
         sut = new ScriptTokenAuthenticator<>(scriptTokenProvider);
     }