<commit_before>package storagenode

import (
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/jacksontj/dataman/src/query"
	"github.com/jacksontj/dataman/src/storage_node/datasource"
	"github.com/jacksontj/dataman/src/storage_node/metadata"
	"github.com/rcrowley/go-metrics"
)

// TODO: remove-- and just have as config options
func NewLocalDatasourceInstance(config *DatasourceInstanceConfig, meta *metadata.Meta) (*DatasourceInstance, error) {
	return NewDatasourceInstance(config, NewStaticMetadataStore(meta))
}

// Create a DatasourceInstance with a default MetadataStore (based on the same config as the storagenode)
func NewDatasourceInstanceDefault(config *DatasourceInstanceConfig) (*DatasourceInstance, error) {
	// Create the meta store
	metaStore, err := NewMetadataStore(config)
	if err != nil {
		return nil, err
	}
	return NewDatasourceInstance(config, metaStore)

}

func NewDatasourceInstance(config *DatasourceInstanceConfig, metaStore StorageMetadataStore) (*DatasourceInstance, error) {
	datasourceInstance := &DatasourceInstance{
		Config:    config,
		MetaStore: metaStore,
		syncChan:  make(chan chan error),
		registry:  config.GetRegistry(),
	}

	datasourceInstance.MutableMetaStore, _ = datasourceInstance.MetaStore.(MutableStorageMetadataStore)

	go datasourceInstance.background()

	if err := datasourceInstance.Sync(); err != nil {
		return nil, err
	}

	var err error
	datasourceInstance.Store, err = config.GetStore(datasourceInstance.GetActiveMeta)
	if err != nil {
		return nil, err
	}

	if StoreSchema, ok := datasourceInstance.Store.(datasource.SchemaInterface); ok {
		datasourceInstance.StoreSchema = StoreSchema
	}

	return datasourceInstance, nil
}

type DatasourceInstance struct {
	Config           *DatasourceInstanceConfig
	MetaStore        StorageMetadataStore
	MutableMetaStore MutableStorageMetadataStore

	StoreSchema datasource.SchemaInterface
	Store       datasource.DataInterface

	// All metadata
	meta atomic.Value
	// Only active objects in metadata
	activeMeta atomic.Value

	// TODO: stop mechanism
	// background sync stuff
	syncChan chan chan error

	// TODO: this should be pluggable, presumably in the datasource
	schemaLock sync.Mutex

	registry metrics.Registry
}

// TODO: remove? since we need to do this while holding the lock it seems useless
func (s *DatasourceInstance) Sync() error {
	errChan := make(chan error, 1)
	s.syncChan <- errChan
	return <-errChan
}

func (s *DatasourceInstance) GetActiveMeta() *metadata.Meta {
	return s.activeMeta.Load().(*metadata.Meta)
}

func (s *DatasourceInstance) GetMeta() *metadata.Meta {
	return s.meta.Load().(*metadata.Meta)
}

func (s *DatasourceInstance) background() {
	interval := time.Second // TODO: configurable interval
	ticker := time.NewTicker(interval)

	for {
		select {
		case <-ticker.C: // time based trigger, in case of error etc.
			s.RefreshMeta()
		case retChan := <-s.syncChan: // event based trigger, so we can get stuff to disk ASAP
			err := s.RefreshMeta()
			retChan <- err
			// since we where just triggered, lets reset the interval
			ticker = time.NewTicker(interval)
		}
	}
}

func (s *DatasourceInstance) RefreshMeta() error {
	s.schemaLock.Lock()
	defer s.schemaLock.Unlock()
	return s.refreshMeta()
}

