Improve slfj logger input for MDC field:datasource (#6787)

* improve slfj logger  MDC datasource input

* add some UT and isNested field

server/src/main/java/org/apache/druid/server/log/LoggingRequestLogger.java

@@ -20,14 +20,20 @@
 package org.apache.druid.server.log;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.DataSource;
 import org.apache.druid.query.Query;
+import org.apache.druid.query.QueryDataSource;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.query.UnionDataSource;
 import org.apache.druid.server.RequestLogLine;
 import org.slf4j.MDC;
 
 import java.io.IOException;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 public class LoggingRequestLogger implements RequestLogger
 {
@@ -58,8 +64,9 @@ public class LoggingRequestLogger implements RequestLogger
         try {
           final Query query = requestLogLine.getQuery();
           MDC.put("queryId", query.getId());
-          MDC.put("dataSource", query.getDataSource().toString());
+          MDC.put("dataSource", findInnerDatasource(query).toString());
           MDC.put("queryType", query.getType());
+          MDC.put("isNested", String.valueOf(!(query.getDataSource() instanceof TableDataSource)));
           MDC.put("hasFilters", Boolean.toString(query.hasFilters()));
           MDC.put("remoteAddr", requestLogLine.getRemoteAddr());
           MDC.put("duration", query.getDuration().toString());
@@ -103,6 +110,30 @@ public class LoggingRequestLogger implements RequestLogger
     return setContextMDC;
   }
 
+  private Object findInnerDatasource(Query query)
+  {
+    DataSource _ds = query.getDataSource();
+    if (_ds instanceof TableDataSource) {
+      return ((TableDataSource) _ds).getName();
+    }
+    if (_ds instanceof QueryDataSource) {
+      return findInnerDatasource(((QueryDataSource) _ds).getQuery());
+    }
+    if (_ds instanceof UnionDataSource) {
+      return Joiner.on(",")
+                   .join(
+                       ((UnionDataSource) _ds)
+                           .getDataSources()
+                           .stream()
+                           .map(TableDataSource::getName)
+                           .collect(Collectors.toList())
+                   );
+    } else {
+      // should not come here
+      return query.getDataSource();
+    }
+  }
+
   @Override
   public String toString()
   {

server/src/test/java/org/apache/druid/server/log/LoggingRequestLoggerTest.java

@@ -21,6 +21,7 @@ package org.apache.druid.server.log;
 
 import com.fasterxml.jackson.annotation.JsonTypeName;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import org.apache.druid.jackson.DefaultObjectMapper;
 import org.apache.druid.java.util.common.DateTimes;
@@ -30,8 +31,10 @@ import org.apache.druid.query.BaseQuery;
 import org.apache.druid.query.DataSource;
 import org.apache.druid.query.LegacyDataSource;
 import org.apache.druid.query.Query;
+import org.apache.druid.query.QueryDataSource;
 import org.apache.druid.query.QueryRunner;
 import org.apache.druid.query.QuerySegmentWalker;
+import org.apache.druid.query.UnionDataSource;
 import org.apache.druid.query.filter.DimFilter;
 import org.apache.druid.query.spec.QuerySegmentSpec;
 import org.apache.druid.server.QueryStats;
@@ -84,6 +87,62 @@ public class LoggingRequestLoggerTest
         }
       }, false, queryContext
   );
+
+  final Query nestedQuery = new FakeQuery(
+      new QueryDataSource(query),
+      new QuerySegmentSpec()
+      {
+        @Override
+        public List<Interval> getIntervals()
+        {
+          return Collections.singletonList(Intervals.of("2016-01-01T00Z/2016-01-02T00Z"));
+        }
+
+        @Override
+        public <T> QueryRunner<T> lookup(Query<T> query, QuerySegmentWalker walker)
+        {
+          return null;
+        }
+      }, false, queryContext
+  );
+
+  final Query nestedNestedQuery = new FakeQuery(
+      new QueryDataSource(nestedQuery),
+      new QuerySegmentSpec()
+      {
+        @Override
+        public List<Interval> getIntervals()
+        {
+          return Collections.singletonList(Intervals.of("2016-01-01T00Z/2016-01-02T00Z"));
+        }
+
+        @Override
+        public <T> QueryRunner<T> lookup(Query<T> query, QuerySegmentWalker walker)
+        {
+          return null;
+        }
+      }, false, queryContext
+  );
+
+  final Query unionQuery = new FakeQuery(
+      new UnionDataSource(ImmutableList.of(new LegacyDataSource("A"), new LegacyDataSource("B"))),
+      new QuerySegmentSpec()
+      {
+        @Override
+        public List<Interval> getIntervals()
+        {
+          return Collections.singletonList(Intervals.of("2016-01-01T00Z/2016-01-02T00Z"));
+        }
+
+        @Override
+        public <T> QueryRunner<T> lookup(Query<T> query, QuerySegmentWalker walker)
+        {
+          return null;
+        }
+      }, false, queryContext
+  );
+
+
   final QueryStats queryStats = new QueryStats(ImmutableMap.of());
   final RequestLogLine logLine = new RequestLogLine(
       timestamp,
@@ -141,6 +200,7 @@ public class LoggingRequestLoggerTest
     Assert.assertEquals("fake", map.get("queryType"));
     Assert.assertEquals("some.host.tld", map.get("remoteAddr"));
     Assert.assertEquals("false", map.get("descending"));
+    Assert.assertEquals("false", map.get("isNested"));
     Assert.assertNull(map.get("foo"));
   }
 
