<commit_before>// Datasource package contains database/source type related.  A few datasources
// are implemented here (test, csv).  This package also includes
// schema base services (datasource registry).
package datasource

import (
	"fmt"
	"strings"
	"sync"

	u "github.com/araddon/gou"

	"github.com/araddon/qlbridge/schema"
)

var (
	// the global data sources registry mutex
	registryMu sync.RWMutex
	// registry for data sources
	registry = newRegistry()

	// If disableRecover=true, we will not capture/suppress panics
	// Test only feature hopefully
	DisableRecover bool
)

// Register makes a datasource available by the provided @sourceName
// If Register is called twice with the same name or if source is nil, it panics.
//
//  Sources are specific schemas of type csv, elasticsearch, etc containing
//    multiple tables
func Register(sourceName string, source schema.Source) {
	registryMu.Lock()
	defer registryMu.Unlock()
	sourceName = strings.ToLower(sourceName)
	registerNeedsLock(sourceName, source)
}

func registerNeedsLock(sourceName string, source schema.Source) {
	if source == nil {
		panic("qlbridge/datasource: Register Source is nil")
	}

	if _, dupe := registry.sources[sourceName]; dupe {
		panic("qlbridge/datasource: Register called twice for source " + sourceName)
	}
	registry.sources[sourceName] = source
}

// Register makes a datasource available by the provided @sourceName
// If Register is called twice with the same name or if source is nil, it panics.
//
//  Sources are specific schemas of type csv, elasticsearch, etc containing
//    multiple tables
func RegisterSchemaSource(schema, sourceName string, source schema.Source) *schema.Schema {
	sourceName = strings.ToLower(sourceName)
	registryMu.Lock()
	defer registryMu.Unlock()
	registerNeedsLock(sourceName, source)
	s, _ := createSchema(sourceName)
	registry.schemas[sourceName] = s
	return s
}

// DataSourcesRegistry get access to the shared/global
// registry of all datasource implementations
func DataSourcesRegistry() *Registry {
	return registry
}

// Open a datasource, Global open connection function using
//  default schema registry
//
func OpenConn(sourceName, sourceConfig string) (schema.Conn, error) {
	sourcei, ok := registry.sources[sourceName]
	if !ok {
		return nil, fmt.Errorf("datasource: unknown source %q (forgotten import?)", sourceName)
	}
	source, err := sourcei.Open(sourceConfig)
	if err != nil {
		return nil, err
	}
	return source, nil
}

// Our internal map of different types of datasources that are registered
// for our runtime system to use
type Registry struct {
	// Map of source name, each source name is name of db in a specific source
	//   such as elasticsearch, mongo, csv etc
	sources map[string]schema.Source
	schemas map[string]*schema.Schema
	// We need to be able to flatten all tables across all sources into single keyspace
	//tableSources map[string]schema.DataSource
	tables []string
}

func newRegistry() *Registry {
	return &Registry{
		sources: make(map[string]schema.Source),
		schemas: make(map[string]*schema.Schema),
		//tableSources: make(map[string]schema.DataSource),
		tables: make([]string, 0),
	}
}

// Init pre-schema load call any sources that need pre-schema init
func (m *Registry) Init() {
	registryMu.RLock()
	defer registryMu.RUnlock()

	for _, src := range m.sources {
		src.Init()
	}
}

