github.com/instill-ai/component@v0.16.0-beta/pkg/operator/main.go (about)

     1  package operator
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/gofrs/uuid"
     8  	"go.uber.org/zap"
     9  
    10  	"github.com/instill-ai/component/pkg/base"
    11  	"github.com/instill-ai/component/pkg/operator/base64/v0"
    12  	"github.com/instill-ai/component/pkg/operator/end/v0"
    13  	"github.com/instill-ai/component/pkg/operator/image/v0"
    14  	"github.com/instill-ai/component/pkg/operator/json/v0"
    15  	"github.com/instill-ai/component/pkg/operator/start/v0"
    16  	"github.com/instill-ai/component/pkg/operator/text/v0"
    17  
    18  	pipelinePB "github.com/instill-ai/protogen-go/vdp/pipeline/v1beta"
    19  )
    20  
    21  var (
    22  	once    sync.Once
    23  	opStore *OperatorStore
    24  )
    25  
    26  // Operator is the derived operator
    27  type OperatorStore struct {
    28  	operatorUIDs   []uuid.UUID
    29  	operatorUIDMap map[uuid.UUID]*operator
    30  	operatorIDMap  map[string]*operator
    31  }
    32  
    33  type operator struct {
    34  	op base.IOperator
    35  }
    36  
    37  // Init initializes the operator
    38  func Init(logger *zap.Logger, usageHandler base.UsageHandler) *OperatorStore {
    39  	once.Do(func() {
    40  		opStore = &OperatorStore{
    41  			operatorUIDMap: map[uuid.UUID]*operator{},
    42  			operatorIDMap:  map[string]*operator{},
    43  		}
    44  		opStore.Import(start.Init(logger, usageHandler)) // deprecated
    45  		opStore.Import(end.Init(logger, usageHandler))   // deprecated
    46  		opStore.Import(base64.Init(logger, usageHandler))
    47  		opStore.Import(json.Init(logger, usageHandler))
    48  		opStore.Import(image.Init(logger, usageHandler))
    49  		opStore.Import(text.Init(logger, usageHandler))
    50  
    51  	})
    52  	return opStore
    53  }
    54  
    55  // Imports imports the operator definitions
    56  func (os *OperatorStore) Import(op base.IOperator) {
    57  	o := &operator{op: op}
    58  	os.operatorUIDMap[op.GetUID()] = o
    59  	os.operatorIDMap[op.GetID()] = o
    60  	os.operatorUIDs = append(os.operatorUIDs, op.GetUID())
    61  }
    62  
    63  func (os *OperatorStore) CreateExecution(defUID uuid.UUID, sysVars map[string]any, task string) (*base.ExecutionWrapper, error) {
    64  	if op, ok := os.operatorUIDMap[defUID]; ok {
    65  		return op.op.CreateExecution(sysVars, task)
    66  	}
    67  	return nil, fmt.Errorf("operator definition not found")
    68  }
    69  
    70  func (os *OperatorStore) GetOperatorDefinitionByUID(defUID uuid.UUID, sysVars map[string]any, component *pipelinePB.OperatorComponent) (*pipelinePB.OperatorDefinition, error) {
    71  	if op, ok := os.operatorUIDMap[defUID]; ok {
    72  		return op.op.GetOperatorDefinition(sysVars, component)
    73  	}
    74  	return nil, fmt.Errorf("operator definition not found")
    75  }
    76  
    77  // Get the operator definition by definition id
    78  func (os *OperatorStore) GetOperatorDefinitionByID(defID string, sysVars map[string]any, component *pipelinePB.OperatorComponent) (*pipelinePB.OperatorDefinition, error) {
    79  	if op, ok := os.operatorIDMap[defID]; ok {
    80  		return op.op.GetOperatorDefinition(sysVars, component)
    81  	}
    82  	return nil, fmt.Errorf("operator definition not found")
    83  }
    84  
    85  // Get the list of operator definitions under this operator
    86  func (os *OperatorStore) ListOperatorDefinitions(sysVars map[string]any, returnTombstone bool) []*pipelinePB.OperatorDefinition {
    87  	defs := []*pipelinePB.OperatorDefinition{}
    88  	for _, uid := range os.operatorUIDs {
    89  		op := os.operatorUIDMap[uid]
    90  		def, err := op.op.GetOperatorDefinition(sysVars, nil)
    91  		if err == nil {
    92  			if !def.Tombstone || returnTombstone {
    93  				defs = append(defs, def)
    94  			}
    95  		}
    96  	}
    97  	return defs
    98  }