github.com/honeycombio/honeytail@v1.9.0/parsers/mongodb/queryshape/shape.go (about)

     1  package queryshape
     2  
     3  import (
     4  	"sort"
     5  	"strings"
     6  )
     7  
     8  // rough rules XXX(toshok) needs editing
     9  // 1. if key is not an op:
    10  //    1a. if value is a primitive, set value = 1
    11  //    1b. if value is an aggregate, walk subtree flattening everything but ops and their values if necessary
    12  // 2. if key is an op:
    13  //    2a. if value is a primitive, set value = 1
    14  //    2b. if value is a map, keep map + all keys, and process keys (starting at step 1)
    15  //    2c. if value is a list, walk list.
    16  //        2c1. if all values are primitive, set value = 1
    17  //        2c2. if any values are maps/lists, keep map + all keys, and process keys (starting at step 1)
    18  
    19  func GetQueryShape(q map[string]interface{}) string {
    20  	if q_, ok := q["$query"].(map[string]interface{}); ok {
    21  		return GetQueryShape(q_)
    22  	}
    23  
    24  	pruned := make(map[string]interface{})
    25  	for k, v := range q {
    26  		if k[0] == '$' {
    27  			pruned[k] = flattenOp(v)
    28  		} else {
    29  			pruned[k] = flatten(v)
    30  		}
    31  	}
    32  
    33  	// flatten pruned to a string, sorting keys alphabetically ($ coming before a/A)
    34  	return serializeShape(pruned)
    35  }
    36  
    37  func isAggregate(v interface{}) bool {
    38  	if _, ok := v.([]interface{}); ok {
    39  		return true
    40  	} else if _, ok := v.(map[string]interface{}); ok {
    41  		return true
    42  	}
    43  	return false
    44  }
    45  
    46  func flattenSlice(slice []interface{}, fromOp bool) interface{} {
    47  	var rv []interface{}
    48  	for _, v := range slice {
    49  		if s, ok := v.([]interface{}); ok {
    50  			sv := flattenSlice(s, false)
    51  			if isAggregate(sv) {
    52  				rv = append(rv, sv)
    53  			}
    54  		} else if m, ok := v.(map[string]interface{}); ok {
    55  			mv := flattenMap(m, fromOp)
    56  			if isAggregate(mv) || fromOp {
    57  				rv = append(rv, mv)
    58  			}
    59  		}
    60  	}
    61  	// if the slice is empty, return 1 (since it's entirely primitives).
    62  	// otherwise return the slice
    63  	if len(rv) == 0 {
    64  		return 1
    65  	}
    66  	return rv
    67  }
    68  
    69  func flattenMap(m map[string]interface{}, fromOp bool) interface{} {
    70  	rv := make(map[string]interface{})
    71  	for k, v := range m {
    72  		if k[0] == '$' {
    73  			rv[k] = flattenOp(v)
    74  		} else {
    75  			flattened := flatten(v)
    76  			if isAggregate(flattened) || fromOp {
    77  				rv[k] = flattened
    78  			}
    79  		}
    80  	}
    81  	// if the slice is empty, return 1 (since it's entirely primitives).
    82  	// otherwise return the slice
    83  	if len(rv) == 0 {
    84  		return 1
    85  	}
    86  	return rv
    87  }
    88  
    89  func flatten(v interface{}) interface{} {
    90  	if s, ok := v.([]interface{}); ok {
    91  		return flattenSlice(s, false)
    92  	} else if m, ok := v.(map[string]interface{}); ok {
    93  		return flattenMap(m, false)
    94  	} else {
    95  		return 1
    96  	}
    97  }
    98  
    99  func flattenOp(v interface{}) interface{} {
   100  	if s, ok := v.([]interface{}); ok {
   101  		return flattenSlice(s, true)
   102  	} else if m, ok := v.(map[string]interface{}); ok {
   103  		return flattenMap(m, true)
   104  	} else {
   105  		return 1
   106  	}
   107  }
   108  
   109  func serializeShape(shape interface{}) string {
   110  	// we can't just json marshal, since we need ordered keys
   111  	if m, ok := shape.(map[string]interface{}); ok {
   112  		var keys []string
   113  		var keyAndVal []string
   114  		for k := range m {
   115  			keys = append(keys, k)
   116  		}
   117  
   118  		sort.Strings(keys)
   119  		for _, k := range keys {
   120  			keyAndVal = append(keyAndVal, "\""+k+"\": "+serializeShape(m[k]))
   121  		}
   122  
   123  		return "{ " + strings.Join(keyAndVal, ", ") + " }"
   124  
   125  	} else if s, ok := shape.([]interface{}); ok {
   126  		var vals []string
   127  		for _, v := range s {
   128  			vals = append(vals, serializeShape(v))
   129  		}
   130  		return "[ " + strings.Join(vals, ", ") + " ]"
   131  	} else {
   132  		return "1"
   133  	}
   134  }