@@ -156,9 +216,73 @@ public class LoggingRequestLoggerTest
     Assert.assertEquals("fake", map.get("queryType"));
     Assert.assertEquals("some.host.tld", map.get("remoteAddr"));
     Assert.assertEquals("false", map.get("descending"));
+    Assert.assertEquals("false", map.get("isNested"));
     Assert.assertEquals("bar", map.get("foo"));
   }
 
+  @Test
+  public void testNestedQueryLoggingMDC() throws Exception
+  {
+    final LoggingRequestLogger requestLogger = new LoggingRequestLogger(new DefaultObjectMapper(), true, false);
+    requestLogger.log(new RequestLogLine(
+        timestamp,
+        remoteAddr,
+        nestedQuery,
+        queryStats
+    ));
+    final Map<String, Object> map = readContextMap(baos.toByteArray());
+    Assert.assertEquals("datasource", map.get("dataSource"));
+    Assert.assertEquals("PT86400S", map.get("duration"));
+    Assert.assertEquals("false", map.get("hasFilters"));
+    Assert.assertEquals("fake", map.get("queryType"));
+    Assert.assertEquals("some.host.tld", map.get("remoteAddr"));
+    Assert.assertEquals("false", map.get("descending"));
+    Assert.assertEquals("true", map.get("isNested"));
+    Assert.assertNull(map.get("foo"));
+  }
+
+  @Test
+  public void testNestedNestedQueryLoggingMDC() throws Exception
+  {
+    final LoggingRequestLogger requestLogger = new LoggingRequestLogger(new DefaultObjectMapper(), true, false);
+    requestLogger.log(new RequestLogLine(
+        timestamp,
+        remoteAddr,
+        nestedNestedQuery,
+        queryStats
+    ));
+    final Map<String, Object> map = readContextMap(baos.toByteArray());
+    Assert.assertEquals("datasource", map.get("dataSource"));
+    Assert.assertEquals("PT86400S", map.get("duration"));
+    Assert.assertEquals("false", map.get("hasFilters"));
+    Assert.assertEquals("fake", map.get("queryType"));
+    Assert.assertEquals("some.host.tld", map.get("remoteAddr"));
+    Assert.assertEquals("true", map.get("isNested"));
+    Assert.assertEquals("false", map.get("descending"));
+    Assert.assertNull(map.get("foo"));
+  }
+
+  @Test
+  public void testUnionQueryLoggingMDC() throws Exception
+  {
+    final LoggingRequestLogger requestLogger = new LoggingRequestLogger(new DefaultObjectMapper(), true, false);
+    requestLogger.log(new RequestLogLine(
+        timestamp,
+        remoteAddr,
+        unionQuery,
+        queryStats
+    ));
+    final Map<String, Object> map = readContextMap(baos.toByteArray());
+    Assert.assertEquals("A,B", map.get("dataSource"));
+    Assert.assertEquals("true", map.get("isNested"));
+    Assert.assertEquals("PT86400S", map.get("duration"));
+    Assert.assertEquals("false", map.get("hasFilters"));
+    Assert.assertEquals("fake", map.get("queryType"));
+    Assert.assertEquals("some.host.tld", map.get("remoteAddr"));
+    Assert.assertEquals("false", map.get("descending"));
+    Assert.assertNull(map.get("foo"));
+  }
+
   private static Map<String, Object> readContextMap(byte[] bytes) throws Exception
   {
     final Map<String, Object> rawMap = mapper.readValue(bytes, JacksonUtils.TYPE_REFERENCE_MAP_STRING_OBJECT);