<commit_before>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*
Package topo is the module responsible for interacting with the topology
service. It uses one Conn connection to the global topo service (with
possibly another one to a read-only version of the global topo service),
and one to each cell topo service.

It contains the plug-in interfaces Conn, Factory and Version that topo
implementations will use. We support Zookeeper, etcd, consul as real
topo servers, and in-memory, tee as test and utility topo servers.
Implementations are in sub-directories here.

In tests, we do not mock this package. Instead, we just use a memorytopo.

We also support copying data across topo servers (using helpers/copy.go
and the topo2topo cmd binary), and writing to two topo servers at the same
time (using helpers/tee.go). This is to facilitate migrations between
topo servers.

There are two test sub-packages associated with this code:
- test/ contains a test suite that is run against all of our implementations.
  It just performs a bunch of common topo server activities (create, list,
  delete various objects, ...). If a topo implementation passes all these
  tests, it most likely will work as expected in a real deployment.
- topotests/ contains tests that use a memorytopo to test the code in this
  package.
*/
package topo

import (
	"flag"
	"fmt"
	"sync"

	"vitess.io/vitess/go/vt/proto/topodata"

	"context"

	"vitess.io/vitess/go/vt/vterrors"

	"vitess.io/vitess/go/vt/log"
)

const (
	// GlobalCell is the name of the global cell.  It is special
	// as it contains the global topology, and references the other cells.
	GlobalCell = "global"

	// GlobalReadOnlyCell is the name of the global read-only cell
	// connection cell name.
	GlobalReadOnlyCell = "global-read-only"
)

// Filenames for all object types.
const (
	CellInfoFile         = "CellInfo"
	CellsAliasFile       = "CellsAlias"
	KeyspaceFile         = "Keyspace"
	ShardFile            = "Shard"
	VSchemaFile          = "VSchema"
	ShardReplicationFile = "ShardReplication"
	TabletFile           = "Tablet"
	SrvVSchemaFile       = "SrvVSchema"
	SrvKeyspaceFile      = "SrvKeyspace"
	RoutingRulesFile     = "RoutingRules"
	ExternalClustersFile = "ExternalClusters"
)

// Path for all object types.
const (
	CellsPath        = "cells"
	CellsAliasesPath = "cells_aliases"
	KeyspacesPath    = "keyspaces"
	ShardsPath       = "shards"
	TabletsPath      = "tablets"
	MetadataPath     = "metadata"

	ExternalClusterMySQL  = "mysql"
	ExternalClusterVitess = "vitess"
)

// Factory is a factory interface to create Conn objects.
// Topo implementations will provide an implementation for this.
type Factory interface {
	// HasGlobalReadOnlyCell returns true if the global cell
	// has read-only replicas of the topology data. The global topology
	// is usually more expensive to read from / write to, as it is
	// replicated over many cells. Some topology services provide
	// more efficient way to read the data, like Observer servers
	// for Zookeeper. If this returns true, we will maintain
	// two connections for the global topology: the 'global' cell
	// for consistent reads and writes, and the 'global-read-only'
	// cell for reads only.
	HasGlobalReadOnlyCell(serverAddr, root string) bool

	// Create creates a topo.Conn object.
	Create(cell, serverAddr, root string) (Conn, error)
}

// Server is the main topo.Server object. We support two ways of creating one:
// 1. From an implementation, server address, and root path.
//    This uses a plugin mechanism, and we have implementations for
//    etcd, zookeeper and consul.
// 2. Specific implementations may have higher level creation methods
//    (in which case they may provide a more complex Factory).
//    We support memorytopo (for tests and processes that only need an
//    in-memory server), and tee (a helper implementation to transition
//    between one server implementation and another).
type Server struct {
	// globalCell is the main connection to the global topo service.
	// It is created once at construction time.
	globalCell Conn

	// globalReadOnlyCell is the read-only connection to the global
	// topo service. It will be equal to globalCell if we don't distinguish
	// the two.
	globalReadOnlyCell Conn

	// factory allows the creation of connections to various backends.
	// It is set at construction time.
	factory Factory

	// mu protects the following fields.
	mu sync.Mutex
	// cellConns contains clients configured to talk to a list of
	// topo instances representing local topo clusters. These
	// should be accessed with the ConnForCell() method, which
	// will read the list of addresses for that cell from the
	// global cluster and create clients as needed.
	cellConns map[string]cellConn
}

type cellConn struct {
	cellInfo *topodata.CellInfo
	conn     Conn
}

type cellsToAliasesMap struct {
	mu sync.Mutex
	// cellsToAliases contains all cell->alias mappings
	cellsToAliases map[string]string
}

var (
	// topoImplementation is the flag for which implementation to use.
	topoImplementation = flag.String("topo_implementation", "", "the topology implementation to use")

	// topoGlobalServerAddress is the address of the global topology
	// server.
	topoGlobalServerAddress = flag.String("topo_global_server_address", "", "the address of the global topology server")

	// topoGlobalRoot is the root path to use for the global topology
	// server.
	topoGlobalRoot = flag.String("topo_global_root", "", "the path of the global topology data in the global topology server")

	// factories has the factories for the Conn objects.
	factories = make(map[string]Factory)

	cellsAliases = cellsToAliasesMap{
		cellsToAliases: make(map[string]string),
	}
)

// RegisterFactory registers a Factory for an implementation for a Server.
// If an implementation with that name already exists, it log.Fatals out.
// Call this in the 'init' function in your topology implementation module.
func RegisterFactory(name string, factory Factory) {
	if factories[name] != nil {
		log.Fatalf("Duplicate topo.Factory registration for %v", name)
	}
	factories[name] = factory
}

// NewWithFactory creates a new Server based on the given Factory.
// It also opens the global cell connection.
func NewWithFactory(factory Factory, serverAddress, root string) (*Server, error) {
	conn, err := factory.Create(GlobalCell, serverAddress, root)
	if err != nil {
		return nil, err
	}
	conn = NewStatsConn(GlobalCell, conn)

	var connReadOnly Conn
	if factory.HasGlobalReadOnlyCell(serverAddress, root) {
		connReadOnly, err = factory.Create(GlobalReadOnlyCell, serverAddress, root)
		if err != nil {
			return nil, err
		}
		connReadOnly = NewStatsConn(GlobalReadOnlyCell, connReadOnly)
	} else {
		connReadOnly = conn
	}

	return &Server{
		globalCell:         conn,
		globalReadOnlyCell: connReadOnly,
		factory:            factory,
		cellConns:          make(map[string]cellConn),
	}, nil
}

// OpenServer returns a Server using the provided implementation,
// address and root for the global server.
func OpenServer(implementation, serverAddress, root string) (*Server, error) {
	factory, ok := factories[implementation]
	if !ok {
		return nil, NewError(NoImplementation, implementation)
	}
	return NewWithFactory(factory, serverAddress, root)
}

// Open returns a Server using the command line parameter flags
// for implementation, address and root. It log.Exits out if an error occurs.
func Open() *Server {
	if *topoGlobalServerAddress == "" && *topoImplementation != "k8s" {
		log.Exitf("topo_global_server_address must be configured")
	}
	if *topoGlobalRoot == "" {
		log.Exit("topo_global_root must be non-empty")
	}
	ts, err := OpenServer(*topoImplementation, *topoGlobalServerAddress, *topoGlobalRoot)
	if err != nil {
		log.Exitf("Failed to open topo server (%v,%v,%v): %v", *topoImplementation, *topoGlobalServerAddress, *topoGlobalRoot, err)
	}
	return ts
}

// ConnForCell returns a Conn object for the given cell.
// It caches Conn objects from previously requested cells.
func (ts *Server) ConnForCell(ctx context.Context, cell string) (Conn, error) {
	// Global cell is the easy case.
	if cell == GlobalCell {
		return ts.globalCell, nil
	}

	// Fetch cell cluster addresses from the global cluster.
	// We can use the GlobalReadOnlyCell for this call.
	ci, err := ts.GetCellInfo(ctx, cell, false /*strongRead*/)
	if err != nil {
		return nil, err
	}

	// Return a cached client if present.
	ts.mu.Lock()
	defer ts.mu.Unlock()
	cc, ok := ts.cellConns[cell]
	if ok {
		// Client exists in cache.
		// Let's verify that it is the same cell as we are looking for.
		// The cell name can be re-used with a different ServerAddress and/or Root
		// in which case we should get a new connection and update the cache
		if ci.ServerAddress == cc.cellInfo.ServerAddress && ci.Root == cc.cellInfo.Root {
			return cc.conn, nil
		}
		// Close the cached connection, we don't need it anymore
		if cc.conn != nil {
			cc.conn.Close()
		}
	}

	// Connect to the cell topo server, while holding the lock.
	// This ensures only one connection is established at any given time.
	// Create the connection and cache it
	conn, err := ts.factory.Create(cell, ci.ServerAddress, ci.Root)
	switch {
	case err == nil:
		conn = NewStatsConn(cell, conn)
		ts.cellConns[cell] = cellConn{ci, conn}
		return conn, nil
	case IsErrType(err, NoNode):
		err = vterrors.Wrap(err, fmt.Sprintf("failed to create topo connection to %v, %v", ci.ServerAddress, ci.Root))
		return nil, NewError(NoNode, err.Error())
	default:
		return nil, vterrors.Wrap(err, fmt.Sprintf("failed to create topo connection to %v, %v", ci.ServerAddress, ci.Root))
	}
}

// GetAliasByCell returns the alias group this `cell` belongs to, if there's none, it returns the `cell` as alias.
func GetAliasByCell(ctx context.Context, ts *Server, cell string) string {
	cellsAliases.mu.Lock()
	defer cellsAliases.mu.Unlock()
	if region, ok := cellsAliases.cellsToAliases[cell]; ok {
		return region
	}
	if ts != nil {
		// lazily get the region from cell info if `aliases` are available
		cellAliases, err := ts.GetCellsAliases(ctx, false)
		if err != nil {
			// for backward compatibility
			return cell
		}

		for alias, cellsAlias := range cellAliases {
			for _, cellAlias := range cellsAlias.Cells {
				if cellAlias == cell {
					cellsAliases.cellsToAliases[cell] = alias
					return alias
				}
			}
		}
	}
	// for backward compatibility
	return cell
}

// Close will close all connections to underlying topo Server.
// It will nil all member variables, so any further access will panic.
func (ts *Server) Close() {
	ts.globalCell.Close()
	if ts.globalReadOnlyCell != ts.globalCell {
		ts.globalReadOnlyCell.Close()
	}
	ts.globalCell = nil
	ts.globalReadOnlyCell = nil
	ts.mu.Lock()
	defer ts.mu.Unlock()
	for _, cc := range ts.cellConns {
		cc.conn.Close()
	}
	ts.cellConns = make(map[string]cellConn)
}

func (ts *Server) clearCellAliasesCache() {
	cellsAliases.mu.Lock()
	defer cellsAliases.mu.Unlock()
	cellsAliases.cellsToAliases = make(map[string]string)
}

