feat(system-metrics): track api usage by user, client, api

Author: david-leifker

metadata-io/build.gradle

@@ -31,8 +31,6 @@ dependencies {
 
   implementation externalDependency.guava
   implementation externalDependency.reflections
-  // https://mvnrepository.com/artifact/nl.basjes.parse.useragent/yauaa
-  implementation 'nl.basjes.parse.useragent:yauaa:7.27.0'
 
   api(externalDependency.dgraph4j) {
     exclude group: 'com.google.guava', module: 'guava'

metadata-io/src/main/java/com/linkedin/metadata/dao/throttle/APIThrottle.java

@@ -12,17 +12,9 @@
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import nl.basjes.parse.useragent.UserAgent;
-import nl.basjes.parse.useragent.UserAgentAnalyzer;
 
 public class APIThrottle {
   private static final Set<String> AGENT_EXEMPTIONS = Set.of("Browser");
-  private static final UserAgentAnalyzer UAA =
-      UserAgentAnalyzer.newBuilder()
-          .hideMatcherLoadStats()
-          .withField(UserAgent.AGENT_CLASS)
-          .withCache(1000)
-          .build();
 
   private APIThrottle() {}
 
@@ -56,13 +48,11 @@ public static void evaluate(
   private static boolean isExempt(@Nullable RequestContext requestContext) {
     // Exclude internal calls
     if (requestContext == null
-        || requestContext.getUserAgent() == null
+        || requestContext.getAgentClass() == null
         || requestContext.getUserAgent().isEmpty()) {
       return true;
     }
-
-    UserAgent ua = UAA.parse(requestContext.getUserAgent());
-    return AGENT_EXEMPTIONS.contains(ua.get(UserAgent.AGENT_CLASS).getValue());
+    return AGENT_EXEMPTIONS.contains(requestContext.getAgentClass());
   }
 
   private static Set<Long> eventMatchMaxWaitMs(

metadata-io/src/test/java/com/linkedin/metadata/dao/throttle/APIThrottleTest.java

@@ -12,6 +12,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import nl.basjes.parse.useragent.UserAgent;
 import org.testng.Assert;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
@@ -61,6 +62,7 @@ public void testExemptions() {
 
     for (ThrottleEvent event : ALL_EVENTS) {
       when(mockRequestContext.getUserAgent()).thenReturn(null);
+      when(mockRequestContext.getAgentClass()).thenReturn(null);
       try {
         APIThrottle.evaluate(opContext, Set.of(event), false);
       } catch (Exception ex) {
@@ -76,12 +78,16 @@ public void testExemptions() {
       for (String ua : exemptions) {
         try {
           when(mockRequestContext.getUserAgent()).thenReturn(ua);
+          when(mockRequestContext.getAgentClass())
+              .thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
           APIThrottle.evaluate(opContext, Set.of(event), true);
         } catch (Exception ex) {
           Assert.fail("Exception was thrown and NOT expected! " + event);
         }
         try {
           when(mockRequestContext.getUserAgent()).thenReturn(ua);
+          when(mockRequestContext.getAgentClass())
+              .thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
           APIThrottle.evaluate(opContext, Set.of(event), false);
         } catch (Exception ex) {
           Assert.fail("Exception was thrown and NOT expected! " + event);
@@ -106,6 +112,8 @@ public void testThrottleException() {
             && !event.getActiveThrottles().contains(MANUAL)) {
           try {
             when(mockRequestContext.getUserAgent()).thenReturn(ua);
+            when(mockRequestContext.getAgentClass())
+                .thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
             APIThrottle.evaluate(opContext, Set.of(event), true);
             Assert.fail(String.format("Exception WAS expected! %s %s", ua, event));
           } catch (Exception ignored) {
@@ -115,6 +123,8 @@ public void testThrottleException() {
             && !event.getActiveThrottles().contains(MANUAL)) {
           try {
             when(mockRequestContext.getUserAgent()).thenReturn(ua);
+            when(mockRequestContext.getAgentClass())
+                .thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
             APIThrottle.evaluate(opContext, Set.of(event), true);
           } catch (Exception ex) {
             Assert.fail(String.format("Exception was thrown and NOT expected! %s %s", ua, event));
@@ -126,6 +136,8 @@ public void testThrottleException() {
             && !event.getActiveThrottles().contains(MANUAL)) {
           try {
             when(mockRequestContext.getUserAgent()).thenReturn(ua);
+            when(mockRequestContext.getAgentClass())
+                .thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
             APIThrottle.evaluate(opContext, Set.of(event), false);
             Assert.fail(String.format("Exception WAS expected! %s %s", ua, event));
           } catch (Exception ignored) {
@@ -135,6 +147,8 @@ public void testThrottleException() {
             && !event.getActiveThrottles().contains(MANUAL)) {
           try {
             when(mockRequestContext.getUserAgent()).thenReturn(ua);
+            when(mockRequestContext.getAgentClass())
+                .thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
             APIThrottle.evaluate(opContext, Set.of(event), false);
           } catch (Exception ex) {
             Assert.fail(String.format("Exception was thrown and NOT expected! %s %s", ua, event));
@@ -145,12 +159,16 @@ public void testThrottleException() {
         if (event.getActiveThrottles().contains(MANUAL)) {
           try {
             when(mockRequestContext.getUserAgent()).thenReturn(ua);
+            when(mockRequestContext.getAgentClass())
+                .thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
             APIThrottle.evaluate(opContext, Set.of(event), true);
             Assert.fail(String.format("Exception WAS expected! %s %s", ua, event));
           } catch (Exception ignored) {
           }
           try {
             when(mockRequestContext.getUserAgent()).thenReturn(ua);
+            when(mockRequestContext.getAgentClass())
+                .thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
             APIThrottle.evaluate(opContext, Set.of(event), false);
             Assert.fail(String.format("Exception WAS expected! %s %s", ua, event));
           } catch (Exception ignored) {

metadata-operation-context/src/test/java/io/datahubproject/metadata/context/RequestContextTest.java

@@ -0,0 +1,4 @@
+package io.datahubproject.metadata.context;
+
+public class RequestContextTest {
+}