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

     1  package connector
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/gofrs/uuid"
     8  	"go.uber.org/zap"
     9  	"google.golang.org/protobuf/types/known/structpb"
    10  
    11  	"github.com/instill-ai/component/pkg/base"
    12  	"github.com/instill-ai/component/pkg/connector/airbyte/v0"
    13  	"github.com/instill-ai/component/pkg/connector/archetypeai/v0"
    14  	"github.com/instill-ai/component/pkg/connector/bigquery/v0"
    15  	"github.com/instill-ai/component/pkg/connector/googlecloudstorage/v0"
    16  	"github.com/instill-ai/component/pkg/connector/googlesearch/v0"
    17  	"github.com/instill-ai/component/pkg/connector/huggingface/v0"
    18  	"github.com/instill-ai/component/pkg/connector/instill/v0"
    19  	"github.com/instill-ai/component/pkg/connector/numbers/v0"
    20  	"github.com/instill-ai/component/pkg/connector/openai/v0"
    21  	"github.com/instill-ai/component/pkg/connector/pinecone/v0"
    22  	"github.com/instill-ai/component/pkg/connector/redis/v0"
    23  	"github.com/instill-ai/component/pkg/connector/restapi/v0"
    24  	"github.com/instill-ai/component/pkg/connector/stabilityai/v0"
    25  	"github.com/instill-ai/component/pkg/connector/website/v0"
    26  
    27  	pipelinePB "github.com/instill-ai/protogen-go/vdp/pipeline/v1beta"
    28  )
    29  
    30  var (
    31  	once     sync.Once
    32  	conStore *ConnectorStore
    33  )
    34  
    35  type ConnectorStore struct {
    36  	connectorUIDs   []uuid.UUID
    37  	connectorUIDMap map[uuid.UUID]*connector
    38  	connectorIDMap  map[string]*connector
    39  }
    40  
    41  type connector struct {
    42  	con base.IConnector
    43  }
    44  
    45  func Init(logger *zap.Logger, usageHandler base.UsageHandler) *ConnectorStore {
    46  	once.Do(func() {
    47  
    48  		conStore = &ConnectorStore{
    49  			connectorUIDMap: map[uuid.UUID]*connector{},
    50  			connectorIDMap:  map[string]*connector{},
    51  		}
    52  
    53  		conStore.Import(stabilityai.Init(logger, usageHandler))
    54  		conStore.Import(instill.Init(logger, usageHandler))
    55  		conStore.Import(huggingface.Init(logger, usageHandler))
    56  		conStore.Import(openai.Init(logger, usageHandler))
    57  		conStore.Import(archetypeai.Init(logger, usageHandler))
    58  		conStore.Import(numbers.Init(logger, usageHandler))
    59  		conStore.Import(airbyte.Init(logger, usageHandler))
    60  		conStore.Import(bigquery.Init(logger, usageHandler))
    61  		conStore.Import(googlecloudstorage.Init(logger, usageHandler))
    62  		conStore.Import(googlesearch.Init(logger, usageHandler))
    63  		conStore.Import(pinecone.Init(logger, usageHandler))
    64  		conStore.Import(redis.Init(logger, usageHandler))
    65  		conStore.Import(restapi.Init(logger, usageHandler))
    66  		conStore.Import(website.Init(logger, usageHandler))
    67  
    68  	})
    69  	return conStore
    70  }
    71  
    72  // Imports imports the connector definitions
    73  func (cs *ConnectorStore) Import(con base.IConnector) {
    74  	c := &connector{con: con}
    75  	cs.connectorUIDMap[con.GetUID()] = c
    76  	cs.connectorIDMap[con.GetID()] = c
    77  	cs.connectorUIDs = append(cs.connectorUIDs, con.GetUID())
    78  }
    79  
    80  func (cs *ConnectorStore) CreateExecution(defUID uuid.UUID, sysVars map[string]any, connection *structpb.Struct, task string) (*base.ExecutionWrapper, error) {
    81  	if con, ok := cs.connectorUIDMap[defUID]; ok {
    82  		return con.con.CreateExecution(sysVars, connection, task)
    83  	}
    84  	return nil, fmt.Errorf("connector definition not found")
    85  }
    86  
    87  func (cs *ConnectorStore) GetConnectorDefinitionByUID(defUID uuid.UUID, sysVars map[string]any, component *pipelinePB.ConnectorComponent) (*pipelinePB.ConnectorDefinition, error) {
    88  	if con, ok := cs.connectorUIDMap[defUID]; ok {
    89  		return con.con.GetConnectorDefinition(sysVars, component)
    90  	}
    91  	return nil, fmt.Errorf("connector definition not found")
    92  }
    93  
    94  // Get the connector definition by definition id
    95  func (cs *ConnectorStore) GetConnectorDefinitionByID(defID string, sysVars map[string]any, component *pipelinePB.ConnectorComponent) (*pipelinePB.ConnectorDefinition, error) {
    96  	if con, ok := cs.connectorIDMap[defID]; ok {
    97  		return con.con.GetConnectorDefinition(sysVars, component)
    98  	}
    99  	return nil, fmt.Errorf("connector definition not found")
   100  }
   101  
   102  // Get the list of connector definitions under this connector
   103  func (cs *ConnectorStore) ListConnectorDefinitions(sysVars map[string]any, returnTombstone bool) []*pipelinePB.ConnectorDefinition {
   104  	defs := []*pipelinePB.ConnectorDefinition{}
   105  	for _, uid := range cs.connectorUIDs {
   106  		con := cs.connectorUIDMap[uid]
   107  		def, err := con.con.GetConnectorDefinition(sysVars, nil)
   108  		if err == nil {
   109  			if !def.Tombstone || returnTombstone {
   110  				defs = append(defs, def)
   111  			}
   112  		}
   113  	}
   114  	return defs
   115  }
   116  
   117  func (cs *ConnectorStore) IsCredentialField(defUID uuid.UUID, target string) (bool, error) {
   118  	if con, ok := cs.connectorUIDMap[defUID]; ok {
   119  		return con.con.IsCredentialField(target), nil
   120  	}
   121  	return false, fmt.Errorf("connector definition not found")
   122  }