// OpenExternalVitessClusterServer returns the topo server of the external cluster
func (ts *Server) OpenExternalVitessClusterServer(ctx context.Context, clusterName string) (*Server, error) {
	vc, err := ts.GetExternalVitessCluster(ctx, clusterName)
	if err != nil {
		return nil, err
	}
	if vc == nil {
		return nil, fmt.Errorf("no vitess cluster found with name %s", clusterName)
	}
	var externalTopo *Server
	externalTopo, err = OpenServer(vc.TopoConfig.TopoType, vc.TopoConfig.Server, vc.TopoConfig.Root)
	if err != nil {
		return nil, err
	}
	if externalTopo == nil {
		return nil, fmt.Errorf("unable to open external topo for config %s", clusterName)
	}
	return externalTopo, nil
}

// SetReadOnly is initially ONLY implemented by StatsConn and used in ReadOnlyServer
func (ts *Server) SetReadOnly(readOnly bool) error {
	globalCellConn, ok := ts.globalCell.(*StatsConn)
	if !ok {
		return fmt.Errorf("invalid global cell connection type, expected StatsConn but found: %T", ts.globalCell)
	}
	globalCellConn.SetReadOnly(readOnly)

	for _, cc := range ts.cellConns {
		localCellConn, ok := cc.conn.(*StatsConn)
		if !ok {
			return fmt.Errorf("invalid local cell connection type, expected StatsConn but found: %T", cc.conn)
		}
		localCellConn.SetReadOnly(true)
	}

	return nil
}

// IsReadOnly is initially ONLY implemented by StatsConn and used in ReadOnlyServer
func (ts *Server) IsReadOnly() (bool, error) {
	globalCellConn, ok := ts.globalCell.(*StatsConn)
	if !ok {
		return false, fmt.Errorf("invalid global cell connection type, expected StatsConn but found: %T", ts.globalCell)
	}
	if !globalCellConn.IsReadOnly() {
		return false, nil
	}

	for _, cc := range ts.cellConns {
		localCellConn, ok := cc.conn.(*StatsConn)
		if !ok {
			return false, fmt.Errorf("invalid local cell connection type, expected StatsConn but found: %T", cc.conn)
		}
		if !localCellConn.IsReadOnly() {
			return false, nil
		}
	}

	return true, nil
}
<commit_msg>Don't allow expired contexts to retrieve topo connection<commit_after>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*
Package topo is the module responsible for interacting with the topology
service. It uses one Conn connection to the global topo service (with
possibly another one to a read-only version of the global topo service),
and one to each cell topo service.

It contains the plug-in interfaces Conn, Factory and Version that topo
implementations will use. We support Zookeeper, etcd, consul as real
topo servers, and in-memory, tee as test and utility topo servers.
Implementations are in sub-directories here.

In tests, we do not mock this package. Instead, we just use a memorytopo.

We also support copying data across topo servers (using helpers/copy.go
and the topo2topo cmd binary), and writing to two topo servers at the same
time (using helpers/tee.go). This is to facilitate migrations between
topo servers.

There are two test sub-packages associated with this code:
- test/ contains a test suite that is run against all of our implementations.
  It just performs a bunch of common topo server activities (create, list,
  delete various objects, ...). If a topo implementation passes all these
  tests, it most likely will work as expected in a real deployment.
- topotests/ contains tests that use a memorytopo to test the code in this
  package.
*/
package topo

import (
	"context"
	"flag"
	"fmt"
	"sync"

	"vitess.io/vitess/go/vt/proto/topodata"

	"vitess.io/vitess/go/vt/vterrors"

	"vitess.io/vitess/go/vt/log"
)

const (
	// GlobalCell is the name of the global cell.  It is special
	// as it contains the global topology, and references the other cells.
	GlobalCell = "global"

	// GlobalReadOnlyCell is the name of the global read-only cell
	// connection cell name.
	GlobalReadOnlyCell = "global-read-only"
)

// Filenames for all object types.
const (
	CellInfoFile         = "CellInfo"
	CellsAliasFile       = "CellsAlias"
	KeyspaceFile         = "Keyspace"
	ShardFile            = "Shard"
	VSchemaFile          = "VSchema"
	ShardReplicationFile = "ShardReplication"
	TabletFile           = "Tablet"
	SrvVSchemaFile       = "SrvVSchema"
	SrvKeyspaceFile      = "SrvKeyspace"
	RoutingRulesFile     = "RoutingRules"
	ExternalClustersFile = "ExternalClusters"
)

// Path for all object types.
const (
	CellsPath        = "cells"
	CellsAliasesPath = "cells_aliases"
	KeyspacesPath    = "keyspaces"
	ShardsPath       = "shards"
	TabletsPath      = "tablets"
	MetadataPath     = "metadata"

	ExternalClusterMySQL  = "mysql"
	ExternalClusterVitess = "vitess"
)

// Factory is a factory interface to create Conn objects.
// Topo implementations will provide an implementation for this.
type Factory interface {
	// HasGlobalReadOnlyCell returns true if the global cell
	// has read-only replicas of the topology data. The global topology
	// is usually more expensive to read from / write to, as it is
	// replicated over many cells. Some topology services provide
	// more efficient way to read the data, like Observer servers
	// for Zookeeper. If this returns true, we will maintain
	// two connections for the global topology: the 'global' cell
	// for consistent reads and writes, and the 'global-read-only'
	// cell for reads only.
	HasGlobalReadOnlyCell(serverAddr, root string) bool

	// Create creates a topo.Conn object.
	Create(cell, serverAddr, root string) (Conn, error)
}

// Server is the main topo.Server object. We support two ways of creating one:
// 1. From an implementation, server address, and root path.
//    This uses a plugin mechanism, and we have implementations for
//    etcd, zookeeper and consul.
// 2. Specific implementations may have higher level creation methods
//    (in which case they may provide a more complex Factory).
//    We support memorytopo (for tests and processes that only need an
//    in-memory server), and tee (a helper implementation to transition
//    between one server implementation and another).
type Server struct {
	// globalCell is the main connection to the global topo service.
	// It is created once at construction time.
	globalCell Conn

	// globalReadOnlyCell is the read-only connection to the global
	// topo service. It will be equal to globalCell if we don't distinguish
	// the two.
	globalReadOnlyCell Conn

	// factory allows the creation of connections to various backends.
	// It is set at construction time.
	factory Factory

	// mu protects the following fields.
	mu sync.Mutex
	// cellConns contains clients configured to talk to a list of
	// topo instances representing local topo clusters. These
	// should be accessed with the ConnForCell() method, which
	// will read the list of addresses for that cell from the
	// global cluster and create clients as needed.
	cellConns map[string]cellConn
}

type cellConn struct {
	cellInfo *topodata.CellInfo
	conn     Conn
}

type cellsToAliasesMap struct {
	mu sync.Mutex
	// cellsToAliases contains all cell->alias mappings
	cellsToAliases map[string]string
}

var (
	// topoImplementation is the flag for which implementation to use.
	topoImplementation = flag.String("topo_implementation", "", "the topology implementation to use")

	// topoGlobalServerAddress is the address of the global topology
	// server.
	topoGlobalServerAddress = flag.String("topo_global_server_address", "", "the address of the global topology server")

	// topoGlobalRoot is the root path to use for the global topology
	// server.
	topoGlobalRoot = flag.String("topo_global_root", "", "the path of the global topology data in the global topology server")

	// factories has the factories for the Conn objects.
	factories = make(map[string]Factory)

	cellsAliases = cellsToAliasesMap{
		cellsToAliases: make(map[string]string),
	}
)

// RegisterFactory registers a Factory for an implementation for a Server.
// If an implementation with that name already exists, it log.Fatals out.
// Call this in the 'init' function in your topology implementation module.
func RegisterFactory(name string, factory Factory) {
	if factories[name] != nil {
		log.Fatalf("Duplicate topo.Factory registration for %v", name)
	}
	factories[name] = factory
}

// NewWithFactory creates a new Server based on the given Factory.
// It also opens the global cell connection.
func NewWithFactory(factory Factory, serverAddress, root string) (*Server, error) {
	conn, err := factory.Create(GlobalCell, serverAddress, root)
	if err != nil {
		return nil, err
	}
	conn = NewStatsConn(GlobalCell, conn)

	var connReadOnly Conn
	if factory.HasGlobalReadOnlyCell(serverAddress, root) {
		connReadOnly, err = factory.Create(GlobalReadOnlyCell, serverAddress, root)
		if err != nil {
			return nil, err
		}
		connReadOnly = NewStatsConn(GlobalReadOnlyCell, connReadOnly)
	} else {
		connReadOnly = conn
	}

	return &Server{
		globalCell:         conn,
		globalReadOnlyCell: connReadOnly,
		factory:            factory,
		cellConns:          make(map[string]cellConn),
	}, nil
}

// OpenServer returns a Server using the provided implementation,
// address and root for the global server.
func OpenServer(implementation, serverAddress, root string) (*Server, error) {
	factory, ok := factories[implementation]
	if !ok {
		return nil, NewError(NoImplementation, implementation)
	}
	return NewWithFactory(factory, serverAddress, root)
}

// Open returns a Server using the command line parameter flags
// for implementation, address and root. It log.Exits out if an error occurs.
func Open() *Server {
	if *topoGlobalServerAddress == "" && *topoImplementation != "k8s" {
		log.Exitf("topo_global_server_address must be configured")
	}
	if *topoGlobalRoot == "" {
		log.Exit("topo_global_root must be non-empty")
	}
	ts, err := OpenServer(*topoImplementation, *topoGlobalServerAddress, *topoGlobalRoot)
	if err != nil {
		log.Exitf("Failed to open topo server (%v,%v,%v): %v", *topoImplementation, *topoGlobalServerAddress, *topoGlobalRoot, err)
	}
	return ts
}

// ConnForCell returns a Conn object for the given cell.
// It caches Conn objects from previously requested cells.
func (ts *Server) ConnForCell(ctx context.Context, cell string) (Conn, error) {
	// Global cell is the easy case.
	if cell == GlobalCell {
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		return ts.globalCell, nil
	}

	// Fetch cell cluster addresses from the global cluster.
	// We can use the GlobalReadOnlyCell for this call.
	ci, err := ts.GetCellInfo(ctx, cell, false /*strongRead*/)
	if err != nil {
		return nil, err
	}

	// Return a cached client if present.
	ts.mu.Lock()
	defer ts.mu.Unlock()
	cc, ok := ts.cellConns[cell]
	if ok {
		// Client exists in cache.
		// Let's verify that it is the same cell as we are looking for.
		// The cell name can be re-used with a different ServerAddress and/or Root
		// in which case we should get a new connection and update the cache
		if ci.ServerAddress == cc.cellInfo.ServerAddress && ci.Root == cc.cellInfo.Root {
			return cc.conn, nil
		}
		// Close the cached connection, we don't need it anymore
		if cc.conn != nil {
			cc.conn.Close()
		}
	}

	// Connect to the cell topo server, while holding the lock.
	// This ensures only one connection is established at any given time.
	// Create the connection and cache it
	conn, err := ts.factory.Create(cell, ci.ServerAddress, ci.Root)
	switch {
	case err == nil:
		conn = NewStatsConn(cell, conn)
		ts.cellConns[cell] = cellConn{ci, conn}
		return conn, nil
	case IsErrType(err, NoNode):
		err = vterrors.Wrap(err, fmt.Sprintf("failed to create topo connection to %v, %v", ci.ServerAddress, ci.Root))
		return nil, NewError(NoNode, err.Error())
	default:
		return nil, vterrors.Wrap(err, fmt.Sprintf("failed to create topo connection to %v, %v", ci.ServerAddress, ci.Root))
	}
}