// Get connection for given Database
//
//  @db      database name
//
func (m *Registry) Conn(db string) schema.Conn {

	//u.Debugf("Registry.Conn(db='%v') ", db)
	source := m.Get(strings.ToLower(db))
	if source != nil {
		//u.Debugf("found source: db=%s   %T", db, source)
		conn, err := source.Open(db)
		if err != nil {<commit_msg>fix race in datasource registry<commit_after>// Datasource package contains database/source type related.  A few datasources
// are implemented here (test, csv).  This package also includes
// schema base services (datasource registry).
package datasource

import (
	"fmt"
	"strings"
	"sync"

	u "github.com/araddon/gou"

	"github.com/araddon/qlbridge/schema"
)

var (
	// the global data sources registry mutex
	registryMu sync.RWMutex
	// registry for data sources
	registry = newRegistry()

	// If disableRecover=true, we will not capture/suppress panics
	// Test only feature hopefully
	DisableRecover bool
)

// Register makes a datasource available by the provided @sourceName
// If Register is called twice with the same name or if source is nil, it panics.
//
//  Sources are specific schemas of type csv, elasticsearch, etc containing
//    multiple tables
func Register(sourceName string, source schema.Source) {
	registryMu.Lock()
	defer registryMu.Unlock()
	sourceName = strings.ToLower(sourceName)
	registerNeedsLock(sourceName, source)
}

func registerNeedsLock(sourceName string, source schema.Source) {
	if source == nil {
		panic("qlbridge/datasource: Register Source is nil")
	}

	if _, dupe := registry.sources[sourceName]; dupe {
		panic("qlbridge/datasource: Register called twice for source " + sourceName)
	}
	registry.sources[sourceName] = source
}

// Register makes a datasource available by the provided @sourceName
// If Register is called twice with the same name or if source is nil, it panics.
//
//  Sources are specific schemas of type csv, elasticsearch, etc containing
//    multiple tables
func RegisterSchemaSource(schema, sourceName string, source schema.Source) *schema.Schema {
	sourceName = strings.ToLower(sourceName)
	registryMu.Lock()
	defer registryMu.Unlock()
	registerNeedsLock(sourceName, source)
	s, _ := createSchema(sourceName)
	registry.schemas[sourceName] = s
	return s
}

// DataSourcesRegistry get access to the shared/global
// registry of all datasource implementations
func DataSourcesRegistry() *Registry {
	return registry
}

// Open a datasource, Global open connection function using
//  default schema registry
//
func OpenConn(sourceName, sourceConfig string) (schema.Conn, error) {
	sourcei, ok := registry.sources[sourceName]
	if !ok {
		return nil, fmt.Errorf("datasource: unknown source %q (forgotten import?)", sourceName)
	}
	source, err := sourcei.Open(sourceConfig)
	if err != nil {
		return nil, err
	}
	return source, nil
}

// Our internal map of different types of datasources that are registered
// for our runtime system to use
type Registry struct {
	// Map of source name, each source name is name of db in a specific source
	//   such as elasticsearch, mongo, csv etc
	sources map[string]schema.Source
	schemas map[string]*schema.Schema
	// We need to be able to flatten all tables across all sources into single keyspace
	//tableSources map[string]schema.DataSource
	tables []string
}

func newRegistry() *Registry {
	return &Registry{
		sources: make(map[string]schema.Source),
		schemas: make(map[string]*schema.Schema),
		//tableSources: make(map[string]schema.DataSource),
		tables: make([]string, 0),
	}
}

// Init pre-schema load call any sources that need pre-schema init
func (m *Registry) Init() {
	//registryMu.RLock()
	//defer registryMu.RUnlock()
	// TODO:  this is a race, we need a lock on sources
	for _, src := range m.sources {
		src.Init()
	}
}

// Get connection for given Database
//
//  @db      database name
//
func (m *Registry) Conn(db string) schema.Conn {

	//u.Debugf("Registry.Conn(db='%v') ", db)
	source := m.Get(strings.ToLower(db))
	if source != nil {
		//u.Debugf("found source: db=%s   %T", db, source)
		conn, err := source.Open(db)
		if err != nil {
{'fraction_non_alphanumeric': 0.10479764138300723, 'fraction_numerical': 0.0002680246582685607, 'mean_word_length': 4.161134163208852, 'pattern_counts': {'":': 0, '<': 3, '<?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': '2984971', 'n_tokens_mistral': 2273, 'n_tokens_neox': 2192, 'n_words': 785}