github.com/weaviate/weaviate@v1.24.6/usecases/connstate/manager.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package connstate
    13  
    14  import (
    15  	"context"
    16  	"encoding/json"
    17  	"fmt"
    18  
    19  	"github.com/sirupsen/logrus"
    20  )
    21  
    22  // Manager can save and load a connector's internal state into a remote storage
    23  type Manager struct {
    24  	repo   Repo
    25  	state  json.RawMessage
    26  	logger logrus.FieldLogger
    27  }
    28  
    29  // Repo describes the dependencies of the connector state manager to an
    30  // external storage
    31  type Repo interface {
    32  	Save(ctx context.Context, state json.RawMessage) error
    33  	Load(ctx context.Context) (json.RawMessage, error)
    34  }
    35  
    36  // NewManager for Connector State
    37  func NewManager(repo Repo, logger logrus.FieldLogger) (*Manager, error) {
    38  	m := &Manager{repo: repo, logger: logger}
    39  	if err := m.loadOrInitialize(context.Background()); err != nil {
    40  		return nil, fmt.Errorf("could not load or initialize: %v", err)
    41  	}
    42  
    43  	return m, nil
    44  }
    45  
    46  // GetInitialState is only supposed to be used during initialization of the
    47  // connector.
    48  func (m *Manager) GetInitialState() json.RawMessage {
    49  	return m.state
    50  }
    51  
    52  // SetState form outside (i.e. from the connector)
    53  func (m *Manager) SetState(ctx context.Context, state json.RawMessage) error {
    54  	m.state = state
    55  	return m.save(ctx)
    56  }
    57  
    58  // func (l *etcdSchemaManager) SetStateConnector(stateConnector connector_state.Connector) {
    59  // 	l.connectorStateSetter = stateConnector
    60  // }
    61  
    62  func (m *Manager) loadOrInitialize(ctx context.Context) error {
    63  	state, err := m.repo.Load(ctx)
    64  	if err != nil {
    65  		return fmt.Errorf("could not load connector state: %v", err)
    66  	}
    67  
    68  	if state == nil {
    69  		m.state = json.RawMessage([]byte("{}"))
    70  		return m.save(ctx)
    71  	}
    72  
    73  	m.state = state
    74  	return nil
    75  }
    76  
    77  func (m *Manager) save(ctx context.Context) error {
    78  	m.logger.
    79  		WithField("action", "connector_state_update").
    80  		WithField("configuration_store", "etcd").
    81  		Debug("saving updated connector state to configuration store")
    82  
    83  	err := m.repo.Save(ctx, m.state)
    84  	if err != nil {
    85  		return fmt.Errorf("could not save connector state: %v", err)
    86  	}
    87  
    88  	return nil
    89  }