func (s *DatasourceInstance) refreshMeta() (err error) {
	start := time.Now()
	defer func() {
		end := time.Now()
		if err == nil {
			// Last update time
			c := metrics.GetOrRegisterGauge("fetchMeta.success.last", s.registry)
			c.Update(end.Unix())

			t := metrics.GetOrRegisterTimer("fetchMeta.success.time", s.registry)
			t.Update(end.Sub(start))
		} else {
			// Last update time
			c := metrics.GetOrRegisterGauge("fetchMeta.failure.last", s.registry)
			c.Update(end.Unix())

			t := metrics.GetOrRegisterTimer("fetchMeta.failure.time", s.registry)
			t.Update(end.Sub(start))
		}
	}()

	if meta, err := s.MetaStore.GetMeta(); err == nil {
		s.meta.Store(meta)
	} else {
		return err
	}

	// TODO: filter only active things
	meta, err := s.MetaStore.GetMeta()
	if err != nil {
		return err
	}
	// TODO separate function?
	// TODO: better? We could just do this looking elsewhere, but it is simpler (for the plugins primarily)
	// to just get the ones they expect
	// TODO: maybe have a "trim" method on these?
	for key, database := range meta.Databases {
		if database.ProvisionState != metadata.Active {
			delete(meta.Databases, key)
		} else {
			for key, shardInstance := range database.ShardInstances {
				if shardInstance.ProvisionState != metadata.Active {
					delete(database.ShardInstances, key)
				} else {
					for key, collection := range shardInstance.Collections {
						if collection.ProvisionState != metadata.Active {
							delete(shardInstance.Collections, key)
						} else {
							for key, field := range collection.Fields {
								// TODO: need to recurse
								if field.ProvisionState != metadata.Active {
									delete(collection.Fields, key)
								}
							}

							for key, index := range collection.Indexes {
								if index.ProvisionState != metadata.Active {
									delete(collection.Indexes, key)
								}
							}
						}
					}
				}
			}
		}
	}

	for key, field := range meta.Fields {
		if field.ProvisionState != metadata.Active {
			delete(meta.Fields, key)
		}
	}

	for key, collection := range meta.Collections {
		if collection.ProvisionState != metadata.Active {
			delete(meta.Collections, key)
		}
	}

	s.activeMeta.Store(meta)

	// TODO: elsewhere?
	metadata.FieldTypeRegistry.Merge(meta.FieldTypeRegistry)
	return nil
}

// TODO: remove? need some mechanism to override the meta func for schema migrations, imports, etc.
func (s *DatasourceInstance) SetActiveMeta(meta *metadata.Meta) {
	s.activeMeta.Store(meta)
}

// TODO: switch this to the query.Query struct? If not then we should probably support both query formats? Or remove that Query struct
func (s *DatasourceInstance) HandleQuery(q map[query.QueryType]query.QueryArgs) *query.Result {
	return s.HandleQueries([]map[query.QueryType]query.QueryArgs{q})[0]
}