// GetAliasByCell returns the alias group this `cell` belongs to, if there's none, it returns the `cell` as alias.
func GetAliasByCell(ctx context.Context, ts *Server, cell string) string {
	cellsAliases.mu.Lock()
	defer cellsAliases.mu.Unlock()
	if region, ok := cellsAliases.cellsToAliases[cell]; ok {
		return region
	}
	if ts != nil {
		// lazily get the region from cell info if `aliases` are available
		cellAliases, err := ts.GetCellsAliases(ctx, false)
		if err != nil {
			// for backward compatibility
			return cell
		}

		for alias, cellsAlias := range cellAliases {
			for _, cellAlias := range cellsAlias.Cells {
				if cellAlias == cell {
					cellsAliases.cellsToAliases[cell] = alias
					return alias
				}
			}
		}
	}
	// for backward compatibility
	return cell
}

// Close will close all connections to underlying topo Server.
// It will nil all member variables, so any further access will panic.
func (ts *Server) Close() {
	ts.globalCell.Close()
	if ts.globalReadOnlyCell != ts.globalCell {
		ts.globalReadOnlyCell.Close()
	}
	ts.globalCell = nil
	ts.globalReadOnlyCell = nil
	ts.mu.Lock()
	defer ts.mu.Unlock()
	for _, cc := range ts.cellConns {
		cc.conn.Close()
	}
	ts.cellConns = make(map[string]cellConn)
}

func (ts *Server) clearCellAliasesCache() {
	cellsAliases.mu.Lock()
	defer cellsAliases.mu.Unlock()
	cellsAliases.cellsToAliases = make(map[string]string)
}

// OpenExternalVitessClusterServer returns the topo server of the external cluster
func (ts *Server) OpenExternalVitessClusterServer(ctx context.Context, clusterName string) (*Server, error) {
	vc, err := ts.GetExternalVitessCluster(ctx, clusterName)
	if err != nil {
		return nil, err
	}
	if vc == nil {
		return nil, fmt.Errorf("no vitess cluster found with name %s", clusterName)
	}
	var externalTopo *Server
	externalTopo, err = OpenServer(vc.TopoConfig.TopoType, vc.TopoConfig.Server, vc.TopoConfig.Root)
	if err != nil {
		return nil, err
	}
	if externalTopo == nil {
		return nil, fmt.Errorf("unable to open external topo for config %s", clusterName)
	}
	return externalTopo, nil
}

// SetReadOnly is initially ONLY implemented by StatsConn and used in ReadOnlyServer
func (ts *Server) SetReadOnly(readOnly bool) error {
	globalCellConn, ok := ts.globalCell.(*StatsConn)
	if !ok {
		return fmt.Errorf("invalid global cell connection type, expected StatsConn but found: %T", ts.globalCell)
	}
	globalCellConn.SetReadOnly(readOnly)

	for _, cc := range ts.cellConns {
		localCellConn, ok := cc.conn.(*StatsConn)
		if !ok {
			return fmt.Errorf("invalid local cell connection type, expected StatsConn but found: %T", cc.conn)
		}
		localCellConn.SetReadOnly(true)
	}

	return nil
}

// IsReadOnly is initially ONLY implemented by StatsConn and used in ReadOnlyServer
func (ts *Server) IsReadOnly() (bool, error) {
	globalCellConn, ok := ts.globalCell.(*StatsConn)
	if !ok {
		return false, fmt.Errorf("invalid global cell connection type, expected StatsConn but found: %T", ts.globalCell)
	}
	if !globalCellConn.IsReadOnly() {
		return false, nil
	}

	for _, cc := range ts.cellConns {
		localCellConn, ok := cc.conn.(*StatsConn)
		if !ok {
			return false, fmt.Errorf("invalid local cell connection type, expected StatsConn but found: %T", cc.conn)
		}
		if !localCellConn.IsReadOnly() {
			return false, nil
		}
	}

	return true, nil
}
<|endoftext|>


<commit_before>package gocbcore

import (
	"encoding/binary"
)

type SnapshotState uint32

func (s SnapshotState) HasInMemory() bool {
	return uint32(s)&1 != 0
}
func (s SnapshotState) HasOnDisk() bool {
	return uint32(s)&2 != 0
}

