github.com/siglens/siglens@v0.0.0-20240328180423-f7ce9ae441ed/pkg/segment/query/pqs/pqs.go (about)

     1  /*
     2  Copyright 2023.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package pqs
    18  
    19  import (
    20  	"errors"
    21  	"fmt"
    22  	"sync"
    23  
    24  	blob "github.com/siglens/siglens/pkg/blob"
    25  	"github.com/siglens/siglens/pkg/blob/ssutils"
    26  	"github.com/siglens/siglens/pkg/config"
    27  	"github.com/siglens/siglens/pkg/segment/pqmr"
    28  	"github.com/siglens/siglens/pkg/segment/structs"
    29  	log "github.com/sirupsen/logrus"
    30  )
    31  
    32  type PersistentQueryResults struct {
    33  	segKey       string // segment key
    34  	tableName    string // table name for segKey
    35  	pqid         string // persistent query id
    36  	pqmrFilePath string // raw file path for pqmr file
    37  }
    38  
    39  // segKey -> [qid -> PersistentQueryResults]
    40  var allPersistentQueryResults map[string]map[string]*PersistentQueryResults
    41  var allPersistentQueryResultsLock sync.RWMutex
    42  
    43  func init() {
    44  	allPersistentQueryResults = make(map[string]map[string]*PersistentQueryResults)
    45  	allPersistentQueryResultsLock = sync.RWMutex{}
    46  }
    47  
    48  // base func to add & read from segmeta updates
    49  func AddPersistentQueryResult(segKey string, tableName string, pqid string) {
    50  
    51  	if !config.IsPQSEnabled() {
    52  		return
    53  	}
    54  	ssFile := structs.SegSetFile{
    55  		SegKey:     segKey,
    56  		Identifier: "",
    57  		FileType:   structs.Pqmr,
    58  	}
    59  	fName := ssutils.GetFileNameFromSegSetFile(ssFile)
    60  	allPersistentQueryResultsLock.Lock()
    61  	if _, ok := allPersistentQueryResults[segKey]; !ok {
    62  		allPersistentQueryResults[segKey] = make(map[string]*PersistentQueryResults)
    63  	}
    64  
    65  	allPersistentQueryResults[segKey][pqid] = &PersistentQueryResults{
    66  		segKey:       segKey,
    67  		pqid:         pqid,
    68  		pqmrFilePath: fName,
    69  		tableName:    tableName,
    70  	}
    71  	allPersistentQueryResultsLock.Unlock()
    72  }
    73  
    74  func getPQResults(segKey string, pqid string) (*PersistentQueryResults, error) {
    75  	allPersistentQueryResultsLock.RLock()
    76  	defer allPersistentQueryResultsLock.RUnlock()
    77  	if _, ok := allPersistentQueryResults[segKey]; !ok {
    78  		return nil, errors.New("segKey does not have any persistent query results")
    79  	}
    80  
    81  	var pqResults *PersistentQueryResults
    82  	var ok bool
    83  	if pqResults, ok = allPersistentQueryResults[segKey][pqid]; !ok {
    84  		return nil, errors.New("pqid does not exist for segKey")
    85  	}
    86  	return pqResults, nil
    87  }
    88  
    89  func GetAllPersistentQueryResults(segKey string, pqid string) (*pqmr.SegmentPQMRResults, error) {
    90  
    91  	pqResults, err := getPQResults(segKey, pqid)
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  	fName := fmt.Sprintf("%v/pqmr/%v.pqmr", segKey, pqResults.pqid)
    96  	err = blob.DownloadSegmentBlob(fName, false)
    97  	if err != nil {
    98  		log.Errorf("Failed to download PQMR results! SegKey: %+v, pqid: %+v", segKey, pqResults.pqid)
    99  		return nil, errors.New("failed to download PQMR results")
   100  	}
   101  
   102  	pqmrResults, err := pqmr.ReadPqmr(&fName)
   103  	if err != nil {
   104  		log.Errorf("Failed to read PQMR results! From file %+v. Error: %+v", fName, err)
   105  		return nil, err
   106  	}
   107  	return pqmrResults, nil
   108  }
   109  
   110  // Returns if segKey, pqid combination exists
   111  func DoesSegKeyHavePqidResults(segKey string, pqid string) bool {
   112  
   113  	_, err := getPQResults(segKey, pqid)
   114  	return err == nil
   115  }