github.com/snowflakedb/gosnowflake@v1.9.0/htap.go (about)

     1  package gosnowflake
     2  
     3  import (
     4  	"sort"
     5  	"strconv"
     6  	"sync"
     7  )
     8  
     9  const (
    10  	queryContextCacheSizeParamName = "QUERY_CONTEXT_CACHE_SIZE"
    11  	defaultQueryContextCacheSize   = 5
    12  )
    13  
    14  type queryContext struct {
    15  	Entries []queryContextEntry `json:"entries,omitempty"`
    16  }
    17  
    18  type queryContextEntry struct {
    19  	ID        int    `json:"id"`
    20  	Timestamp int64  `json:"timestamp"`
    21  	Priority  int    `json:"priority"`
    22  	Context   string `json:"context,omitempty"`
    23  }
    24  
    25  type queryContextCache struct {
    26  	mutex   *sync.Mutex
    27  	entries []queryContextEntry
    28  }
    29  
    30  func (qcc *queryContextCache) init() *queryContextCache {
    31  	qcc.mutex = &sync.Mutex{}
    32  	return qcc
    33  }
    34  
    35  func (qcc *queryContextCache) add(sc *snowflakeConn, qces ...queryContextEntry) {
    36  	qcc.mutex.Lock()
    37  	defer qcc.mutex.Unlock()
    38  	if len(qces) == 0 {
    39  		qcc.prune(0)
    40  	} else {
    41  		for _, newQce := range qces {
    42  			logger.Debugf("adding query context: %v", newQce)
    43  			newQceProcessed := false
    44  			for existingQceIdx, existingQce := range qcc.entries {
    45  				if newQce.ID == existingQce.ID {
    46  					newQceProcessed = true
    47  					if newQce.Timestamp > existingQce.Timestamp {
    48  						qcc.entries[existingQceIdx] = newQce
    49  					} else if newQce.Timestamp == existingQce.Timestamp {
    50  						if newQce.Priority != existingQce.Priority {
    51  							qcc.entries[existingQceIdx] = newQce
    52  						}
    53  					}
    54  				}
    55  			}
    56  			if !newQceProcessed {
    57  				for existingQceIdx, existingQce := range qcc.entries {
    58  					if newQce.Priority == existingQce.Priority {
    59  						qcc.entries[existingQceIdx] = newQce
    60  						newQceProcessed = true
    61  					}
    62  				}
    63  			}
    64  			if !newQceProcessed {
    65  				qcc.entries = append(qcc.entries, newQce)
    66  			}
    67  		}
    68  		sort.Slice(qcc.entries, func(idx1, idx2 int) bool {
    69  			return qcc.entries[idx1].Priority < qcc.entries[idx2].Priority
    70  		})
    71  		qcc.prune(qcc.getQueryContextCacheSize(sc))
    72  	}
    73  }
    74  
    75  func (qcc *queryContextCache) prune(size int) {
    76  	if len(qcc.entries) > size {
    77  		qcc.entries = qcc.entries[0:size]
    78  	}
    79  }
    80  
    81  func (qcc *queryContextCache) getQueryContextCacheSize(sc *snowflakeConn) int {
    82  	paramsMutex.Lock()
    83  	sizeStr, ok := sc.cfg.Params[queryContextCacheSizeParamName]
    84  	paramsMutex.Unlock()
    85  	if ok {
    86  		size, err := strconv.Atoi(*sizeStr)
    87  		if err != nil {
    88  			logger.Warnf("cannot parse %v as int as query context cache size: %v", sizeStr, err)
    89  		} else {
    90  			return size
    91  		}
    92  	}
    93  	return defaultQueryContextCacheSize
    94  }