type StreamObserver interface {
	SnapshotMarker(startSeqNo, endSeqNo uint64, vbId uint16, snapshotType SnapshotState)
	Mutation(seqNo, revNo uint64, flags, expiry, lockTime uint32, cas uint64, datatype uint8, vbId uint16, key, value []byte)
	Deletion(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	Expiration(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	End(vbId uint16, err error)
}

type FailoverEntry struct {
	VbUuid VbUuid
	SeqNo SeqNo
}

type OpenStreamCallback func(error)
type CloseStreamCallback func(error)
type GetFailoverLogCallback func([]FailoverEntry, error)
type GetLastCheckpointCallback func(SeqNo, error)

func (c *Agent) OpenStream(vbId uint16, vbUuid VbUuid, startSeqNo, endSeqNo, snapStartSeqNo, snapEndSeqNo SeqNo, evtHandler StreamObserver, cb OpenStreamCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			// All client errors are handled by the StreamObserver
			cb(err)
			return
		}

		if resp.Magic == ResMagic {
			// This is the response to the open stream request.
			cb(nil)
			return
		}

		// This is one of the stream events
		switch resp.Opcode {
		case CmdDcpSnapshotMarker:
			vbId := uint16(resp.Status)
			newStartSeqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			newEndSeqNo := binary.BigEndian.Uint64(resp.Extras[8:])
			snapshotType := binary.BigEndian.Uint32(resp.Extras[16:])
			evtHandler.SnapshotMarker(newStartSeqNo, newEndSeqNo, vbId, SnapshotState(snapshotType))
		case CmdDcpMutation:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			flags := binary.BigEndian.Uint32(resp.Extras[16:])
			expiry := binary.BigEndian.Uint32(resp.Extras[20:])
			lockTime := binary.BigEndian.Uint32(resp.Extras[24:])
			evtHandler.Mutation(seqNo, revNo, flags, expiry, lockTime, resp.Cas, resp.Datatype, vbId, resp.Key, resp.Value)
		case CmdDcpDeletion:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Deletion(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpExpiration:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Expiration(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpStreamEnd:
			vbId := uint16(resp.Status)
			code := StreamEndStatus(binary.BigEndian.Uint32(resp.Extras[0:]))
			evtHandler.End(vbId, streamEndError{code})
		}
	}

	extraBuf := make([]byte, 48)
	binary.BigEndian.PutUint32(extraBuf[0:], 0)
	binary.BigEndian.PutUint32(extraBuf[4:], 0)
	binary.BigEndian.PutUint64(extraBuf[8:], uint64(startSeqNo))
	binary.BigEndian.PutUint64(extraBuf[16:], uint64(endSeqNo))
	binary.BigEndian.PutUint64(extraBuf[24:], uint64(vbUuid))
	binary.BigEndian.PutUint64(extraBuf[32:], uint64(snapStartSeqNo))
	binary.BigEndian.PutUint64(extraBuf[40:], uint64(snapEndSeqNo))

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpStreamReq,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

func (c *Agent) CloseStream(vbId uint16, cb CloseStreamCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		cb(err)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpCloseStream,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

func (c *Agent) GetFailoverLog(vbId uint16, cb GetFailoverLogCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			cb(nil, err)
			return
		}

		numEntries := len(resp.Value) / 16
		entries := make([]FailoverEntry, numEntries)
		for i := 0; i < numEntries; i++ {
			entries[i] = FailoverEntry{
				VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
				SeqNo: SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
			}
		}
		cb(entries, nil)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpGetFailoverLog,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

func (c *Agent) GetLastCheckpoint(vbId uint16, cb GetLastCheckpointCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			cb(0, err)
			return
		}

		seqNo := SeqNo(binary.BigEndian.Uint64(resp.Value[0:]))
		cb(seqNo, err)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdGetLastCheckpoint,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

<commit_msg>Return the failover log when opening a DCP stream<commit_after>package gocbcore

import (
	"encoding/binary"
)

type SnapshotState uint32

func (s SnapshotState) HasInMemory() bool {
	return uint32(s)&1 != 0
}
func (s SnapshotState) HasOnDisk() bool {
	return uint32(s)&2 != 0
}

type StreamObserver interface {
	SnapshotMarker(startSeqNo, endSeqNo uint64, vbId uint16, snapshotType SnapshotState)
	Mutation(seqNo, revNo uint64, flags, expiry, lockTime uint32, cas uint64, datatype uint8, vbId uint16, key, value []byte)
	Deletion(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	Expiration(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	End(vbId uint16, err error)
}

type FailoverEntry struct {
	VbUuid VbUuid
	SeqNo SeqNo
}

type OpenStreamCallback func([]FailoverEntry, error)
type CloseStreamCallback func(error)
type GetFailoverLogCallback func([]FailoverEntry, error)
type GetLastCheckpointCallback func(SeqNo, error)

func (c *Agent) OpenStream(vbId uint16, vbUuid VbUuid, startSeqNo, endSeqNo, snapStartSeqNo, snapEndSeqNo SeqNo, evtHandler StreamObserver, cb OpenStreamCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			// All client errors are handled by the StreamObserver
			cb(nil, err)
			return
		}

		if resp.Magic == ResMagic {
			// This is the response to the open stream request.

			numEntries := len(resp.Value) / 16
			entries := make([]FailoverEntry, numEntries)
			for i := 0; i < numEntries; i++ {
				entries[i] = FailoverEntry{
					VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
					SeqNo: SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
				}
			}

			cb(entries, nil)
			return
		}

		// This is one of the stream events
		switch resp.Opcode {
		case CmdDcpSnapshotMarker:
			vbId := uint16(resp.Status)
			newStartSeqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			newEndSeqNo := binary.BigEndian.Uint64(resp.Extras[8:])
			snapshotType := binary.BigEndian.Uint32(resp.Extras[16:])
			evtHandler.SnapshotMarker(newStartSeqNo, newEndSeqNo, vbId, SnapshotState(snapshotType))
		case CmdDcpMutation:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			flags := binary.BigEndian.Uint32(resp.Extras[16:])
			expiry := binary.BigEndian.Uint32(resp.Extras[20:])
			lockTime := binary.BigEndian.Uint32(resp.Extras[24:])
			evtHandler.Mutation(seqNo, revNo, flags, expiry, lockTime, resp.Cas, resp.Datatype, vbId, resp.Key, resp.Value)
		case CmdDcpDeletion:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Deletion(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpExpiration:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Expiration(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpStreamEnd:
			vbId := uint16(resp.Status)
			code := StreamEndStatus(binary.BigEndian.Uint32(resp.Extras[0:]))
			evtHandler.End(vbId, streamEndError{code})
		}
	}

	extraBuf := make([]byte, 48)
	binary.BigEndian.PutUint32(extraBuf[0:], 0)
	binary.BigEndian.PutUint32(extraBuf[4:], 0)
	binary.BigEndian.PutUint64(extraBuf[8:], uint64(startSeqNo))
	binary.BigEndian.PutUint64(extraBuf[16:], uint64(endSeqNo))
	binary.BigEndian.PutUint64(extraBuf[24:], uint64(vbUuid))
	binary.BigEndian.PutUint64(extraBuf[32:], uint64(snapStartSeqNo))
	binary.BigEndian.PutUint64(extraBuf[40:], uint64(snapEndSeqNo))

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpStreamReq,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

func (c *Agent) CloseStream(vbId uint16, cb CloseStreamCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		cb(err)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpCloseStream,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

func (c *Agent) GetFailoverLog(vbId uint16, cb GetFailoverLogCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			cb(nil, err)
			return
		}

		numEntries := len(resp.Value) / 16
		entries := make([]FailoverEntry, numEntries)
		for i := 0; i < numEntries; i++ {
			entries[i] = FailoverEntry{
				VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
				SeqNo: SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
			}
		}
		cb(entries, nil)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpGetFailoverLog,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

func (c *Agent) GetLastCheckpoint(vbId uint16, cb GetLastCheckpointCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			cb(0, err)
			return
		}

		seqNo := SeqNo(binary.BigEndian.Uint64(resp.Value[0:]))
		cb(seqNo, err)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdGetLastCheckpoint,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"

	"github.com/jimjibone/goopenzwave"
)

type NodeInfo struct {
	HomeID uint32
	NodeID uint8
	Node   *goopenzwave.Node
	Values map[uint64]*goopenzwave.ValueID
}

var Nodes = make(map[uint8]*NodeInfo)
var initialQueryComplete = make(chan bool)
var sentinitialQueryComplete = false

func main() {
	var controllerPath string
	flag.StringVar(&controllerPath, "controller", "/dev/ttyUSB0", "the path to your controller device")
	flag.Parse()

	fmt.Println("gominozw started with openzwave version:", goopenzwave.GetVersionLongAsString())

	// Setup the OpenZWave library.
	options := goopenzwave.CreateOptions("/usr/local/etc/openzwave/", "", "")
	options.AddOptionLogLevel("SaveLogLevel", goopenzwave.LogLevelNone)
	options.AddOptionLogLevel("QueueLogLevel", goopenzwave.LogLevelNone)
	options.AddOptionInt("DumpTrigger", 4)
	options.AddOptionInt("PollInterval", 500)
	options.AddOptionBool("IntervalBetweenPolls", true)
	options.AddOptionBool("ValidateValueChanges", true)
	options.Lock()

	// Start the library and listen for notifications.
	err := goopenzwave.Start(handleNotification)
	if err != nil {
		fmt.Println("ERROR: failed to start goopenzwave package:", err)
		return
	}

	// Add a driver using the supplied controller path.
	err = goopenzwave.AddDriver(controllerPath)
	if err != nil {
		fmt.Println("ERROR: failed to add goopenzwave driver:", err)
		return
	}

	// Wait here until the initial node query has completed.
	<-initialQueryComplete
	fmt.Println("Finished initial scan, now setting up polling...")

	// Now we will enable polling for a variable. In this simple example, it
	// has been hardwired to poll COMMAND_CLASS_BASIC on each node that
	// supports this setting.
	for _, node := range Nodes {
		// Skip the controller (most likely node 1).
		if node.NodeID == 1 {
			continue
		}

		// For each value for this node, set up polling.
		for i := range node.Values {
			valueid := node.Values[i]

			if valueid.CommandClassID == 0x20 {
				// Enable polling with "intensity" of 2. Though, this is
				// irrelevant with only one value polled.
				valueid.EnablePoll(2)
			}
		}
	}

	fmt.Println("Initial scan complete. Now polling for updates...")
	fmt.Println("Hit ctrl-c to quit")

	// Now wait for the user to quit.
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt)
	<-sig

	// Print out what we know about the network.
	fmt.Println("Nodes:")
	for id, node := range Nodes {
		fmt.Printf("\t%3d: Node: %s\n", id, node.Node)
		fmt.Printf("\t     Values: (%d)\n", len(node.Values))
		for i := range node.Values {
			fmt.Printf("\t\t0x%x: %s\n", i, node.Values[i])
		}
	}

	// All done now finish up.
	goopenzwave.RemoveDriver(controllerPath)
	goopenzwave.Stop()
	goopenzwave.DestroyOptions()
}

func handleNotification(notification *goopenzwave.Notification) {
	fmt.Println("Received notification:", notification)

	// Switch based on notification type.
	switch notification.Type {
	case goopenzwave.NotificationTypeNodeAdded:
		// Create a NodeInfo from the notification then add it to the
		// map.
		nodeinfo := &NodeInfo{
			HomeID: notification.HomeID,
			NodeID: notification.NodeID,
			Node:   goopenzwave.NewNode(notification.HomeID, notification.NodeID),
			Values: make(map[uint64]*goopenzwave.ValueID),
		}
		Nodes[nodeinfo.NodeID] = nodeinfo

	case goopenzwave.NotificationTypeNodeRemoved:
		// Remove the NodeInfo from the map.
		if _, found := Nodes[notification.NodeID]; found {
			delete(Nodes, notification.NodeID)
		}

	case goopenzwave.NotificationTypeValueAdded, goopenzwave.NotificationTypeValueChanged:
		// Find the NodeInfo in the map and add/change the ValueID.
		if node, found := Nodes[notification.NodeID]; found {
			node.Values[notification.ValueID.ID] = notification.ValueID
		}

	case goopenzwave.NotificationTypeValueRemoved:
		// Find the NodeInfo in the map and remove the ValueID.
		if node, found := Nodes[notification.NodeID]; found {
			if _, foundVal := node.Values[notification.ValueID.ID]; foundVal {
				delete(node.Values, notification.ValueID.ID)
			}
		}

	case goopenzwave.NotificationTypeAwakeNodesQueried, goopenzwave.NotificationTypeAllNodesQueried, goopenzwave.NotificationTypeAllNodesQueriedSomeDead:
		// The initial node query has completed.
		if sentinitialQueryComplete == false {
			initialQueryComplete <- true
		}
	}
}
<commit_msg>gominozw support passing config dir as an argument<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"

	"github.com/jimjibone/goopenzwave"
)

type NodeInfo struct {
	HomeID uint32
	NodeID uint8
	Node   *goopenzwave.Node
	Values map[uint64]*goopenzwave.ValueID
}

var Nodes = make(map[uint8]*NodeInfo)
var initialQueryComplete = make(chan bool)
var sentinitialQueryComplete = false

func main() {
	var controllerPath string
	var configPath string
	flag.StringVar(&controllerPath, "controller", "/dev/ttyUSB0", "the path to your controller device")
	flag.StringVar(&configPath, "config", "/usr/local/etc/openzwave/", "the path to open-zwave config")
	flag.Parse()

	fmt.Println("gominozw started with openzwave version:", goopenzwave.GetVersionLongAsString())

	// Setup the OpenZWave library.
	options := goopenzwave.CreateOptions(configPath, "", "")
	options.AddOptionLogLevel("SaveLogLevel", goopenzwave.LogLevelNone)
	options.AddOptionLogLevel("QueueLogLevel", goopenzwave.LogLevelNone)
	options.AddOptionInt("DumpTrigger", 4)
	options.AddOptionInt("PollInterval", 500)
	options.AddOptionBool("IntervalBetweenPolls", true)
	options.AddOptionBool("ValidateValueChanges", true)
	options.Lock()

	// Start the library and listen for notifications.
	err := goopenzwave.Start(handleNotification)
	if err != nil {
		fmt.Println("ERROR: failed to start goopenzwave package:", err)
		return
	}

	// Add a driver using the supplied controller path.
	err = goopenzwave.AddDriver(controllerPath)
	if err != nil {
		fmt.Println("ERROR: failed to add goopenzwave driver:", err)
		return
	}

	// Wait here until the initial node query has completed.
	<-initialQueryComplete
	fmt.Println("Finished initial scan, now setting up polling...")

	// Now we will enable polling for a variable. In this simple example, it
	// has been hardwired to poll COMMAND_CLASS_BASIC on each node that
	// supports this setting.
	for _, node := range Nodes {
		// Skip the controller (most likely node 1).
		if node.NodeID == 1 {
			continue
		}

		// For each value for this node, set up polling.
		for i := range node.Values {
			valueid := node.Values[i]

			if valueid.CommandClassID == 0x20 {
				// Enable polling with "intensity" of 2. Though, this is
				// irrelevant with only one value polled.
				valueid.EnablePoll(2)
			}
		}
	}

	fmt.Println("Initial scan complete. Now polling for updates...")
	fmt.Println("Hit ctrl-c to quit")

	// Now wait for the user to quit.
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt)
	<-sig

	// Print out what we know about the network.
	fmt.Println("Nodes:")
	for id, node := range Nodes {
		fmt.Printf("\t%3d: Node: %s\n", id, node.Node)
		fmt.Printf("\t     Values: (%d)\n", len(node.Values))
		for i := range node.Values {
			fmt.Printf("\t\t0x%x: %s\n", i, node.Values[i])
		}
	}

	// All done now finish up.
	goopenzwave.RemoveDriver(controllerPath)
	goopenzwave.Stop()
	goopenzwave.DestroyOptions()
}

func handleNotification(notification *goopenzwave.Notification) {
	fmt.Println("Received notification:", notification)

	// Switch based on notification type.
	switch notification.Type {
	case goopenzwave.NotificationTypeNodeAdded:
		// Create a NodeInfo from the notification then add it to the
		// map.
		nodeinfo := &NodeInfo{
			HomeID: notification.HomeID,
			NodeID: notification.NodeID,
			Node:   goopenzwave.NewNode(notification.HomeID, notification.NodeID),
			Values: make(map[uint64]*goopenzwave.ValueID),
		}
		Nodes[nodeinfo.NodeID] = nodeinfo

	case goopenzwave.NotificationTypeNodeRemoved:
		// Remove the NodeInfo from the map.
		if _, found := Nodes[notification.NodeID]; found {
			delete(Nodes, notification.NodeID)
		}

	case goopenzwave.NotificationTypeValueAdded, goopenzwave.NotificationTypeValueChanged:
		// Find the NodeInfo in the map and add/change the ValueID.
		if node, found := Nodes[notification.NodeID]; found {
			node.Values[notification.ValueID.ID] = notification.ValueID
		}

	case goopenzwave.NotificationTypeValueRemoved:
		// Find the NodeInfo in the map and remove the ValueID.
		if node, found := Nodes[notification.NodeID]; found {
			if _, foundVal := node.Values[notification.ValueID.ID]; foundVal {
				delete(node.Values, notification.ValueID.ID)
			}
		}

	case goopenzwave.NotificationTypeAwakeNodesQueried, goopenzwave.NotificationTypeAllNodesQueried, goopenzwave.NotificationTypeAllNodesQueriedSomeDead:
		// The initial node query has completed.
		if sentinitialQueryComplete == false {
			initialQueryComplete <- true
		}
	}
}
<|endoftext|>


<commit_before>package flect

var pluralRules = []rule{}

// AddPlural adds a rule that will replace the given suffix with the replacement suffix.
func AddPlural(suffix string, repl string) {
	pluralMoot.Lock()
	defer pluralMoot.Unlock()
	pluralRules = append(pluralRules, rule{
		suffix: suffix,
		fn: func(s string) string {
			s = s[:len(s)-len(suffix)]
			return s + repl
		},
	})

	pluralRules = append(pluralRules, rule{
		suffix: repl,
		fn:     noop,
	})
}

var singleToPlural = map[string]string{
	"aircraft":    "aircraft",
	"alias":       "aliases",
	"alumna":      "alumnae",
	"alumnus":     "alumni",
	"analysis":    "analyses",
	"antenna":     "antennas",
	"antithesis":  "antitheses",
	"apex":        "apexes",
	"appendix":    "appendices",
	"axis":        "axes",
	"bacillus":    "bacilli",
	"bacterium":   "bacteria",
	"basis":       "bases",
	"beau":        "beaus",
	"bison":       "bison",
	"bureau":      "bureaus",
	"bus":         "buses",
	"campus":      "campuses",
	"caucus":      "caucuses",
	"child":       "children",
	"château":     "châteaux",
	"circus":      "circuses",
	"codex":       "codices",
	"concerto":    "concertos",
	"corpus":      "corpora",
	"crisis":      "crises",
	"curriculum":  "curriculums",
	"datum":       "data",
	"dear":        "dear",
	"deer":        "deer",
	"diagnosis":   "diagnoses",
	"die":         "dice",
	"dwarf":       "dwarves",
	"ellipsis":    "ellipses",
	"equipment":   "equipment",
	"erratum":     "errata",
	"faux pas":    "faux pas",
	"fez":         "fezzes",
	"fish":        "fish",
	"focus":       "foci",
	"foo":         "foos",
	"foot":        "feet",
	"formula":     "formulas",
	"fungus":      "fungi",
	"genus":       "genera",
	"goose":       "geese",
	"graffito":    "graffiti",
	"grouse":      "grouse",
	"half":        "halves",
	"halo":        "halos",
	"hoof":        "hooves",
	"human":       "humans",
	"hypothesis":  "hypotheses",
	"index":       "indices",
	"information": "information",
	"jeans":       "jeans",
	"larva":       "larvae",
	"libretto":    "librettos",
	"loaf":        "loaves",
	"locus":       "loci",
	"louse":       "lice",
	"matrix":      "matrices",
	"minutia":     "minutiae",
	"money":       "money",
	"moose":       "moose",
	"mouse":       "mice",
	"nebula":      "nebulae",
	"news":        "news",
	"nucleus":     "nuclei",
	"oasis":       "oases",
	"octopus":     "octopi",
	"offspring":   "offspring",
	"opus":        "opera",
	"ovum":        "ova",
	"ox":          "oxen",
	"parenthesis": "parentheses",
	"phenomenon":  "phenomena",
	"photo":       "photos",
	"phylum":      "phyla",
	"piano":       "pianos",
	"plus":        "pluses",
	"police":      "police",
	"prognosis":   "prognoses",
	"prometheus":  "prometheuses",
	"quiz":        "quizzes",
	"radius":      "radiuses",
	"referendum":  "referendums",
	"ress":        "resses",
	"rice":        "rice",
	"salmon":      "salmon",
	"series":      "series",
	"sheep":       "sheep",
	"shoe":        "shoes",
	"shrimp":      "shrimp",
	"species":     "species",
	"stimulus":    "stimuli",
	"stratum":     "strata",
	"swine":       "swine",
	"syllabus":    "syllabi",
	"symposium":   "symposiums",
	"synopsis":    "synopses",
	"tableau":     "tableaus",
	"testis":      "testes",
	"thesis":      "theses",
	"thief":       "thieves",
	"tooth":       "teeth",
	"trout":       "trout",
	"tuna":        "tuna",
	"vertebra":    "vertebrae",
	"vertix":      "vertices",
	"vertex":      "vertices",
	"vita":        "vitae",
	"vortex":      "vortices",
	"wharf":       "wharves",
	"wife":        "wives",
	"woman":       "women",
	"wolf":        "wolves",
	"you":         "you",
}

var pluralToSingle = map[string]string{}

func init() {
	for k, v := range singleToPlural {
		pluralToSingle[v] = k
	}
}

type singularToPluralSuffix struct {
	singular string
	plural   string
}

var singularToPluralSuffixList = []singularToPluralSuffix{
	{"iterion", "iteria"},
	{"campus", "campuses"},
	{"genera", "genus"},
	{"person", "people"},
	{"phylum", "phyla"},
	{"randum", "randa"},
	{"actus", "acti"},
	{"adium", "adia"},
	{"alias", "aliases"},
	{"basis", "basis"},
	{"child", "children"},
	{"chive", "chives"},
	{"focus", "foci"},
	{"hello", "hellos"},
	{"jeans", "jeans"},
	{"louse", "lice"},
	{"mouse", "mice"},
	{"movie", "movies"},
	{"oasis", "oasis"},
	{"atum", "ata"},
	{"atus", "atuses"},
	{"base", "bases"},
	{"cess", "cesses"},
	{"dium", "diums"},
	{"eses", "esis"},
	{"half", "halves"},
	{"hive", "hives"},
	{"iano", "ianos"},
	{"irus", "iri"},
	{"isis", "ises"},
	{"leus", "li"},
	{"mnus", "mni"},
	{"move", "moves"},
	{"news", "news"},
	{"odex", "odice"},
	{"oose", "eese"},
	{"ouse", "ouses"},
	{"ovum", "ova"},
	{"rion", "ria"},
	{"shoe", "shoes"},
	{"stis", "stes"},
	{"tive", "tives"},
	{"wife", "wives"},
	{"afe", "aves"},
	{"bfe", "bves"},
	{"box", "boxes"},
	{"cfe", "cves"},
	{"dfe", "dves"},
	{"dge", "dges"},
	{"efe", "eves"},
	{"gfe", "gves"},
	{"hfe", "hves"},
	{"ife", "ives"},
	{"itz", "itzes"},
	{"ium", "ia"},
	{"ize", "izes"},
	{"jfe", "jves"},
	{"kfe", "kves"},
	{"man", "men"},
	{"mfe", "mves"},
	{"nfe", "nves"},
	{"nna", "nnas"},
	{"oaf", "oaves"},
	{"oci", "ocus"},
	{"ode", "odes"},
	{"ofe", "oves"},
	{"oot", "eet"},
	{"pfe", "pves"},
	{"pse", "psis"},
	{"qfe", "qves"},
	{"quy", "quies"},
	{"rfe", "rves"},
	{"sfe", "sves"},
	{"tfe", "tves"},
	{"tum", "ta"},
	{"tus", "tuses"},
	{"ufe", "uves"},
	{"ula", "ulae"},
	{"ula", "ulas"},
	{"uli", "ulus"},
	{"use", "uses"},
	{"uss", "usses"},
	{"vfe", "vves"},
	{"wfe", "wves"},
	{"xfe", "xves"},
	{"yfe", "yves"},
	{"you", "you"},
	{"zfe", "zves"},
	{"by", "bies"},
	{"ch", "ches"},
	{"cy", "cies"},
	{"dy", "dies"},
	{"ex", "ices"},
	{"fy", "fies"},
	{"gy", "gies"},
	{"hy", "hies"},
	{"io", "ios"},
	{"jy", "jies"},
	{"ky", "kies"},
	{"lf", "lves"},
	{"ly", "lies"},
	{"my", "mies"},
	{"ny", "nies"},
	{"ox", "oxen"},
	{"py", "pies"},
	{"qy", "qies"},
	{"rf", "rves"},
	{"ry", "ries"},
	{"sh", "shes"},
	{"ss", "sses"},
	{"sy", "sies"},
	{"ty", "ties"},
	{"tz", "tzes"},
	{"va", "vae"},
	{"vy", "vies"},
	{"wy", "wies"},
	{"xy", "xies"},
	{"zy", "zies"},
	{"zz", "zzes"},
	{"o", "oes"},
	{"x", "xes"},
}

func init() {
	for _, suffix := range singularToPluralSuffixList {
		AddPlural(suffix.singular, suffix.plural)
		AddSingular(suffix.plural, suffix.singular)
	}
}
<commit_msg>adding few more words<commit_after>package flect

var pluralRules = []rule{}

// AddPlural adds a rule that will replace the given suffix with the replacement suffix.
func AddPlural(suffix string, repl string) {
	pluralMoot.Lock()
	defer pluralMoot.Unlock()
	pluralRules = append(pluralRules, rule{
		suffix: suffix,
		fn: func(s string) string {
			s = s[:len(s)-len(suffix)]
			return s + repl
		},
	})

	pluralRules = append(pluralRules, rule{
		suffix: repl,
		fn:     noop,
	})
}

var singleToPlural = map[string]string{
	"aircraft":    "aircraft",
	"alias":       "aliases",
	"alumna":      "alumnae",
	"alumnus":     "alumni",
	"analysis":    "analyses",
	"antenna":     "antennas",
	"antithesis":  "antitheses",
	"apex":        "apexes",
	"appendix":    "appendices",
	"axis":        "axes",
	"bacillus":    "bacilli",
	"bacterium":   "bacteria",
	"basis":       "bases",
	"beau":        "beaus",
	"bison":       "bison",
	"bureau":      "bureaus",
	"bus":         "buses",
	"campus":      "campuses",
	"caucus":      "caucuses",
	"child":       "children",
	"château":     "châteaux",
	"circus":      "circuses",
	"codex":       "codices",
	"concerto":    "concertos",
	"corpus":      "corpora",
	"crisis":      "crises",
	"curriculum":  "curriculums",
	"datum":       "data",
	"dear":        "dear",
	"deer":        "deer",
	"diagnosis":   "diagnoses",
	"die":         "dice",
	"dwarf":       "dwarves",
	"ellipsis":    "ellipses",
	"equipment":   "equipment",
	"erratum":     "errata",
	"faux pas":    "faux pas",
	"fez":         "fezzes",
	"fish":        "fish",
	"focus":       "foci",
	"foo":         "foos",
	"foot":        "feet",
	"formula":     "formulas",
	"fungus":      "fungi",
	"genus":       "genera",
	"goose":       "geese",
	"graffito":    "graffiti",
	"grouse":      "grouse",
	"half":        "halves",
	"halo":        "halos",
	"hoof":        "hooves",
	"human":       "humans",
	"hypothesis":  "hypotheses",
	"index":       "indices",
	"information": "information",
	"jeans":       "jeans",
	"larva":       "larvae",
	"libretto":    "librettos",
	"loaf":        "loaves",
	"locus":       "loci",
	"louse":       "lice",
	"matrix":      "matrices",
	"minutia":     "minutiae",
	"money":       "money",
	"moose":       "moose",
	"mouse":       "mice",
	"nebula":      "nebulae",
	"news":        "news",
	"nucleus":     "nuclei",
	"oasis":       "oases",
	"octopus":     "octopi",
	"offspring":   "offspring",
	"opus":        "opera",
	"ovum":        "ova",
	"ox":          "oxen",
	"parenthesis": "parentheses",
	"phenomenon":  "phenomena",
	"photo":       "photos",
	"phylum":      "phyla",
	"piano":       "pianos",
	"plus":        "pluses",
	"police":      "police",
	"prognosis":   "prognoses",
	"prometheus":  "prometheuses",
	"quiz":        "quizzes",
	"radius":      "radiuses",
	"referendum":  "referendums",
	"ress":        "resses",
	"rice":        "rice",
	"salmon":      "salmon",
	"sex":         "sexes",
	"series":      "series",
	"sheep":       "sheep",
	"shoe":        "shoes",
	"shrimp":      "shrimp",
	"species":     "species",
	"stimulus":    "stimuli",
	"stratum":     "strata",
	"swine":       "swine",
	"syllabus":    "syllabi",
	"symposium":   "symposiums",
	"synopsis":    "synopses",
	"tableau":     "tableaus",
	"testis":      "testes",
	"thesis":      "theses",
	"thief":       "thieves",
	"tooth":       "teeth",
	"trout":       "trout",
	"tuna":        "tuna",
	"vertebra":    "vertebrae",
	"vertix":      "vertices",
	"vertex":      "vertices",
	"vita":        "vitae",
	"vortex":      "vortices",
	"wharf":       "wharves",
	"wife":        "wives",
	"woman":       "women",
	"wolf":        "wolves",
	"you":         "you",
}

var pluralToSingle = map[string]string{}

func init() {
	for k, v := range singleToPlural {
		pluralToSingle[v] = k
	}
}

type singularToPluralSuffix struct {
	singular string
	plural   string
}

var singularToPluralSuffixList = []singularToPluralSuffix{
	{"iterion", "iteria"},
	{"campus", "campuses"},
	{"genera", "genus"},
	{"person", "people"},
	{"phylum", "phyla"},
	{"randum", "randa"},
	{"actus", "acti"},
	{"adium", "adia"},
	{"alias", "aliases"},
	{"basis", "basis"},
	{"child", "children"},
	{"chive", "chives"},
	{"focus", "foci"},
	{"hello", "hellos"},
	{"jeans", "jeans"},
	{"louse", "lice"},
	{"mouse", "mice"},
	{"movie", "movies"},
	{"oasis", "oasis"},
	{"atum", "ata"},
	{"atus", "atuses"},
	{"base", "bases"},
	{"cess", "cesses"},
	{"dium", "diums"},
	{"eses", "esis"},
	{"half", "halves"},
	{"hive", "hives"},
	{"iano", "ianos"},
	{"irus", "iri"},
	{"isis", "ises"},
	{"leus", "li"},
	{"mnus", "mni"},
	{"move", "moves"},
	{"news", "news"},
	{"odex", "odice"},
	{"oose", "eese"},
	{"ouse", "ouses"},
	{"ovum", "ova"},
	{"rion", "ria"},
	{"shoe", "shoes"},
	{"stis", "stes"},
	{"tive", "tives"},
	{"wife", "wives"},
	{"afe", "aves"},
	{"bfe", "bves"},
	{"box", "boxes"},
	{"cfe", "cves"},
	{"dfe", "dves"},
	{"dge", "dges"},
	{"efe", "eves"},
	{"gfe", "gves"},
	{"hfe", "hves"},
	{"ife", "ives"},
	{"itz", "itzes"},
	{"ium", "ia"},
	{"ize", "izes"},
	{"jfe", "jves"},
	{"kfe", "kves"},
	{"man", "men"},
	{"mfe", "mves"},
	{"nfe", "nves"},
	{"nna", "nnas"},
	{"oaf", "oaves"},
	{"oci", "ocus"},
	{"ode", "odes"},
	{"ofe", "oves"},
	{"oot", "eet"},
	{"pfe", "pves"},
	{"pse", "psis"},
	{"qfe", "qves"},
	{"quy", "quies"},
	{"rfe", "rves"},
	{"sfe", "sves"},
	{"tfe", "tves"},
	{"tum", "ta"},
	{"tus", "tuses"},
	{"ufe", "uves"},
	{"ula", "ulae"},
	{"ula", "ulas"},
	{"uli", "ulus"},
	{"use", "uses"},
	{"uss", "usses"},
	{"vfe", "vves"},
	{"wfe", "wves"},
	{"xfe", "xves"},
	{"yfe", "yves"},
	{"you", "you"},
	{"zfe", "zves"},
	{"by", "bies"},
	{"ch", "ches"},
	{"cy", "cies"},
	{"dy", "dies"},
	{"ex", "ices"},
	{"fy", "fies"},
	{"gy", "gies"},
	{"hy", "hies"},
	{"io", "ios"},
	{"jy", "jies"},
	{"ky", "kies"},
	{"lf", "lves"},
	{"ly", "lies"},
	{"my", "mies"},
	{"ny", "nies"},
	{"ox", "oxen"},
	{"py", "pies"},
	{"qy", "qies"},
	{"rf", "rves"},
	{"ry", "ries"},
	{"sh", "shes"},
	{"ss", "sses"},
	{"sy", "sies"},
	{"ty", "ties"},
	{"tz", "tzes"},
	{"va", "vae"},
	{"vy", "vies"},
	{"wy", "wies"},
	{"xy", "xies"},
	{"zy", "zies"},
	{"zz", "zzes"},
	{"o", "oes"},
	{"x", "xes"},
}

func init() {
	for _, suffix := range singularToPluralSuffixList {
		AddPlural(suffix.singular, suffix.plural)
		AddSingular(suffix.plural, suffix.singular)
	}
}
<|endoftext|>


<commit_before>package migration

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"regexp"
	"sort"
	"strconv"

	"github.com/Boostport/migration/parser"
)

// Direction type up/down
type Direction int

// String returns a string representation of the direction
func (d Direction) String() string {
	switch d {
	case Up:
		return "up"
	case Down:
		return "down"
	default:
		return "directionless"
	}
}

// Constants for direction
const (
	Up Direction = iota
	Down
)

var numberPrefixRegex = regexp.MustCompile(`^(\d+).*$`)

// Migration represents a migration, containing statements for migrating up and down.
type Migration struct {
	ID   string
	Up   *parser.ParsedMigration
	Down *parser.ParsedMigration
}

// PlannedMigration is a migration with a direction defined. This allows the driver to
// work out how to apply the migration.
type PlannedMigration struct {
	*Migration
	Direction Direction
}

// Less compares two migrations to determine how they should be ordered.
func (m Migration) Less(other *Migration) bool {
	switch {
	case m.isNumeric() && other.isNumeric() && m.VersionInt() != other.VersionInt():
		return m.VersionInt() < other.VersionInt()
	case m.isNumeric() && !other.isNumeric():
		return true
	case !m.isNumeric() && other.isNumeric():
		return false
	default:
		return m.ID < other.ID
	}
}

func (m Migration) isNumeric() bool {
	return len(m.NumberPrefixMatches()) > 0
}

// NumberPrefixMatches returns a list of string matches
func (m Migration) NumberPrefixMatches() []string {
	return numberPrefixRegex.FindStringSubmatch(m.ID)
}

// VersionInt converts the migration version to an 64-bit integer.
func (m Migration) VersionInt() int64 {
	v := m.NumberPrefixMatches()[1]
	value, err := strconv.ParseInt(v, 10, 64)
	if err != nil {
		panic(fmt.Sprintf("Could not parse %q into int64: %s", v, err))
	}
	return value
}

type byID []*Migration

func (b byID) Len() int           { return len(b) }
func (b byID) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
func (b byID) Less(i, j int) bool { return b[i].Less(b[j]) }

// Migrate runs a migration using a given driver and MigrationSource. The direction defines whether
// the migration is up or down, and max is the maximum number of migrations to apply. If max is set to 0,
// then there is no limit on the number of migrations to apply.
func Migrate(driver Driver, migrations Source, direction Direction, max int) (int, error) {
	count := 0

	m, err := getMigrations(migrations)
	if err != nil {
		return count, err
	}

	appliedMigrations, err := driver.Versions()
	if err != nil {
		return count, err
	}

	migrationsToApply := planMigrations(m, appliedMigrations, direction, max)
	for _, plannedMigration := range migrationsToApply {
		logPrintf("Applying migration (%s) named '%s'...", direction.String(), plannedMigration.ID)

		err = driver.Migrate(plannedMigration)
		if err != nil {
			errorMessage := "Error while running migration " + plannedMigration.ID

			if plannedMigration.Direction == Up {
				errorMessage += " (up)"
			} else {
				errorMessage += " (down)"
			}
			return count, fmt.Errorf(errorMessage+": %s", err)
		}

		logPrintf("Applied migration (%s) named '%s'", direction.String(), plannedMigration.ID)
		count++
	}

	err = driver.Close()
	return count, err
}

func getMigrations(migrations Source) ([]*Migration, error) {
	var m []*Migration
	tempMigrations := map[string]*Migration{}

	files, err := migrations.ListMigrationFiles()
	if err != nil {
		return m, err
	}

	regex := regexp.MustCompile(`(\d*_.*)\.(up|down)\..*`)

	for _, file := range files {
		matches := regex.FindStringSubmatch(file)

		if len(matches) > 0 && file == matches[0] {
			id := matches[1]
			direction := matches[2]

			if _, ok := tempMigrations[id]; !ok {
				tempMigrations[id] = &Migration{
					ID: id,
				}
			}

			reader, err := migrations.GetMigrationFile(file)
			if err != nil {
				return m, fmt.Errorf("Error getting migrations: %s", err)
			}

			contents, err := ioutil.ReadAll(reader)
			if err != nil {
				return m, fmt.Errorf("Error getting migration content: %s", err)
			}

			parsed, err := parser.Parse(bytes.NewReader(contents))
			if err != nil {
				return m, fmt.Errorf("Error parsing migration %s: %s", id, err)
			}

			if direction == "up" {
				tempMigrations[id].Up = parsed
			} else {
				tempMigrations[id].Down = parsed
			}
		}
	}

	for _, migration := range tempMigrations {
		m = append(m, migration)
	}

	sort.Sort(byID(m))

	return m, nil
}

func planMigrations(migrations []*Migration, appliedMigrations []string, direction Direction, max int) []*PlannedMigration {
	var applied []*Migration

	for _, appliedMigration := range appliedMigrations {
		applied = append(applied, &Migration{
			ID: appliedMigration,
		})
	}

	sort.Sort(byID(applied))

	// Get last migration that was run
	record := &Migration{}

	if len(applied) > 0 {
		record = applied[len(applied)-1]
	}

	var result []*PlannedMigration

	// Add missing migrations up to the last run migration.
	// This can happen for example when merges happened.
	if len(applied) > 0 {
		result = append(result, toCatchup(migrations, applied, record)...)
	}

	// Figure out which migrations to apply
	toApply := toApply(migrations, record.ID, direction)
	toApplyCount := len(toApply)

	if max > 0 && max < toApplyCount {
		toApplyCount = max
	}

	for _, v := range toApply[0:toApplyCount] {
		result = append(result, &PlannedMigration{
			Migration: v,
			Direction: direction,
		})
	}

	return result
}

// Filter a slice of migrations into ones that should be applied.
func toApply(migrations []*Migration, current string, direction Direction) []*Migration {
	var index = -1

	if current != "" {
		for index < len(migrations)-1 {
			index++
			if migrations[index].ID == current {
				break
			}
		}
	}

	if direction == Up {
		return migrations[index+1:]
	} else if direction == Down {
		if index == -1 {
			return []*Migration{}
		}

		// Add in reverse order
		toApply := make([]*Migration, index+1)
		for i := 0; i < index+1; i++ {
			toApply[index-i] = migrations[i]
		}
		return toApply
	}

	panic("Not possible")
}

// Get migrations that we need to apply regardless of whether the direction is up or down. This is
// because there may be migration "holes" due to merges.
func toCatchup(migrations, existingMigrations []*Migration, lastRun *Migration) []*PlannedMigration {
	var missing []*PlannedMigration

	for _, migration := range migrations {
		found := false

		for _, existing := range existingMigrations {
			if existing.ID == migration.ID {
				found = true
				break
			}
		}

		if !found && migration.Less(lastRun) {
			missing = append(missing, &PlannedMigration{Migration: migration, Direction: Up})
		}
	}

	return missing
}
<commit_msg>Replace ioutil with io equivalents<commit_after>package migration

import (
	"bytes"
	"fmt"
	"io"
	"regexp"
	"sort"
	"strconv"

	"github.com/Boostport/migration/parser"
)

// Direction type up/down
type Direction int

// String returns a string representation of the direction
func (d Direction) String() string {
	switch d {
	case Up:
		return "up"
	case Down:
		return "down"
	default:
		return "directionless"
	}
}

// Constants for direction
const (
	Up Direction = iota
	Down
)

var numberPrefixRegex = regexp.MustCompile(`^(\d+).*$`)

// Migration represents a migration, containing statements for migrating up and down.
type Migration struct {
	ID   string
	Up   *parser.ParsedMigration
	Down *parser.ParsedMigration
}

// PlannedMigration is a migration with a direction defined. This allows the driver to
// work out how to apply the migration.
type PlannedMigration struct {
	*Migration
	Direction Direction
}

// Less compares two migrations to determine how they should be ordered.
func (m Migration) Less(other *Migration) bool {
	switch {
	case m.isNumeric() && other.isNumeric() && m.VersionInt() != other.VersionInt():
		return m.VersionInt() < other.VersionInt()
	case m.isNumeric() && !other.isNumeric():
		return true
	case !m.isNumeric() && other.isNumeric():
		return false
	default:
		return m.ID < other.ID
	}
}

func (m Migration) isNumeric() bool {
	return len(m.NumberPrefixMatches()) > 0
}

// NumberPrefixMatches returns a list of string matches
func (m Migration) NumberPrefixMatches() []string {
	return numberPrefixRegex.FindStringSubmatch(m.ID)
}

// VersionInt converts the migration version to an 64-bit integer.
func (m Migration) VersionInt() int64 {
	v := m.NumberPrefixMatches()[1]
	value, err := strconv.ParseInt(v, 10, 64)
	if err != nil {
		panic(fmt.Sprintf("Could not parse %q into int64: %s", v, err))
	}
	return value
}

type byID []*Migration

func (b byID) Len() int           { return len(b) }
func (b byID) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
func (b byID) Less(i, j int) bool { return b[i].Less(b[j]) }

// Migrate runs a migration using a given driver and MigrationSource. The direction defines whether
// the migration is up or down, and max is the maximum number of migrations to apply. If max is set to 0,
// then there is no limit on the number of migrations to apply.
func Migrate(driver Driver, migrations Source, direction Direction, max int) (int, error) {
	count := 0

	m, err := getMigrations(migrations)
	if err != nil {
		return count, err
	}

	appliedMigrations, err := driver.Versions()
	if err != nil {
		return count, err
	}

	migrationsToApply := planMigrations(m, appliedMigrations, direction, max)
	for _, plannedMigration := range migrationsToApply {
		logPrintf("Applying migration (%s) named '%s'...", direction.String(), plannedMigration.ID)

		err = driver.Migrate(plannedMigration)
		if err != nil {
			errorMessage := "Error while running migration " + plannedMigration.ID

			if plannedMigration.Direction == Up {
				errorMessage += " (up)"
			} else {
				errorMessage += " (down)"
			}
			return count, fmt.Errorf(errorMessage+": %s", err)
		}

		logPrintf("Applied migration (%s) named '%s'", direction.String(), plannedMigration.ID)
		count++
	}

	err = driver.Close()
	return count, err
}

func getMigrations(migrations Source) ([]*Migration, error) {
	var m []*Migration
	tempMigrations := map[string]*Migration{}

	files, err := migrations.ListMigrationFiles()
	if err != nil {
		return m, err
	}

	regex := regexp.MustCompile(`(\d*_.*)\.(up|down)\..*`)

	for _, file := range files {
		matches := regex.FindStringSubmatch(file)

		if len(matches) > 0 && file == matches[0] {
			id := matches[1]
			direction := matches[2]

			if _, ok := tempMigrations[id]; !ok {
				tempMigrations[id] = &Migration{
					ID: id,
				}
			}

			reader, err := migrations.GetMigrationFile(file)
			if err != nil {
				return m, fmt.Errorf("Error getting migrations: %s", err)
			}

			contents, err := io.ReadAll(reader)
			if err != nil {
				return m, fmt.Errorf("Error getting migration content: %s", err)
			}

			parsed, err := parser.Parse(bytes.NewReader(contents))
			if err != nil {
				return m, fmt.Errorf("Error parsing migration %s: %s", id, err)
			}

			if direction == "up" {
				tempMigrations[id].Up = parsed
			} else {
				tempMigrations[id].Down = parsed
			}
		}
	}

	for _, migration := range tempMigrations {
		m = append(m, migration)
	}

	sort.Sort(byID(m))

	return m, nil
}

func planMigrations(migrations []*Migration, appliedMigrations []string, direction Direction, max int) []*PlannedMigration {
	var applied []*Migration

	for _, appliedMigration := range appliedMigrations {
		applied = append(applied, &Migration{
			ID: appliedMigration,
		})
	}

	sort.Sort(byID(applied))

	// Get last migration that was run
	record := &Migration{}

	if len(applied) > 0 {
		record = applied[len(applied)-1]
	}

	var result []*PlannedMigration

	// Add missing migrations up to the last run migration.
	// This can happen for example when merges happened.
	if len(applied) > 0 {
		result = append(result, toCatchup(migrations, applied, record)...)
	}

	// Figure out which migrations to apply
	toApply := toApply(migrations, record.ID, direction)
	toApplyCount := len(toApply)

	if max > 0 && max < toApplyCount {
		toApplyCount = max
	}

	for _, v := range toApply[0:toApplyCount] {
		result = append(result, &PlannedMigration{
			Migration: v,
			Direction: direction,
		})
	}

	return result
}

// Filter a slice of migrations into ones that should be applied.
func toApply(migrations []*Migration, current string, direction Direction) []*Migration {
	var index = -1

	if current != "" {
		for index < len(migrations)-1 {
			index++
			if migrations[index].ID == current {
				break
			}
		}
	}

	if direction == Up {
		return migrations[index+1:]
	} else if direction == Down {
		if index == -1 {
			return []*Migration{}
		}

		// Add in reverse order
		toApply := make([]*Migration, index+1)
		for i := 0; i < index+1; i++ {
			toApply[index-i] = migrations[i]
		}
		return toApply
	}

	panic("Not possible")
}

// Get migrations that we need to apply regardless of whether the direction is up or down. This is
// because there may be migration "holes" due to merges.
func toCatchup(migrations, existingMigrations []*Migration, lastRun *Migration) []*PlannedMigration {
	var missing []*PlannedMigration

	for _, migration := range migrations {
		found := false

		for _, existing := range existingMigrations {
			if existing.ID == migration.ID {
				found = true
				break
			}
		}

		if !found && migration.Less(lastRun) {
			missing = append(missing, &PlannedMigration{Migration: migration, Direction: Up})
		}
	}

	return missing
}
<|endoftext|>


<commit_before>package migration

import (
	"database/sql"
	"fmt"
)

var ef = fmt.Errorf

// GetVersion is any function that can retrieve the migration version of a
// particular database. It is exposed in case a client wants to override the
// default behavior of this package. (For example, by using the `user_version`
// PRAGMA in SQLite.)
//
// The DefaultGetVersion function provided with this package creates its own
// table with a single column and a single row.
//
// The version returned should be equivalent to the number of migrations
// applied to this database. It should be 0 if no migrations have been applied
// yet.
//
// If an error is returned, the migration automatically fails.
type GetVersion func(*sql.Tx) (int, error)

// The default way to get the version from a database. If the database has
// had no migrations performed, then it creates a table with a single row and
// a single column storing the version as 0. It then returns 0.
//
// If the table exists, then the version stored in the table is returned.
var DefaultGetVersion GetVersion = defaultGetVersion

// SetVersion is the dual of GetVersion. It allows the client to define a
// different mechanism for setting the database version than the one used by
// DefaultSetVersion in this package.
//
// If an error is returned, the migration that tried to set the version
// automatically fails.
type SetVersion func(*sql.Tx, int) error

// The default way to set the version of the database. If the database has had
// no migrations performed, then it creates a table with a single row and a
// single column and storing the version given there.
//
// If the table exists, then the existing version is overwritten.
var DefaultSetVersion SetVersion = defaultSetVersion

// Migrator corresponds to a function that updates the database by one version.
// Note that a migration should NOT call Rollback or Commit. Instead, this
// package will call Rollback for you if your migration returns an error. If
// no error is returned, then Commit is called.
type Migrator func(*sql.Tx) error

// Open wraps the Open function from the database/sql function, but performs
// a series of migrations on a database if they haven't been performed already.
//
// Migrations are tracked by a simple versioning scheme. The version of the
// database is the number of migrations that have been performed on it.
// Similarly, the version of your library is the number of migrations that are
// defined.
//
// If Open returns successfully, then the database and your library will have
// the same versions. If there was a problem migrating---or if the database
// version is greater than your library version---then an error is returned.
// Since all migrations are performed in a single transaction, if an error
// occurs, no changes are made to the database. (Assuming you're using a
// relational database that allows modifications to a schema to be rolled back.)
//
// Note that this versioning scheme includes no semantic analysis. It is up to
// client to ensure that once a migration is defined, it never changes.
//
// The details of how the version is stored are opaque to the client, but in
// general, it will add a table to your database called "migration_version"
// with a single column containing a single row.
func Open(driver, dsn string, migrations []Migrator) (*sql.DB, error) {
	return OpenWith(driver, dsn, migrations, nil, nil)
}

// OpenWith is exactly like Open, except it allows the client to specify their
// own versioning scheme. Note that versionGet and versionSet must BOTH be
// nil or BOTH be non-nil. Otherwise, this function panics. This because the
// implementation of one generally relies on the implementation of the other.
//
// If versionGet and versionSet are both set to nil, then the behavior of this
// function is identical to the behavior of Open.
func OpenWith(
	driver, dsn string,
	migrations []Migrator,
	versionGet GetVersion,
	versionSet SetVersion,
) (*sql.DB, error) {
	if (versionGet == nil && versionSet != nil) ||
		(versionGet != nil && versionSet == nil) {
		panic("versionGet/versionSet must both be nil or both be non-nil")
	}
	if versionGet == nil {
		versionGet = DefaultGetVersion
	}
	if versionSet == nil {
		versionSet = DefaultSetVersion
	}

	db, err := sql.Open(driver, dsn)
	if err != nil {
		return nil, err
	}
	m := migration{db, migrations, versionGet, versionSet}
	if err := m.migrate(); err != nil {
		return nil, err
	}
	return db, nil
}

type migration struct {
	*sql.DB
	migrations []Migrator
	getVersion GetVersion
	setVersion SetVersion
}

func (m migration) migrate() error {
	tx, err := m.Begin()
	if err != nil {
		return ef("Could not start transaction: %s", err)
	}
	libVersion := len(m.migrations)
	dbVersion, err := m.getVersion(tx)
	if err != nil {
		return ef("Could not get DB version: %s", err)
	}
	if dbVersion > libVersion {
		return ef("Database version (%d) is greater than library version (%d).",
			dbVersion, libVersion)
	}
	if dbVersion == libVersion {
		return tx.Commit()
	}

	for i := dbVersion; i < libVersion; i++ {
		if err := m.migrations[i](tx); err != nil {
			if err2 := tx.Rollback(); err2 != nil {
				return ef(
					"When migrating from %d to %d, got error '%s' and "+
						"got error '%s' after trying to rollback.",
					i, i+1, err, err2)
			}
			return ef(
				"When migrating from %d to %d, got error '%s' and "+
					"successfully rolled back.", i, i+1, err)
		}
	}
	if err := m.setVersion(tx, libVersion); err != nil {
		if err2 := tx.Rollback(); err2 != nil {
			return ef(
				"When trying to set version to %d (from %d), got error '%s' "+
					"and got error '%s' after trying to rollback.",
				libVersion, dbVersion, err, err2)
		}
		return ef(
			"When trying to set version to %d (from %d), got error '%s' "+
				"and successfully rolled back.",
			libVersion, dbVersion, err)
	}
	if err := tx.Commit(); err != nil {
		return ef("Error committing migration from %d to %d: %s",
			dbVersion, libVersion, err)
	}
	return nil
}

func defaultGetVersion(tx *sql.Tx) (int, error) {
	v, err := getVersion(tx)
	if err != nil {
		if err := createVersionTable(tx); err != nil {
			return 0, err
		}
		return getVersion(tx)
	}
	return v, nil
}

func defaultSetVersion(tx *sql.Tx, version int) error {
	if err := setVersion(tx, version); err != nil {
		if err := createVersionTable(tx); err != nil {
			return err
		}
		return setVersion(tx, version)
	}
	return nil
}

func getVersion(tx *sql.Tx) (int, error) {
	var version int
	r := tx.QueryRow("SELECT version FROM migration_version")
	if err := r.Scan(&version); err != nil {
		return 0, err
	}
	return version, nil
}

func setVersion(tx *sql.Tx, version int) error {
	_, err := tx.Exec("UPDATE migration_version SET version = $1", version)
	return err
}

func createVersionTable(tx *sql.Tx) error {
	_, err := tx.Exec(`
		CREATE TABLE migration_version (
			version INTEGER
		);
		INSERT INTO migration_version (version) VALUES (0)`)
	return err
}
<commit_msg>Clearer and safer public API.<commit_after>package migration

import (
	"database/sql"
	"fmt"
)

var ef = fmt.Errorf

// LimitedTx specifies the behavior of a transaction *without* commit and
// rollback functions. Values with this type are given to client functions.
// In particular, the migration routines in this package
// handle transaction commits and rollbacks. There the functions provided by
// the client should not use them.
type LimitedTx interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) *sql.Row
	Stmt(stmt *sql.Stmt) *sql.Stmt
}

// GetVersion is any function that can retrieve the migration version of a
// particular database. It is exposed in case a client wants to override the
// default behavior of this package. (For example, by using the `user_version`
// PRAGMA in SQLite.)
//
// The DefaultGetVersion function provided with this package creates its own
// table with a single column and a single row.
//
// The version returned should be equivalent to the number of migrations
// applied to this database. It should be 0 if no migrations have been applied
// yet.
//
// If an error is returned, the migration automatically fails.
//
// Note that a LimitedTx is used to emphasize that functions with this type
// MUST NOT call Commit or Rollback. The migration routine in this pacakge will
// do it for you.
type GetVersion func(LimitedTx) (int, error)

// The default way to get the version from a database. If the database has
// had no migrations performed, then it creates a table with a single row and
// a single column storing the version as 0. It then returns 0.
//
// If the table exists, then the version stored in the table is returned.
var DefaultGetVersion GetVersion = defaultGetVersion

// SetVersion is the dual of GetVersion. It allows the client to define a
// different mechanism for setting the database version than the one used by
// DefaultSetVersion in this package.
//
// If an error is returned, the migration that tried to set the version
// automatically fails.
//
// Note that a LimitedTx is used to emphasize that functions with this type
// MUST NOT call Commit or Rollback. The migration routine in this pacakge will
// do it for you.
type SetVersion func(LimitedTx, int) error

// The default way to set the version of the database. If the database has had
// no migrations performed, then it creates a table with a single row and a
// single column and storing the version given there.
//
// If the table exists, then the existing version is overwritten.
var DefaultSetVersion SetVersion = defaultSetVersion

// Migrator corresponds to a function that updates the database by one version.
// Note that a migration should NOT call Rollback or Commit. Instead, this
// package will call Rollback for you if your migration returns an error. If
// no error is returned, then the next migration is applied. When all
// migrations have been applied, the version is updated and the changes are
// committed to the database.
type Migrator func(LimitedTx) error

// Open wraps the Open function from the database/sql package, but performs
// a series of migrations on a database if they haven't been performed already.
//
// Migrations are tracked by a simple versioning scheme. The version of the
// database is the number of migrations that have been performed on it.
// Similarly, the version of your library is the number of migrations that are
// given to this function.
//
// If Open returns successfully, then the database and your library will have
// the same versions. If there was a problem migrating---or if the database
// version is greater than your library version---then an error is returned.
// Since all migrations are performed in a single transaction, if an error
// occurs, no changes are made to the database. (Assuming you're using a
// relational database that allows modifications to a schema to be rolled back.)
//
// Note that this versioning scheme includes no semantic analysis. It is up to
// client to ensure that once a migration is defined, it never changes.
//
// The details of how the version is stored are opaque to the client, but in
// general, it will add a table to your database called "migration_version"
// with a single column containing a single row.
func Open(driver, dsn string, migrations []Migrator) (*sql.DB, error) {
	return OpenWith(driver, dsn, migrations, nil, nil)
}

// OpenWith is exactly like Open, except it allows the client to specify their
// own versioning scheme. Note that versionGet and versionSet must BOTH be
// nil or BOTH be non-nil. Otherwise, this function panics. This is because the
// implementation of one generally relies on the implementation of the other.
//
// If versionGet and versionSet are both set to nil, then the behavior of this
// function is identical to the behavior of Open.
func OpenWith(
	driver, dsn string,
	migrations []Migrator,
	versionGet GetVersion,
	versionSet SetVersion,
) (*sql.DB, error) {
	if (versionGet == nil && versionSet != nil) ||
		(versionGet != nil && versionSet == nil) {
		panic("versionGet/versionSet must both be nil or both be non-nil")
	}
	if versionGet == nil {
		versionGet = DefaultGetVersion
	}
	if versionSet == nil {
		versionSet = DefaultSetVersion
	}

	db, err := sql.Open(driver, dsn)
	if err != nil {
		return nil, err
	}
	m := migration{db, migrations, versionGet, versionSet}
	if err := m.migrate(); err != nil {
		return nil, err
	}
	return db, nil
}

type migration struct {
	*sql.DB
	migrations []Migrator
	getVersion GetVersion
	setVersion SetVersion
}

func (m migration) migrate() error {
	tx, err := m.Begin()
	if err != nil {
		return ef("Could not start transaction: %s", err)
	}
	libVersion := len(m.migrations)
	dbVersion, err := m.getVersion(tx)
	if err != nil {
		return ef("Could not get DB version: %s", err)
	}
	if dbVersion > libVersion {
		return ef("Database version (%d) is greater than library version (%d).",
			dbVersion, libVersion)
	}
	if dbVersion == libVersion {
		return tx.Commit()
	}

	for i := dbVersion; i < libVersion; i++ {
		if err := m.migrations[i](tx); err != nil {
			if err2 := tx.Rollback(); err2 != nil {
				return ef(
					"When migrating from %d to %d, got error '%s' and "+
						"got error '%s' after trying to rollback.",
					i, i+1, err, err2)
			}
			return ef(
				"When migrating from %d to %d, got error '%s' and "+
					"successfully rolled back.", i, i+1, err)
		}
	}
	if err := m.setVersion(tx, libVersion); err != nil {
		if err2 := tx.Rollback(); err2 != nil {
			return ef(
				"When trying to set version to %d (from %d), got error '%s' "+
					"and got error '%s' after trying to rollback.",
				libVersion, dbVersion, err, err2)
		}
		return ef(
			"When trying to set version to %d (from %d), got error '%s' "+
				"and successfully rolled back.",
			libVersion, dbVersion, err)
	}
	if err := tx.Commit(); err != nil {
		return ef("Error committing migration from %d to %d: %s",
			dbVersion, libVersion, err)
	}
	return nil
}

func defaultGetVersion(tx LimitedTx) (int, error) {
	v, err := getVersion(tx)
	if err != nil {
		if err := createVersionTable(tx); err != nil {
			return 0, err
		}
		return getVersion(tx)
	}
	return v, nil
}

func defaultSetVersion(tx LimitedTx, version int) error {
	if err := setVersion(tx, version); err != nil {
		if err := createVersionTable(tx); err != nil {
			return err
		}
		return setVersion(tx, version)
	}
	return nil
}

func getVersion(tx LimitedTx) (int, error) {
	var version int
	r := tx.QueryRow("SELECT version FROM migration_version")
	if err := r.Scan(&version); err != nil {
		return 0, err
	}
	return version, nil
}

func setVersion(tx LimitedTx, version int) error {
	_, err := tx.Exec("UPDATE migration_version SET version = $1", version)
	return err
}

func createVersionTable(tx LimitedTx) error {
	_, err := tx.Exec(`
		CREATE TABLE migration_version (
			version INTEGER
		);
		INSERT INTO migration_version (version) VALUES (0)`)
	return err
}
<|endoftext|>