// TODO: we should actually do these in parallel (potentially with some config of *how* parallel)
func (s *DatasourceInstance) HandleQueries(queries []map[query.QueryType]query.QueryArgs) []*query.Result {
	start := time.Now()
	defer func() {
		end := time.Now()
		t := metrics.GetOrRegisterTimer("handleQueries.time", s.registry)
		t.Update(end.Sub(start))
	}()

	results := make([]*query.Result, len(queries))

	// We specifically want to load this once for the batch so we don't have mixed
	// schema information across this batch of queries
	meta := s.GetActiveMeta()

QUERYLOOP:
	for i, queryMap := range queries {
		// We only allow a single method to be defined per item
		if len(queryMap) == 1 {
			for queryType, queryArgs := range queryMap {

				collection, err := meta.GetCollection(queryArgs["db"].(string), queryArgs["shard_instance"].(string), queryArgs["collection"].(string))
				// Verify that the table is within our domain
				if err != nil {
					results[i] = &query.Result{
						Error: err.Error(),
					}
					continue
				}

				// If this is a write operation, do whatever schema validation is necessary<commit_msg>Add some safeguarding around the join<commit_after>package storagenode

import (
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/jacksontj/dataman/src/query"
	"github.com/jacksontj/dataman/src/storage_node/datasource"
	"github.com/jacksontj/dataman/src/storage_node/metadata"
	"github.com/rcrowley/go-metrics"
)

// TODO: remove-- and just have as config options
func NewLocalDatasourceInstance(config *DatasourceInstanceConfig, meta *metadata.Meta) (*DatasourceInstance, error) {
	return NewDatasourceInstance(config, NewStaticMetadataStore(meta))
}

// Create a DatasourceInstance with a default MetadataStore (based on the same config as the storagenode)
func NewDatasourceInstanceDefault(config *DatasourceInstanceConfig) (*DatasourceInstance, error) {
	// Create the meta store
	metaStore, err := NewMetadataStore(config)
	if err != nil {
		return nil, err
	}
	return NewDatasourceInstance(config, metaStore)

}

func NewDatasourceInstance(config *DatasourceInstanceConfig, metaStore StorageMetadataStore) (*DatasourceInstance, error) {
	datasourceInstance := &DatasourceInstance{
		Config:    config,
		MetaStore: metaStore,
		syncChan:  make(chan chan error),
		registry:  config.GetRegistry(),
	}

	datasourceInstance.MutableMetaStore, _ = datasourceInstance.MetaStore.(MutableStorageMetadataStore)

	go datasourceInstance.background()

	if err := datasourceInstance.Sync(); err != nil {
		return nil, err
	}

	var err error
	datasourceInstance.Store, err = config.GetStore(datasourceInstance.GetActiveMeta)
	if err != nil {
		return nil, err
	}

	if StoreSchema, ok := datasourceInstance.Store.(datasource.SchemaInterface); ok {
		datasourceInstance.StoreSchema = StoreSchema
	}

	return datasourceInstance, nil
}

type DatasourceInstance struct {
	Config           *DatasourceInstanceConfig
	MetaStore        StorageMetadataStore
	MutableMetaStore MutableStorageMetadataStore

	StoreSchema datasource.SchemaInterface
	Store       datasource.DataInterface

	// All metadata
	meta atomic.Value
	// Only active objects in metadata
	activeMeta atomic.Value

	// TODO: stop mechanism
	// background sync stuff
	syncChan chan chan error

	// TODO: this should be pluggable, presumably in the datasource
	schemaLock sync.Mutex

	registry metrics.Registry
}

// TODO: remove? since we need to do this while holding the lock it seems useless
func (s *DatasourceInstance) Sync() error {
	errChan := make(chan error, 1)
	s.syncChan <- errChan
	return <-errChan
}

func (s *DatasourceInstance) GetActiveMeta() *metadata.Meta {
	return s.activeMeta.Load().(*metadata.Meta)
}

func (s *DatasourceInstance) GetMeta() *metadata.Meta {
	return s.meta.Load().(*metadata.Meta)
}

func (s *DatasourceInstance) background() {
	interval := time.Second // TODO: configurable interval
	ticker := time.NewTicker(interval)

	for {
		select {
		case <-ticker.C: // time based trigger, in case of error etc.
			s.RefreshMeta()
		case retChan := <-s.syncChan: // event based trigger, so we can get stuff to disk ASAP
			err := s.RefreshMeta()
			retChan <- err
			// since we where just triggered, lets reset the interval
			ticker = time.NewTicker(interval)
		}
	}
}

func (s *DatasourceInstance) RefreshMeta() error {
	s.schemaLock.Lock()
	defer s.schemaLock.Unlock()
	return s.refreshMeta()
}

func (s *DatasourceInstance) refreshMeta() (err error) {
	start := time.Now()
	defer func() {
		end := time.Now()
		if err == nil {
			// Last update time
			c := metrics.GetOrRegisterGauge("fetchMeta.success.last", s.registry)
			c.Update(end.Unix())

			t := metrics.GetOrRegisterTimer("fetchMeta.success.time", s.registry)
			t.Update(end.Sub(start))
		} else {
			// Last update time
			c := metrics.GetOrRegisterGauge("fetchMeta.failure.last", s.registry)
			c.Update(end.Unix())

			t := metrics.GetOrRegisterTimer("fetchMeta.failure.time", s.registry)
			t.Update(end.Sub(start))
		}
	}()

	if meta, err := s.MetaStore.GetMeta(); err == nil {
		s.meta.Store(meta)
	} else {
		return err
	}

	// TODO: filter only active things
	meta, err := s.MetaStore.GetMeta()
	if err != nil {
		return err
	}
	// TODO separate function?
	// TODO: better? We could just do this looking elsewhere, but it is simpler (for the plugins primarily)
	// to just get the ones they expect
	// TODO: maybe have a "trim" method on these?
	for key, database := range meta.Databases {
		if database.ProvisionState != metadata.Active {
			delete(meta.Databases, key)
		} else {
			for key, shardInstance := range database.ShardInstances {
				if shardInstance.ProvisionState != metadata.Active {
					delete(database.ShardInstances, key)
				} else {
					for key, collection := range shardInstance.Collections {
						if collection.ProvisionState != metadata.Active {
							delete(shardInstance.Collections, key)
						} else {
							for key, field := range collection.Fields {
								// TODO: need to recurse
								if field.ProvisionState != metadata.Active {
									delete(collection.Fields, key)
								}
							}

							for key, index := range collection.Indexes {
								if index.ProvisionState != metadata.Active {
									delete(collection.Indexes, key)
								}
							}
						}
					}
				}
			}
		}
	}

	for key, field := range meta.Fields {
		if field.ProvisionState != metadata.Active {
			delete(meta.Fields, key)
		}
	}

	for key, collection := range meta.Collections {
		if collection.ProvisionState != metadata.Active {
			delete(meta.Collections, key)
		}
	}

	s.activeMeta.Store(meta)

	// TODO: elsewhere?
	metadata.FieldTypeRegistry.Merge(meta.FieldTypeRegistry)
	return nil
}

// TODO: remove? need some mechanism to override the meta func for schema migrations, imports, etc.
func (s *DatasourceInstance) SetActiveMeta(meta *metadata.Meta) {
	s.activeMeta.Store(meta)
}

// TODO: switch this to the query.Query struct? If not then we should probably support both query formats? Or remove that Query struct
func (s *DatasourceInstance) HandleQuery(q map[query.QueryType]query.QueryArgs) *query.Result {
	return s.HandleQueries([]map[query.QueryType]query.QueryArgs{q})[0]
}

// TODO: we should actually do these in parallel (potentially with some config of *how* parallel)
func (s *DatasourceInstance) HandleQueries(queries []map[query.QueryType]query.QueryArgs) []*query.Result {
	start := time.Now()
	defer func() {
		end := time.Now()
		t := metrics.GetOrRegisterTimer("handleQueries.time", s.registry)
		t.Update(end.Sub(start))
	}()

	results := make([]*query.Result, len(queries))

	// We specifically want to load this once for the batch so we don't have mixed
	// schema information across this batch of queries
	meta := s.GetActiveMeta()

QUERYLOOP:
	for i, queryMap := range queries {
		// We only allow a single method to be defined per item
		if len(queryMap) == 1 {
			for queryType, queryArgs := range queryMap {

				collection, err := meta.GetCollection(queryArgs["db"].(string), queryArgs["shard_instance"].(string), queryArgs["collection"].(string))
				// Verify that the table is within our domain
				if err != nil {
					results[i] = &query.Result{
						Error: err.Error(),
					}
					continue
				}

				if _, ok := queryArgs["join"]; ok {
					// TODO: remove? We can only do joins at this layer if there is only one shardInstance
					if meta.Databases[queryArgs["db"].(string)].ShardInstances[queryArgs["shard_instance"].(string)].Count != 1 {
						results[i] = &query.Result{Error: "Joins only supported on collections with one shardInstance"}
						continue QUERYLOOP
					}
				}

				// If this is a write operation, do whatever schema validation is necessary
{'fraction_non_alphanumeric': 0.11222837090025645, 'fraction_numerical': 0.00047239843433661763, 'mean_word_length': 3.8270358306188923, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '100841', 'n_tokens_mistral': 5102, 'n_tokens_neox': 4503, 'n_words': 1275}