<commit_before>/*
Copyright 2021 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 command

import (
	"errors"
	"fmt"
	"time"

	"github.com/spf13/cobra"

	"vitess.io/vitess/go/cmd/vtctldclient/cli"
	"vitess.io/vitess/go/vt/logutil"
	"vitess.io/vitess/go/vt/topo"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
	"vitess.io/vitess/go/vt/proto/vttime"
)

var (
	// CreateKeyspace makes a CreateKeyspace gRPC call to a vtctld.
	CreateKeyspace = &cobra.Command{
		Use:  "CreateKeyspace KEYSPACE_NAME [--force] [--sharding-column-name NAME --sharding-column-type TYPE] [--base-keyspace KEYSPACE --snapshot-timestamp TIME] [--served-from DB_TYPE:KEYSPACE ...]",
		Args: cobra.ExactArgs(1),
		RunE: commandCreateKeyspace,
	}
	// DeleteKeyspace makes a DeleteKeyspace gRPC call to a vtctld.
	DeleteKeyspace = &cobra.Command{
		Use:  "DeleteKeyspace KEYSPACE_NAME",
		Args: cobra.ExactArgs(1),
		RunE: commandDeleteKeyspace,
	}
	// FindAllShardsInKeyspace makes a FindAllShardsInKeyspace gRPC call to a vtctld.
	FindAllShardsInKeyspace = &cobra.Command{
		Use:     "FindAllShardsInKeyspace keyspace",
		Aliases: []string{"findallshardsinkeyspace"},
		Args:    cobra.ExactArgs(1),
		RunE:    commandFindAllShardsInKeyspace,
	}
	// GetKeyspace makes a GetKeyspace gRPC call to a vtctld.
	GetKeyspace = &cobra.Command{
		Use:     "GetKeyspace keyspace",
		Aliases: []string{"getkeyspace"},
		Args:    cobra.ExactArgs(1),
		RunE:    commandGetKeyspace,
	}
	// GetKeyspaces makes a GetKeyspaces gRPC call to a vtctld.
	GetKeyspaces = &cobra.Command{
		Use:     "GetKeyspaces",
		Aliases: []string{"getkeyspaces"},
		Args:    cobra.NoArgs,
		RunE:    commandGetKeyspaces,
	}
	// RemoveKeyspaceCell makes a RemoveKeyspaceCell gRPC call to a vtctld.
	RemoveKeyspaceCell = &cobra.Command{
		Use:  "RemoveKeyspaceCell <keyspace> <cell>",
		Args: cobra.ExactArgs(2),
		RunE: commandRemoveKeyspaceCell,
	}
)

var createKeyspaceOptions = struct {
	Force             bool
	AllowEmptyVSchema bool

	ShardingColumnName string
	ShardingColumnType cli.KeyspaceIDTypeFlag

	ServedFromsMap cli.StringMapValue

	KeyspaceType      cli.KeyspaceTypeFlag
	BaseKeyspace      string
	SnapshotTimestamp string
}{
	KeyspaceType: cli.KeyspaceTypeFlag(topodatapb.KeyspaceType_NORMAL),
}

func commandCreateKeyspace(cmd *cobra.Command, args []string) error {
	name := cmd.Flags().Arg(0)

	switch topodatapb.KeyspaceType(createKeyspaceOptions.KeyspaceType) {
	case topodatapb.KeyspaceType_NORMAL, topodatapb.KeyspaceType_SNAPSHOT:
	default:
		return fmt.Errorf("invalid keyspace type passed to --type: %v", createKeyspaceOptions.KeyspaceType)
	}

	var snapshotTime *vttime.Time
	if topodatapb.KeyspaceType(createKeyspaceOptions.KeyspaceType) == topodatapb.KeyspaceType_SNAPSHOT {
		if createKeyspaceOptions.BaseKeyspace == "" {
			return errors.New("--base-keyspace is required for a snapshot keyspace")
		}

		if createKeyspaceOptions.SnapshotTimestamp == "" {
			return errors.New("--snapshot-timestamp is required for a snapshot keyspace")
		}

		t, err := time.Parse(time.RFC3339, createKeyspaceOptions.SnapshotTimestamp)
		if err != nil {
			return fmt.Errorf("cannot parse --snapshot-timestamp as RFC3339: %w", err)
		}

		if now := time.Now(); t.After(now) {
			return fmt.Errorf("--snapshot-time cannot be in the future; snapshot = %v, now = %v", t, now)
		}

		snapshotTime = logutil.TimeToProto(t)
	}

	cli.FinishedParsing(cmd)

	req := &vtctldatapb.CreateKeyspaceRequest{
		Name:               name,
		Force:              createKeyspaceOptions.Force,
		AllowEmptyVSchema:  createKeyspaceOptions.AllowEmptyVSchema,
		ShardingColumnName: createKeyspaceOptions.ShardingColumnName,
		ShardingColumnType: topodatapb.KeyspaceIdType(createKeyspaceOptions.ShardingColumnType),
		Type:               topodatapb.KeyspaceType(createKeyspaceOptions.KeyspaceType),
		BaseKeyspace:       createKeyspaceOptions.BaseKeyspace,
		SnapshotTime:       snapshotTime,
	}

	for n, v := range createKeyspaceOptions.ServedFromsMap.StringMapValue {
		tt, err := topo.ParseServingTabletType(n)
		if err != nil {
			return err
		}

		req.ServedFroms = append(req.ServedFroms, &topodatapb.Keyspace_ServedFrom{
			TabletType: tt,
			Keyspace:   v,
		})
	}

	resp, err := client.CreateKeyspace(commandCtx, req)
	if err != nil {
		return err
	}

	data, err := cli.MarshalJSON(resp.Keyspace)
	if err != nil {
		return err
	}

	fmt.Printf("Successfully created keyspace %s. Result:\n%s\n", name, data)

	return nil
}

var deleteKeyspaceOptions = struct {
	Recursive bool
}{}

func commandDeleteKeyspace(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	ks := cmd.Flags().Arg(0)
	_, err := client.DeleteKeyspace(commandCtx, &vtctldatapb.DeleteKeyspaceRequest{
		Keyspace:  ks,
		Recursive: deleteKeyspaceOptions.Recursive,
	})

	if err != nil {
		return fmt.Errorf("DeleteKeyspace(%v) error: %w; please check the topo", ks, err)
	}

	fmt.Printf("Successfully deleted keyspace %v.\n", ks)

	return nil
}

func commandFindAllShardsInKeyspace(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	ks := cmd.Flags().Arg(0)
	resp, err := client.FindAllShardsInKeyspace(commandCtx, &vtctldatapb.FindAllShardsInKeyspaceRequest{
		Keyspace: ks,
	})

	if err != nil {
		return err
	}

	data, err := cli.MarshalJSON(resp)
	if err != nil {
		return err
	}

	fmt.Printf("%s\n", data)
	return nil
}

func commandGetKeyspace(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	ks := cmd.Flags().Arg(0)
	resp, err := client.GetKeyspace(commandCtx, &vtctldatapb.GetKeyspaceRequest{
		Keyspace: ks,
	})

	if err != nil {
		return err
	}

	fmt.Printf("%+v\n", resp.Keyspace)

	return nil
}

func commandGetKeyspaces(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	resp, err := client.GetKeyspaces(commandCtx, &vtctldatapb.GetKeyspacesRequest{})
	if err != nil {
		return err
	}

	data, err := cli.MarshalJSON(resp.Keyspaces)
	if err != nil {
		return err
	}

	fmt.Printf("%s\n", data)

	return nil
}

var removeKeyspaceCellOptions = struct {
	Force     bool
	Recursive bool
}{}

func commandRemoveKeyspaceCell(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	keyspace := cmd.Flags().Arg(0)
	cell := cmd.Flags().Arg(1)

	_, err := client.RemoveKeyspaceCell(commandCtx, &vtctldatapb.RemoveKeyspaceCellRequest{
		Keyspace:  keyspace,
		Cell:      cell,
		Force:     removeKeyspaceCellOptions.Force,
		Recursive: removeKeyspaceCellOptions.Recursive,
	})

	if err != nil {
		return err
	}

	fmt.Printf("Successfully removed keyspace %s from cell %s\n", keyspace, cell)

	return nil
}

func init() {
	CreateKeyspace.Flags().BoolVarP(&createKeyspaceOptions.Force, "force", "f", false, "Proceeds even if the keyspace already exists. Does not overwrite the existing keyspace record")
	CreateKeyspace.Flags().BoolVarP(&createKeyspaceOptions.AllowEmptyVSchema, "allow-empty-vschema", "e", false, "Allows a new keyspace to have no vschema")
	CreateKeyspace.Flags().StringVar(&createKeyspaceOptions.ShardingColumnName, "sharding-column-name", "", "The column name to use for sharding operations")
	CreateKeyspace.Flags().Var(&createKeyspaceOptions.ShardingColumnType, "sharding-column-type", "The type of the column to use for sharding operations")
	CreateKeyspace.Flags().Var(&createKeyspaceOptions.ServedFromsMap, "served-from", "TODO")
	CreateKeyspace.Flags().Var(&createKeyspaceOptions.KeyspaceType, "type", "The type of the keyspace")
	CreateKeyspace.Flags().StringVar(&createKeyspaceOptions.BaseKeyspace, "base-keyspace", "", "The base keyspace for a snapshot keyspace.")
	CreateKeyspace.Flags().StringVar(&createKeyspaceOptions.SnapshotTimestamp, "snapshot-timestamp", "", "The snapshot time for a snapshot keyspace, as a timestamp in RFC3339 format.")
	Root.AddCommand(CreateKeyspace)

	DeleteKeyspace.Flags().BoolVarP(&deleteKeyspaceOptions.Recursive, "recursive", "r", false, "Recursively delete all shards in the keyspace, and all tablets in those shards.")
	Root.AddCommand(DeleteKeyspace)

	Root.AddCommand(FindAllShardsInKeyspace)
	Root.AddCommand(GetKeyspace)
	Root.AddCommand(GetKeyspaces)

	RemoveKeyspaceCell.Flags().BoolVarP(&removeKeyspaceCellOptions.Force, "force", "f", false, "Proceed even if the cell's topology server cannot be reached. The assumption is that you turned down the entire cell, and just need to update the global topo data.")
	RemoveKeyspaceCell.Flags().BoolVarP(&removeKeyspaceCellOptions.Recursive, "recursive", "r", false, "Also delete all tablets in that cell beloning to the specified keyspace.")
	Root.AddCommand(RemoveKeyspaceCell)
}
<commit_msg>Add cli entrypoints<commit_after>/*
Copyright 2021 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 command

import (
	"errors"
	"fmt"
	"time"

	"github.com/spf13/cobra"

	"vitess.io/vitess/go/cmd/vtctldclient/cli"
	"vitess.io/vitess/go/vt/key"
	"vitess.io/vitess/go/vt/logutil"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/topoproto"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
	"vitess.io/vitess/go/vt/proto/vttime"
)

var (
	// CreateKeyspace makes a CreateKeyspace gRPC call to a vtctld.
	CreateKeyspace = &cobra.Command{
		Use:  "CreateKeyspace KEYSPACE_NAME [--force] [--sharding-column-name NAME --sharding-column-type TYPE] [--base-keyspace KEYSPACE --snapshot-timestamp TIME] [--served-from DB_TYPE:KEYSPACE ...]",
		Args: cobra.ExactArgs(1),
		RunE: commandCreateKeyspace,
	}
	// DeleteKeyspace makes a DeleteKeyspace gRPC call to a vtctld.
	DeleteKeyspace = &cobra.Command{
		Use:  "DeleteKeyspace KEYSPACE_NAME",
		Args: cobra.ExactArgs(1),
		RunE: commandDeleteKeyspace,
	}
	// FindAllShardsInKeyspace makes a FindAllShardsInKeyspace gRPC call to a vtctld.
	FindAllShardsInKeyspace = &cobra.Command{
		Use:     "FindAllShardsInKeyspace keyspace",
		Aliases: []string{"findallshardsinkeyspace"},
		Args:    cobra.ExactArgs(1),
		RunE:    commandFindAllShardsInKeyspace,
	}
	// GetKeyspace makes a GetKeyspace gRPC call to a vtctld.
	GetKeyspace = &cobra.Command{
		Use:     "GetKeyspace keyspace",
		Aliases: []string{"getkeyspace"},
		Args:    cobra.ExactArgs(1),
		RunE:    commandGetKeyspace,
	}
	// GetKeyspaces makes a GetKeyspaces gRPC call to a vtctld.
	GetKeyspaces = &cobra.Command{
		Use:     "GetKeyspaces",
		Aliases: []string{"getkeyspaces"},
		Args:    cobra.NoArgs,
		RunE:    commandGetKeyspaces,
	}
	// RemoveKeyspaceCell makes a RemoveKeyspaceCell gRPC call to a vtctld.
	RemoveKeyspaceCell = &cobra.Command{
		Use:  "RemoveKeyspaceCell <keyspace> <cell>",
		Args: cobra.ExactArgs(2),
		RunE: commandRemoveKeyspaceCell,
	}
	// SetKeyspaceServedFrom makes a SetKeyspaceServedFrom gRPC call to a vtcltd.
	SetKeyspaceServedFrom = &cobra.Command{
		Use:                   "SetKeyspaceServedFrom [--source <keyspace>] [--remove] [--cells=<cells>] <keyspace> <tablet_type>",
		Short:                 "Updates the ServedFromMap for a keyspace manually. This command is intended for emergency fixes; the map is automatically set by MigrateServedTypes. This command does not rebuild the serving graph.",
		DisableFlagsInUseLine: true,
		Args:                  cobra.ExactArgs(2),
		RunE:                  commandSetKeyspaceServedFrom,
	}
	// SetKeyspaceShardingInfo makes a SetKeyspaceShardingInfo gRPC call to a vtcltd.
	SetKeyspaceShardingInfo = &cobra.Command{
		Use:                   "SetKeyspaceShardingInfo [--force] <keyspace> [<column name> [<column type>]]",
		Short:                 "Updates the sharding information for a keyspace.",
		DisableFlagsInUseLine: true,
		Args:                  cobra.RangeArgs(1, 3),
		RunE:                  commandSetKeyspaceShardingInfo,
	}
)

var createKeyspaceOptions = struct {
	Force             bool
	AllowEmptyVSchema bool

	ShardingColumnName string
	ShardingColumnType cli.KeyspaceIDTypeFlag

	ServedFromsMap cli.StringMapValue

	KeyspaceType      cli.KeyspaceTypeFlag
	BaseKeyspace      string
	SnapshotTimestamp string
}{
	KeyspaceType: cli.KeyspaceTypeFlag(topodatapb.KeyspaceType_NORMAL),
}

func commandCreateKeyspace(cmd *cobra.Command, args []string) error {
	name := cmd.Flags().Arg(0)

	switch topodatapb.KeyspaceType(createKeyspaceOptions.KeyspaceType) {
	case topodatapb.KeyspaceType_NORMAL, topodatapb.KeyspaceType_SNAPSHOT:
	default:
		return fmt.Errorf("invalid keyspace type passed to --type: %v", createKeyspaceOptions.KeyspaceType)
	}

	var snapshotTime *vttime.Time
	if topodatapb.KeyspaceType(createKeyspaceOptions.KeyspaceType) == topodatapb.KeyspaceType_SNAPSHOT {
		if createKeyspaceOptions.BaseKeyspace == "" {
			return errors.New("--base-keyspace is required for a snapshot keyspace")
		}

		if createKeyspaceOptions.SnapshotTimestamp == "" {
			return errors.New("--snapshot-timestamp is required for a snapshot keyspace")
		}

		t, err := time.Parse(time.RFC3339, createKeyspaceOptions.SnapshotTimestamp)
		if err != nil {
			return fmt.Errorf("cannot parse --snapshot-timestamp as RFC3339: %w", err)
		}

		if now := time.Now(); t.After(now) {
			return fmt.Errorf("--snapshot-time cannot be in the future; snapshot = %v, now = %v", t, now)
		}

		snapshotTime = logutil.TimeToProto(t)
	}

	cli.FinishedParsing(cmd)

	req := &vtctldatapb.CreateKeyspaceRequest{
		Name:               name,
		Force:              createKeyspaceOptions.Force,
		AllowEmptyVSchema:  createKeyspaceOptions.AllowEmptyVSchema,
		ShardingColumnName: createKeyspaceOptions.ShardingColumnName,
		ShardingColumnType: topodatapb.KeyspaceIdType(createKeyspaceOptions.ShardingColumnType),
		Type:               topodatapb.KeyspaceType(createKeyspaceOptions.KeyspaceType),
		BaseKeyspace:       createKeyspaceOptions.BaseKeyspace,
		SnapshotTime:       snapshotTime,
	}

	for n, v := range createKeyspaceOptions.ServedFromsMap.StringMapValue {
		tt, err := topo.ParseServingTabletType(n)
		if err != nil {
			return err
		}

		req.ServedFroms = append(req.ServedFroms, &topodatapb.Keyspace_ServedFrom{
			TabletType: tt,
			Keyspace:   v,
		})
	}

	resp, err := client.CreateKeyspace(commandCtx, req)
	if err != nil {
		return err
	}

	data, err := cli.MarshalJSON(resp.Keyspace)
	if err != nil {
		return err
	}

	fmt.Printf("Successfully created keyspace %s. Result:\n%s\n", name, data)

	return nil
}

var deleteKeyspaceOptions = struct {
	Recursive bool
}{}

func commandDeleteKeyspace(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	ks := cmd.Flags().Arg(0)
	_, err := client.DeleteKeyspace(commandCtx, &vtctldatapb.DeleteKeyspaceRequest{
		Keyspace:  ks,
		Recursive: deleteKeyspaceOptions.Recursive,
	})

	if err != nil {
		return fmt.Errorf("DeleteKeyspace(%v) error: %w; please check the topo", ks, err)
	}

	fmt.Printf("Successfully deleted keyspace %v.\n", ks)

	return nil
}

func commandFindAllShardsInKeyspace(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	ks := cmd.Flags().Arg(0)
	resp, err := client.FindAllShardsInKeyspace(commandCtx, &vtctldatapb.FindAllShardsInKeyspaceRequest{
		Keyspace: ks,
	})

	if err != nil {
		return err
	}

	data, err := cli.MarshalJSON(resp)
	if err != nil {
		return err
	}

	fmt.Printf("%s\n", data)
	return nil
}

func commandGetKeyspace(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	ks := cmd.Flags().Arg(0)
	resp, err := client.GetKeyspace(commandCtx, &vtctldatapb.GetKeyspaceRequest{
		Keyspace: ks,
	})

	if err != nil {
		return err
	}

	fmt.Printf("%+v\n", resp.Keyspace)

	return nil
}

func commandGetKeyspaces(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	resp, err := client.GetKeyspaces(commandCtx, &vtctldatapb.GetKeyspacesRequest{})
	if err != nil {
		return err
	}

	data, err := cli.MarshalJSON(resp.Keyspaces)
	if err != nil {
		return err
	}

	fmt.Printf("%s\n", data)

	return nil
}

var removeKeyspaceCellOptions = struct {
	Force     bool
	Recursive bool
}{}

func commandRemoveKeyspaceCell(cmd *cobra.Command, args []string) error {
	cli.FinishedParsing(cmd)

	keyspace := cmd.Flags().Arg(0)
	cell := cmd.Flags().Arg(1)

	_, err := client.RemoveKeyspaceCell(commandCtx, &vtctldatapb.RemoveKeyspaceCellRequest{
		Keyspace:  keyspace,
		Cell:      cell,
		Force:     removeKeyspaceCellOptions.Force,
		Recursive: removeKeyspaceCellOptions.Recursive,
	})

	if err != nil {
		return err
	}

	fmt.Printf("Successfully removed keyspace %s from cell %s\n", keyspace, cell)

	return nil
}

var setKeyspaceServedFromOptions = struct {
	Cells          []string
	SourceKeyspace string
	Remove         bool
}{}

func commandSetKeyspaceServedFrom(cmd *cobra.Command, args []string) error {
	keyspace := cmd.Flags().Arg(0)
	tabletType, err := topoproto.ParseTabletType(cmd.Flags().Arg(1))
	if err != nil {
		return err
	}

	cli.FinishedParsing(cmd)

	resp, err := client.SetKeyspaceServedFrom(commandCtx, &vtctldatapb.SetKeyspaceServedFromRequest{
		Keyspace:       keyspace,
		TabletType:     tabletType,
		Cells:          setKeyspaceServedFromOptions.Cells,
		SourceKeyspace: setKeyspaceServedFromOptions.SourceKeyspace,
		Remove:         setKeyspaceServedFromOptions.Remove,
	})
	if err != nil {
		return err
	}

	data, err := cli.MarshalJSON(resp)
	if err != nil {
		return err
	}

	fmt.Printf("%s\n", data)
	return nil
}

var setKeyspaceShardingInfoOptions = struct {
	Force bool
}{}

func commandSetKeyspaceShardingInfo(cmd *cobra.Command, args []string) error {
	var (
		keyspace   = cmd.Flags().Arg(0)
		columnName string
		columnType = topodatapb.KeyspaceIdType_UNSET
	)

	switch len(cmd.Flags().Args()) {
	case 1:
		// Nothing else to do; we set keyspace already above.
	case 2:
		columnName = cmd.Flags().Arg(1)
	case 3:
		var err error
		columnType, err = key.ParseKeyspaceIDType(cmd.Flags().Arg(2))
		if err != nil {
			return err
		}
	default:
		// This should be impossible due to cobra.RangeArgs, but we handle it
		// explicitly anyway.
		return fmt.Errorf("SetKeyspaceShardingInfo expects between 1 and 3 positional args; have %d", len(cmd.Flags().Args()))
	}

	isColumnNameSet := columnName != ""
	isColumnTypeSet := columnType != topodatapb.KeyspaceIdType_UNSET

	if (isColumnNameSet && !isColumnTypeSet) || (!isColumnNameSet && isColumnTypeSet) {
		return fmt.Errorf("both <column_name:%v> and <column_type:%v> must be set, or both must be unset", columnName, key.KeyspaceIDTypeString(columnType))
	}

	cli.FinishedParsing(cmd)

	resp, err := client.SetKeyspaceShardingInfo(commandCtx, &vtctldatapb.SetKeyspaceShardingInfoRequest{
		Keyspace:   keyspace,
		ColumnName: columnName,
		ColumnType: columnType,
		Force:      setKeyspaceShardingInfoOptions.Force,
	})
	if err != nil {
		return err
	}

	data, err := cli.MarshalJSON(resp)
	if err != nil {
		return err
	}

	fmt.Printf("%s\n", data)
	return nil
}

func init() {
	CreateKeyspace.Flags().BoolVarP(&createKeyspaceOptions.Force, "force", "f", false, "Proceeds even if the keyspace already exists. Does not overwrite the existing keyspace record")
	CreateKeyspace.Flags().BoolVarP(&createKeyspaceOptions.AllowEmptyVSchema, "allow-empty-vschema", "e", false, "Allows a new keyspace to have no vschema")
	CreateKeyspace.Flags().StringVar(&createKeyspaceOptions.ShardingColumnName, "sharding-column-name", "", "The column name to use for sharding operations")
	CreateKeyspace.Flags().Var(&createKeyspaceOptions.ShardingColumnType, "sharding-column-type", "The type of the column to use for sharding operations")
	CreateKeyspace.Flags().Var(&createKeyspaceOptions.ServedFromsMap, "served-from", "TODO")
	CreateKeyspace.Flags().Var(&createKeyspaceOptions.KeyspaceType, "type", "The type of the keyspace")
	CreateKeyspace.Flags().StringVar(&createKeyspaceOptions.BaseKeyspace, "base-keyspace", "", "The base keyspace for a snapshot keyspace.")
	CreateKeyspace.Flags().StringVar(&createKeyspaceOptions.SnapshotTimestamp, "snapshot-timestamp", "", "The snapshot time for a snapshot keyspace, as a timestamp in RFC3339 format.")
	Root.AddCommand(CreateKeyspace)

	DeleteKeyspace.Flags().BoolVarP(&deleteKeyspaceOptions.Recursive, "recursive", "r", false, "Recursively delete all shards in the keyspace, and all tablets in those shards.")
	Root.AddCommand(DeleteKeyspace)

	Root.AddCommand(FindAllShardsInKeyspace)
	Root.AddCommand(GetKeyspace)
	Root.AddCommand(GetKeyspaces)

	RemoveKeyspaceCell.Flags().BoolVarP(&removeKeyspaceCellOptions.Force, "force", "f", false, "Proceed even if the cell's topology server cannot be reached. The assumption is that you turned down the entire cell, and just need to update the global topo data.")
	RemoveKeyspaceCell.Flags().BoolVarP(&removeKeyspaceCellOptions.Recursive, "recursive", "r", false, "Also delete all tablets in that cell beloning to the specified keyspace.")
	Root.AddCommand(RemoveKeyspaceCell)

	SetKeyspaceServedFrom.Flags().StringSliceVarP(&setKeyspaceServedFromOptions.Cells, "cells", "c", nil, "Cells to affect (comma-separated).")
	SetKeyspaceServedFrom.Flags().BoolVarP(&setKeyspaceServedFromOptions.Remove, "remove", "r", false, "If set, remove the ServedFrom record.")
	SetKeyspaceServedFrom.Flags().StringVar(&setKeyspaceServedFromOptions.SourceKeyspace, "source", "", "Specifies the source keyspace name.")
	Root.AddCommand(SetKeyspaceServedFrom)

	SetKeyspaceShardingInfo.Flags().BoolVarP(&setKeyspaceShardingInfoOptions.Force, "force", "f", false, "Updates fields even if they are already set. Use caution before passing force to this command.")
	Root.AddCommand(SetKeyspaceShardingInfo)
}
<|endoftext|>


<commit_before>/*
   Copyright 2010-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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 main

import (
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3"

    "fmt"
)

// Tag S3 bucket MyBucket with cost center tag "123456" and stack tag "MyTestStack".
//
// See:
//    http://docs.aws.amazon.com/awsaccountbilling/latest/aboutv2/cost-alloc-tags.html
func main() {
    // Pre-defined values
    bucket := "MyBucket"
    tagName1 := "Cost Center"
    tagValue1 := "123456"
    tagName2 := "Stack"
    tagValue2 := "MyTestStack"

    // Initialize a session in us-west-2 that the SDK will use to load credentials
    // from the shared credentials file. (~/.aws/ccredentials).
    sess, err := session.NewSession(&aws.Config{
        Region: aws.String("us-west-2")},
    )
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    // Create S3 service client
    svc := s3.New(sess)

    // Create input for PutBucket method
    input := &s3.PutBucketTaggingInput{
        Bucket: aws.String(bucket),
        Tagging: &s3.Tagging{
            TagSet: []*s3.Tag{
                {
                    Key:   aws.String(tagName1),
                    Value: aws.String(tagValue`),
                },
                {
                    Key:   aws.String(tagName2),
                    Value: aws.String(tagValue2),
              },
            },
        },
    }

    _, err = svc.PutBucketTagging(input)
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    // Now show the tags
    // Create input for GetBucket method
    input := &s3.GetBucketTaggingInput{
        Bucket: aws.String(bucket),
    }

    result, err := svc.GetBucketTagging(input)
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    numTags := len(result.TagSet)

    if numTags > 0 {
        fmt.Println("Found", numTags, "Tag(s):")
        fmt.Println("")

        for _, t := range result.TagSet {
            fmt.Println("  Key:  ", *t.Key)
            fmt.Println("  Value:", *t.Value)
            fmt.Println("")
        }
    } else {
        fmt.Println("Did not find any tags")
    }
}
<commit_msg>Fixed type in Go example code<commit_after>/*
   Copyright 2010-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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 main

import (
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3"

    "fmt"
)

// Tag S3 bucket MyBucket with cost center tag "123456" and stack tag "MyTestStack".
//
// See:
//    http://docs.aws.amazon.com/awsaccountbilling/latest/aboutv2/cost-alloc-tags.html
func main() {
    // Pre-defined values
    bucket := "MyBucket"
    tagName1 := "Cost Center"
    tagValue1 := "123456"
    tagName2 := "Stack"
    tagValue2 := "MyTestStack"

    // Initialize a session in us-west-2 that the SDK will use to load credentials
    // from the shared credentials file. (~/.aws/ccredentials).
    sess, err := session.NewSession(&aws.Config{
        Region: aws.String("us-west-2")},
    )
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    // Create S3 service client
    svc := s3.New(sess)

    // Create input for PutBucket method
    input := &s3.PutBucketTaggingInput{
        Bucket: aws.String(bucket),
        Tagging: &s3.Tagging{
            TagSet: []*s3.Tag{
                {
                    Key:   aws.String(tagName1),
                    Value: aws.String(tagValue),
                },
                {
                    Key:   aws.String(tagName2),
                    Value: aws.String(tagValue2),
              },
            },
        },
    }

    _, err = svc.PutBucketTagging(input)
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    // Now show the tags
    // Create input for GetBucket method
    input := &s3.GetBucketTaggingInput{
        Bucket: aws.String(bucket),
    }

    result, err := svc.GetBucketTagging(input)
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    numTags := len(result.TagSet)

    if numTags > 0 {
        fmt.Println("Found", numTags, "Tag(s):")
        fmt.Println("")

        for _, t := range result.TagSet {
            fmt.Println("  Key:  ", *t.Key)
            fmt.Println("  Value:", *t.Value)
            fmt.Println("")
        }
    } else {
        fmt.Println("Did not find any tags")
    }
}
<|endoftext|>


<commit_before>// +build !windows

package terminal

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"koding/kites/config"
	kos "koding/klient/os"

	"github.com/koding/passwd"
)

var (
	sessionPrefix      = "koding"
	defaultShell       = "/bin/bash"
	randomStringLength = 24 // 144 bit hex encoded
	screenEnv          []string
)

var defaultScreenPath = "/usr/bin/screen"

func init() {
	const embeddedScreen = "/opt/kite/klient/embedded/bin/screen"

	term := ""

	if fi, err := os.Stat(embeddedScreen); err == nil && !fi.IsDir() {
		defaultScreenPath = embeddedScreen
		term = "screen-256color"
	}

	SetTerm(term)
}

func SetTerm(term string) {
	if term == "" {
		term = guessTerm()
	}

	screenEnv = (kos.Environ{
		"TERM":      term,
		"HOME":      config.CurrentUser.HomeDir,
		"SCREENDIR": "/var/run/screen",
	}).Encode(nil)
}

func guessTerm() string {
	terms := [][2]string{
		{"xterm-256color", "/usr/share/terminfo/x/xterm-256color"},
		{"xterm-256color", "/usr/share/terminfo/78/xterm-256color"},
		{"xterm-color", "/usr/share/terminfo/x/xterm-color"},
	}

	for _, term := range terms {
		if _, err := os.Stat(term[1]); err == nil {
			return term[0]
		}
	}

	return "xterm"
}

type Command struct {
	// Name is used for starting the terminal instance, it's the program path
	// usually
	Name string

	// Args is passed to the program name
	Args []string

	// Session id used for reconnections, used by screen or tmux
	Session string
}

var (
	ErrNoSession     = errors.New("session doesn't exists")
	ErrSessionExists = errors.New("session with the same name exists already")
)

func getUserEntry(username string) (*passwd.Entry, error) {
	entries, err := passwd.Parse()
	if err != nil {
		return nil, err
	}

	user, ok := entries[username]
	if !ok {
		return nil, err
	}

	if user.Shell == "" {
		return nil, err
	}

	return &user, nil
}

func getDefaultShell(username string) string {
	if runtime.GOOS == "darwin" {
		return defaultShell
	}

	entry, err := getUserEntry(username)
	if err != nil {
		log.Println("terminal: couldn't get default shell ", err)
		return defaultShell
	}

	return entry.Shell
}

// newCmd returns a new command instance that is used to start the terminal.
// The command line is created differently based on the incoming mode.
func (t *terminal) newCommand(mode, session, username string) (*Command, error) {
	// let's assume by default its Screen
	name := defaultScreenPath
	defaultShell := getDefaultShell(username)
	args := []string{"-e^Bb", "-s", defaultShell, "-S"}

	// TODO: resume and create are backwards compatible modes. Remove then once
	// the client side switched to use the "attach" mode which does both,
	// resume or create.
	switch mode {
	case "shared", "resume":
		if session == "" {
			return nil, errors.New("session is needed for 'shared' or 'resume' mode")
		}

		if !t.sessionExists(session, username) {
			return nil, ErrNoSession
		}

		args = append(args, sessionPrefix+"."+session)
		if mode == "shared" {
			args = append(args, "-x") // multiuser mode
		} else if mode == "resume" {
			args = append(args, "-raAd") // resume
		}
	case "noscreen":
		name = defaultShell
		args = []string{}
	case "attach", "create":
		if session == "" {
			// if the user didn't send a session name, create a custom
			// randomized
			session = randomString()
			args = append(args, sessionPrefix+"."+session)
		} else {
			// -a  : includes all capabilities
			// -A  : adapts the sizes of all windows to the current terminal
			// -DR : if session is running, re attach. If not create a new one
			args = append(args, sessionPrefix+"."+session, "-aADR")
		}
	default:
		return nil, fmt.Errorf("mode '%s' is unknown. Valid modes are:  [shared|noscreen|resume|create]", mode)
	}

	c := &Command{
		Name:    name,
		Args:    args,
		Session: session,
	}

	return c, nil
}

// screenSessions returns a list of sessions that belongs to the given
// username.  The sessions are in the form of ["k7sdjv12344", "askIj12sas12",
// ...]
// TODO: socket directory is different under darwin, it will not work probably
func (t *terminal) screenSessions(username string) []string {
	// Do not include dead sessions in our result
	t.run(defaultScreenPath, "-wipe")

	// We need to use ls here, because /var/run/screen mount is only
	// visible from inside of container. Errors are ignored.
	stdout, stderr, err := t.run("ls", "/var/run/screen/S-"+username)
	if err != nil {
		t.Log.Error("terminal: listing sessions failed: %s:\n%s\n", err, stderr)
		return nil
	}

	shellOut := string(bytes.TrimSpace(stdout))
	if shellOut == "" {
		return nil
	}

	names := strings.Split(shellOut, "\n")
	sessions := make([]string, len(names))

	prefix := sessionPrefix + "."
	for i, name := range names {
		segments := strings.SplitN(name, ".", 2)
		sessions[i] = strings.TrimPrefix(segments[1], prefix)
	}

	return sessions
}

// screenExists checks whether the given session exists in the running list of
// screen sessions.
func (t *terminal) sessionExists(session, username string) bool {
	for _, s := range t.screenSessions(username) {
		if s == session {
			return true
		}
	}

	return false
}

// killSessions kills all screen sessions for given username
func (t *terminal) killSessions(username string) error {
	for _, session := range t.screenSessions(username) {
		if err := t.killSession(session); err != nil {
			return err
		}
	}

	return nil
}

// killSession kills the given SessionID
func (t *terminal) killSession(session string) error {
	stdout, stderr, err := t.run(defaultScreenPath, "-X", "-S", sessionPrefix+"."+session, "kill")
	if err != nil {
		return commandError("screen kill failed", err, stdout, stderr)
	}

	return nil
}

func (t *terminal) renameSession(oldName, newName string) error {
	stdout, stderr, err := t.run(defaultScreenPath, "-X", "-S", sessionPrefix+"."+oldName, "sessionname", sessionPrefix+"."+newName)
	if err != nil {
		return commandError("screen renaming failed", err, stdout, stderr)
	}

	return nil
}

func commandError(message string, err error, stdout, stderr []byte) error {
	return fmt.Errorf("%s\n%s\n%s\n%s\n", message, err, stdout, stderr)
}

func (t *terminal) run(cmd string, args ...string) (stdout, stderr []byte, err error) {
	var bufout, buferr bytes.Buffer

	c := exec.Command(cmd, args...)
	c.Stdout = &bufout
	c.Stderr = &buferr
	c.Env = screenEnv

	t.Log.Debug("terminal: running: %v (%v)", c.Args, screenEnv)

	if err := c.Run(); err != nil {
		return nil, buferr.Bytes(), err
	}

	return bufout.Bytes(), nil, nil
}

func randomString() string {
	p := make([]byte, randomStringLength/2+1)
	rand.Read(p)
	return hex.EncodeToString(p)[:randomStringLength]
}
<commit_msg>klient/terminal: add inline doc<commit_after>// +build !windows

package terminal

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"koding/kites/config"
	kos "koding/klient/os"

	"github.com/koding/passwd"
)

var (
	sessionPrefix      = "koding"
	defaultShell       = "/bin/bash"
	randomStringLength = 24 // 144 bit hex encoded
	screenEnv          []string
)

var defaultScreenPath = "/usr/bin/screen"

func init() {
	const embeddedScreen = "/opt/kite/klient/embedded/bin/screen"

	term := ""

	if fi, err := os.Stat(embeddedScreen); err == nil && !fi.IsDir() {
		defaultScreenPath = embeddedScreen
		term = "screen-256color"
	}

	SetTerm(term)
}

// SetTerm changes the TERM environment variable used with
// screen processes.
//
// The function cannot be called after Terminal starts
// accepting kite requets.
func SetTerm(term string) {
	if term == "" {
		term = guessTerm()
	}

	screenEnv = (kos.Environ{
		"TERM":      term,
		"HOME":      config.CurrentUser.HomeDir,
		"SCREENDIR": "/var/run/screen",
	}).Encode(nil)
}

func guessTerm() string {
	terms := [][2]string{
		{"xterm-256color", "/usr/share/terminfo/x/xterm-256color"},
		{"xterm-256color", "/usr/share/terminfo/78/xterm-256color"},
		{"xterm-color", "/usr/share/terminfo/x/xterm-color"},
	}

	for _, term := range terms {
		if _, err := os.Stat(term[1]); err == nil {
			return term[0]
		}
	}

	return "xterm"
}

type Command struct {
	// Name is used for starting the terminal instance, it's the program path
	// usually
	Name string

	// Args is passed to the program name
	Args []string

	// Session id used for reconnections, used by screen or tmux
	Session string
}

var (
	ErrNoSession     = errors.New("session doesn't exists")
	ErrSessionExists = errors.New("session with the same name exists already")
)

func getUserEntry(username string) (*passwd.Entry, error) {
	entries, err := passwd.Parse()
	if err != nil {
		return nil, err
	}

	user, ok := entries[username]
	if !ok {
		return nil, err
	}

	if user.Shell == "" {
		return nil, err
	}

	return &user, nil
}

func getDefaultShell(username string) string {
	if runtime.GOOS == "darwin" {
		return defaultShell
	}

	entry, err := getUserEntry(username)
	if err != nil {
		log.Println("terminal: couldn't get default shell ", err)
		return defaultShell
	}

	return entry.Shell
}

// newCmd returns a new command instance that is used to start the terminal.
// The command line is created differently based on the incoming mode.
func (t *terminal) newCommand(mode, session, username string) (*Command, error) {
	// let's assume by default its Screen
	name := defaultScreenPath
	defaultShell := getDefaultShell(username)
	args := []string{"-e^Bb", "-s", defaultShell, "-S"}

	// TODO: resume and create are backwards compatible modes. Remove then once
	// the client side switched to use the "attach" mode which does both,
	// resume or create.
	switch mode {
	case "shared", "resume":
		if session == "" {
			return nil, errors.New("session is needed for 'shared' or 'resume' mode")
		}

		if !t.sessionExists(session, username) {
			return nil, ErrNoSession
		}

		args = append(args, sessionPrefix+"."+session)
		if mode == "shared" {
			args = append(args, "-x") // multiuser mode
		} else if mode == "resume" {
			args = append(args, "-raAd") // resume
		}
	case "noscreen":
		name = defaultShell
		args = []string{}
	case "attach", "create":
		if session == "" {
			// if the user didn't send a session name, create a custom
			// randomized
			session = randomString()
			args = append(args, sessionPrefix+"."+session)
		} else {
			// -a  : includes all capabilities
			// -A  : adapts the sizes of all windows to the current terminal
			// -DR : if session is running, re attach. If not create a new one
			args = append(args, sessionPrefix+"."+session, "-aADR")
		}
	default:
		return nil, fmt.Errorf("mode '%s' is unknown. Valid modes are:  [shared|noscreen|resume|create]", mode)
	}

	c := &Command{
		Name:    name,
		Args:    args,
		Session: session,
	}

	return c, nil
}

// screenSessions returns a list of sessions that belongs to the given
// username.  The sessions are in the form of ["k7sdjv12344", "askIj12sas12",
// ...]
// TODO: socket directory is different under darwin, it will not work probably
func (t *terminal) screenSessions(username string) []string {
	// Do not include dead sessions in our result
	t.run(defaultScreenPath, "-wipe")

	// We need to use ls here, because /var/run/screen mount is only
	// visible from inside of container. Errors are ignored.
	stdout, stderr, err := t.run("ls", "/var/run/screen/S-"+username)
	if err != nil {
		t.Log.Error("terminal: listing sessions failed: %s:\n%s\n", err, stderr)
		return nil
	}

	shellOut := string(bytes.TrimSpace(stdout))
	if shellOut == "" {
		return nil
	}

	names := strings.Split(shellOut, "\n")
	sessions := make([]string, len(names))

	prefix := sessionPrefix + "."
	for i, name := range names {
		segments := strings.SplitN(name, ".", 2)
		sessions[i] = strings.TrimPrefix(segments[1], prefix)
	}

	return sessions
}

// screenExists checks whether the given session exists in the running list of
// screen sessions.
func (t *terminal) sessionExists(session, username string) bool {
	for _, s := range t.screenSessions(username) {
		if s == session {
			return true
		}
	}

	return false
}

// killSessions kills all screen sessions for given username
func (t *terminal) killSessions(username string) error {
	for _, session := range t.screenSessions(username) {
		if err := t.killSession(session); err != nil {
			return err
		}
	}

	return nil
}

// killSession kills the given SessionID
func (t *terminal) killSession(session string) error {
	stdout, stderr, err := t.run(defaultScreenPath, "-X", "-S", sessionPrefix+"."+session, "kill")
	if err != nil {
		return commandError("screen kill failed", err, stdout, stderr)
	}

	return nil
}

func (t *terminal) renameSession(oldName, newName string) error {
	stdout, stderr, err := t.run(defaultScreenPath, "-X", "-S", sessionPrefix+"."+oldName, "sessionname", sessionPrefix+"."+newName)
	if err != nil {
		return commandError("screen renaming failed", err, stdout, stderr)
	}

	return nil
}

func commandError(message string, err error, stdout, stderr []byte) error {
	return fmt.Errorf("%s\n%s\n%s\n%s\n", message, err, stdout, stderr)
}

func (t *terminal) run(cmd string, args ...string) (stdout, stderr []byte, err error) {
	var bufout, buferr bytes.Buffer

	c := exec.Command(cmd, args...)
	c.Stdout = &bufout
	c.Stderr = &buferr
	c.Env = screenEnv

	t.Log.Debug("terminal: running: %v (%v)", c.Args, screenEnv)

	if err := c.Run(); err != nil {
		return nil, buferr.Bytes(), err
	}

	return bufout.Bytes(), nil, nil
}

func randomString() string {
	p := make([]byte, randomStringLength/2+1)
	rand.Read(p)
	return hex.EncodeToString(p)[:randomStringLength]
}
<|endoftext|>


<commit_before>// Package Cephbackupstorage implements the BackupStorage interface
// for Ceph Cloud Storage.
package cephbackupstorage

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"sync"

	"github.com/minio/minio-go"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/mysqlctl/backupstorage"
)

var (
	// bucket is where the backups will go.
	bucket string
	// configFilePath is where the configs/credentials for backups will be stored.
	configFilePath = flag.String("ceph_backup_storage_config", "ceph_backup_config.json",
		"Path to JSON config file for ceph backup storage")
)

var StorageConfig struct {
	AccessKey string `json:"accessKey"`
	SecretKey string `json:"secretKey"`
	EndPoint  string `json:"endPoint"`
	Bucket    string `json:"bucket"`
}

// CephBackupHandle implements BackupHandle for Ceph Cloud Storage.
type CephBackupHandle struct {
	client    *minio.Client
	bs        *CephBackupStorage
	dir       string
	name      string
	readOnly  bool
	errors    concurrency.AllErrorRecorder
	waitGroup sync.WaitGroup
}

// Directory implements BackupHandle.
func (bh *CephBackupHandle) Directory() string {
	return bh.dir
}

// Name implements BackupHandle.
func (bh *CephBackupHandle) Name() string {
	return bh.name
}

// AddFile implements BackupHandle.
func (bh *CephBackupHandle) AddFile(filename string) (io.WriteCloser, error) {
	if bh.readOnly {
		return nil, fmt.Errorf("AddFile cannot be called on read-only backup")
	}
	reader, writer := io.Pipe()
	bh.waitGroup.Add(1)
	go func() {
		defer bh.waitGroup.Done()
		// Give PutObject() the read end of the pipe.
		object := objName(bh.dir, bh.name, filename)
		_, err := bh.client.PutObject(bucket, object, reader, "application/octet-stream")
		if err != nil {
			// Signal the writer that an error occurred, in case it's not done writing yet.
			reader.CloseWithError(err)
			// In case the error happened after the writer finished, we need to remember it.
			bh.errors.RecordError(err)
		}
	}()
	// Give our caller the write end of the pipe.
	return writer, nil
}

// EndBackup implements BackupHandle.
func (bh *CephBackupHandle) EndBackup() error {
	if bh.readOnly {
		return fmt.Errorf("EndBackup cannot be called on read-only backup")
	}
	bh.waitGroup.Wait()
	// Return the saved PutObject() errors, if any.
	return bh.errors.Error()
}

// AbortBackup implements BackupHandle.
func (bh *CephBackupHandle) AbortBackup() error {
	if bh.readOnly {
		return fmt.Errorf("AbortBackup cannot be called on read-only backup")
	}
	return bh.bs.RemoveBackup(bh.dir, bh.name)
}

// ReadFile implements BackupHandle.
func (bh *CephBackupHandle) ReadFile(filename string) (io.ReadCloser, error) {
	if !bh.readOnly {
		return nil, fmt.Errorf("ReadFile cannot be called on read-write backup")
	}
	object := objName(bh.dir, bh.name, filename)
	return bh.client.GetObject(bucket, object)
}

// CephBackupStorage implements BackupStorage for Ceph Cloud Storage.
type CephBackupStorage struct {
	// client is the instance of the Ceph Cloud Storage Go client.
	// Once this field is set, it must not be written again/unset to nil.
	_client *minio.Client
	// mu guards all fields.
	mu sync.Mutex
}

// ListBackups implements BackupStorage.
func (bs *CephBackupStorage) ListBackups(dir string) ([]backupstorage.BackupHandle, error) {
	c, err := bs.client()
	if err != nil {
		return nil, err
	}

	// List prefixes that begin with dir (i.e. list subdirs).
	var subdirs []string
	searchPrefix := objName(dir, "")

	doneCh := make(chan struct{})
	for object := range c.ListObjects(bucket, searchPrefix, false, doneCh) {
		if object.Err != nil {
			return nil, object.Err
		}
		subdir := strings.TrimPrefix(object.Key, searchPrefix)
		subdir = strings.TrimSuffix(subdir, "/")
		subdirs = append(subdirs, subdir)
	}

	// Backups must be returned in order, oldest first.
	sort.Strings(subdirs)

	result := make([]backupstorage.BackupHandle, 0, len(subdirs))
	for _, subdir := range subdirs {
		result = append(result, &CephBackupHandle{
			client:   c,
			bs:       bs,
			dir:      dir,
			name:     subdir,
			readOnly: true,
		})
	}
	return result, nil
}

// StartBackup implements BackupStorage.
func (bs *CephBackupStorage) StartBackup(dir, name string) (backupstorage.BackupHandle, error) {
	c, err := bs.client()
	if err != nil {
		return nil, err
	}

	return &CephBackupHandle{
		client:   c,
		bs:       bs,
		dir:      dir,
		name:     name,
		readOnly: false,
	}, nil
}

// RemoveBackup implements BackupStorage.
func (bs *CephBackupStorage) RemoveBackup(dir, name string) error {
	c, err := bs.client()
	if err != nil {
		return err
	}
	fullName := objName(dir, name, "")
	//	err = c.RemoveObject(bucket, fullName)
	//      if err != nil {
	//              return err
	//      }
	//      return nil
	var arr []string
	doneCh := make(chan struct{})
	defer close(doneCh)
	for object := range c.ListObjects(bucket, fullName, true, doneCh) {
		if object.Err != nil {
			return object.Err
		}
		arr = append(arr, object.Key)
	}
	for _, obj := range arr {
		err = c.RemoveObject(bucket, obj)
		if err != nil {
			return err
		}
	}
	return nil
}

// Close implements BackupStorage.
func (bs *CephBackupStorage) Close() error {
	bs.mu.Lock()
	defer bs.mu.Unlock()

	if bs._client != nil {
		// a new client the next time one is needed.
		bs._client = nil
	}
	return nil
}

// client returns the Ceph Storage client instance.
// If there isn't one yet, it tries to create one.
func (bs *CephBackupStorage) client() (*minio.Client, error) {
	bs.mu.Lock()
	defer bs.mu.Unlock()

	if bs._client == nil {
		configFile, err := os.Open(*configFilePath)
		if err != nil {
			return nil, fmt.Errorf("file not present : %v", err)
		}
		defer configFile.Close()
		jsonParser := json.NewDecoder(configFile)
		if err = jsonParser.Decode(&StorageConfig); err != nil {
			return nil, fmt.Errorf("Error parsing the json file : %v", err)
		}

		bucket = StorageConfig.Bucket
		accessKey := StorageConfig.AccessKey
		secretKey := StorageConfig.SecretKey
		url := StorageConfig.EndPoint

		ceph_client, err := minio.NewV2(url, accessKey, secretKey, true)
		if err != nil {
			return nil, err
		}
		bs._client = ceph_client
	}
	return bs._client, nil
}

func init() {
	backupstorage.BackupStorageMap["ceph"] = &CephBackupStorage{}
}

// objName joins path parts into an object name.
// Unlike path.Join, it doesn't collapse ".." or strip trailing slashes.
// It also adds the value of the -gcs_backup_storage_root flag if set.
func objName(parts ...string) string {
	return strings.Join(parts, "/")
}
<commit_msg>Fix golint issues in cephbackupstorage.<commit_after>// Package cephbackupstorage implements the BackupStorage interface
// for Ceph Cloud Storage.
package cephbackupstorage

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"sync"

	minio "github.com/minio/minio-go"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/mysqlctl/backupstorage"
)

var (
	// bucket is where the backups will go.
	bucket string
	// configFilePath is where the configs/credentials for backups will be stored.
	configFilePath = flag.String("ceph_backup_storage_config", "ceph_backup_config.json",
		"Path to JSON config file for ceph backup storage")
)

var storageConfig struct {
	AccessKey string `json:"accessKey"`
	SecretKey string `json:"secretKey"`
	EndPoint  string `json:"endPoint"`
	Bucket    string `json:"bucket"`
}

// CephBackupHandle implements BackupHandle for Ceph Cloud Storage.
type CephBackupHandle struct {
	client    *minio.Client
	bs        *CephBackupStorage
	dir       string
	name      string
	readOnly  bool
	errors    concurrency.AllErrorRecorder
	waitGroup sync.WaitGroup
}

// Directory implements BackupHandle.
func (bh *CephBackupHandle) Directory() string {
	return bh.dir
}

// Name implements BackupHandle.
func (bh *CephBackupHandle) Name() string {
	return bh.name
}

// AddFile implements BackupHandle.
func (bh *CephBackupHandle) AddFile(filename string) (io.WriteCloser, error) {
	if bh.readOnly {
		return nil, fmt.Errorf("AddFile cannot be called on read-only backup")
	}
	reader, writer := io.Pipe()
	bh.waitGroup.Add(1)
	go func() {
		defer bh.waitGroup.Done()
		// Give PutObject() the read end of the pipe.
		object := objName(bh.dir, bh.name, filename)
		_, err := bh.client.PutObject(bucket, object, reader, "application/octet-stream")
		if err != nil {
			// Signal the writer that an error occurred, in case it's not done writing yet.
			reader.CloseWithError(err)
			// In case the error happened after the writer finished, we need to remember it.
			bh.errors.RecordError(err)
		}
	}()
	// Give our caller the write end of the pipe.
	return writer, nil
}

// EndBackup implements BackupHandle.
func (bh *CephBackupHandle) EndBackup() error {
	if bh.readOnly {
		return fmt.Errorf("EndBackup cannot be called on read-only backup")
	}
	bh.waitGroup.Wait()
	// Return the saved PutObject() errors, if any.
	return bh.errors.Error()
}

// AbortBackup implements BackupHandle.
func (bh *CephBackupHandle) AbortBackup() error {
	if bh.readOnly {
		return fmt.Errorf("AbortBackup cannot be called on read-only backup")
	}
	return bh.bs.RemoveBackup(bh.dir, bh.name)
}

// ReadFile implements BackupHandle.
func (bh *CephBackupHandle) ReadFile(filename string) (io.ReadCloser, error) {
	if !bh.readOnly {
		return nil, fmt.Errorf("ReadFile cannot be called on read-write backup")
	}
	object := objName(bh.dir, bh.name, filename)
	return bh.client.GetObject(bucket, object)
}

// CephBackupStorage implements BackupStorage for Ceph Cloud Storage.
type CephBackupStorage struct {
	// client is the instance of the Ceph Cloud Storage Go client.
	// Once this field is set, it must not be written again/unset to nil.
	_client *minio.Client
	// mu guards all fields.
	mu sync.Mutex
}

// ListBackups implements BackupStorage.
func (bs *CephBackupStorage) ListBackups(dir string) ([]backupstorage.BackupHandle, error) {
	c, err := bs.client()
	if err != nil {
		return nil, err
	}

	// List prefixes that begin with dir (i.e. list subdirs).
	var subdirs []string
	searchPrefix := objName(dir, "")

	doneCh := make(chan struct{})
	for object := range c.ListObjects(bucket, searchPrefix, false, doneCh) {
		if object.Err != nil {
			return nil, object.Err
		}
		subdir := strings.TrimPrefix(object.Key, searchPrefix)
		subdir = strings.TrimSuffix(subdir, "/")
		subdirs = append(subdirs, subdir)
	}

	// Backups must be returned in order, oldest first.
	sort.Strings(subdirs)

	result := make([]backupstorage.BackupHandle, 0, len(subdirs))
	for _, subdir := range subdirs {
		result = append(result, &CephBackupHandle{
			client:   c,
			bs:       bs,
			dir:      dir,
			name:     subdir,
			readOnly: true,
		})
	}
	return result, nil
}

// StartBackup implements BackupStorage.
func (bs *CephBackupStorage) StartBackup(dir, name string) (backupstorage.BackupHandle, error) {
	c, err := bs.client()
	if err != nil {
		return nil, err
	}

	return &CephBackupHandle{
		client:   c,
		bs:       bs,
		dir:      dir,
		name:     name,
		readOnly: false,
	}, nil
}

// RemoveBackup implements BackupStorage.
func (bs *CephBackupStorage) RemoveBackup(dir, name string) error {
	c, err := bs.client()
	if err != nil {
		return err
	}
	fullName := objName(dir, name, "")
	//	err = c.RemoveObject(bucket, fullName)
	//      if err != nil {
	//              return err
	//      }
	//      return nil
	var arr []string
	doneCh := make(chan struct{})
	defer close(doneCh)
	for object := range c.ListObjects(bucket, fullName, true, doneCh) {
		if object.Err != nil {
			return object.Err
		}
		arr = append(arr, object.Key)
	}
	for _, obj := range arr {
		err = c.RemoveObject(bucket, obj)
		if err != nil {
			return err
		}
	}
	return nil
}

// Close implements BackupStorage.
func (bs *CephBackupStorage) Close() error {
	bs.mu.Lock()
	defer bs.mu.Unlock()

	if bs._client != nil {
		// a new client the next time one is needed.
		bs._client = nil
	}
	return nil
}

// client returns the Ceph Storage client instance.
// If there isn't one yet, it tries to create one.
func (bs *CephBackupStorage) client() (*minio.Client, error) {
	bs.mu.Lock()
	defer bs.mu.Unlock()

	if bs._client == nil {
		configFile, err := os.Open(*configFilePath)
		if err != nil {
			return nil, fmt.Errorf("file not present : %v", err)
		}
		defer configFile.Close()
		jsonParser := json.NewDecoder(configFile)
		if err = jsonParser.Decode(&storageConfig); err != nil {
			return nil, fmt.Errorf("Error parsing the json file : %v", err)
		}

		bucket = storageConfig.Bucket
		accessKey := storageConfig.AccessKey
		secretKey := storageConfig.SecretKey
		url := storageConfig.EndPoint

		client, err := minio.NewV2(url, accessKey, secretKey, true)
		if err != nil {
			return nil, err
		}
		bs._client = client
	}
	return bs._client, nil
}

func init() {
	backupstorage.BackupStorageMap["ceph"] = &CephBackupStorage{}
}

// objName joins path parts into an object name.
// Unlike path.Join, it doesn't collapse ".." or strip trailing slashes.
// It also adds the value of the -gcs_backup_storage_root flag if set.
func objName(parts ...string) string {
	return strings.Join(parts, "/")
}
<|endoftext|>


<commit_before>package loopback

import (
	"sync"
	"io"
	"os"
	"time"
)

type block struct {
	t    int64
	data []byte
	prev *block
	next *block
}

type streamReader stream
type streamWriter stream

func (r *streamReader) Read(data []byte) (int, os.Error) {
	return (*stream)(r).Read(data)
}

func (r *streamReader) Close() os.Error {
	return (*stream)(r).closeInput()
}

func (w *streamWriter) Write(data []byte) (int, os.Error) {
	return (*stream)(w).Write(data)
}

func (w *streamWriter) Close() os.Error {
	return (*stream)(w).closeOutput()
}

type stream struct {
	mu sync.Mutex

	outClosed bool
	inClosed  bool

	outTail     *block // sentinel.
	outHead     *block // also transitTail.
	transitHead *block // also inTail.
	inHead      *block // overall head of list.

	outLimit int // total size of output queue.
	outAvail int // free bytes in output queue.

	inLimit int // total size of input queue.
	inAvail int // free bytes in input queue.

	byteDelay int64
	latency   int64
	mtu       int

	notEmpty     sync.Cond
	notFull      sync.Cond
}

// Loopback options for use with Pipe.
type Options struct {
	// ByteDelay controls the time a packet takes in the link.  A packet
	// n bytes long takes ByteDelay * n nanoseconds to exit
	// the output queue and is available for reading Latency
	// nanoseconds later.
	ByteDelay int64
	Latency   int64

	// MTU gives the maximum packet size that can
	// be tranferred atomically across the link.
	// Larger packet will be split.
	// If this is zero, a default of 32768 is assumed
	MTU int

	// InLimit and OutLimit gives the size of the input and output queues.
	// If either is zero, a default of 10*MTU is assumed.
	InLimit  int
	OutLimit int
}

// Pipe creates an asynchronous in-memory pipe,
// Writes are divided into packets of at most opts.MTU bytes
// written to a flow-controlled output queue, transferred across the link,
// and put into an input queue where it is readable with the r.
// The options determine when and how the data will be transferred.
func Pipe(opt Options) (r io.ReadCloser, w io.WriteCloser) {
	if opt.MTU == 0 {
		opt.MTU = 32768
	}
	if opt.InLimit == 0 {
		opt.InLimit = 10 * opt.MTU
	}
	if opt.OutLimit == 0 {
		opt.OutLimit = 10 * opt.MTU
	}
	if opt.InLimit < opt.MTU {
		opt.InLimit = opt.MTU
	}
	if opt.OutLimit < opt.MTU {
		opt.OutLimit = opt.MTU
	}
	sentinel := &block{}
	s := &stream{
		outLimit:    opt.OutLimit,
		outAvail:    opt.OutLimit,
		inLimit:     opt.InLimit,
		inAvail:     opt.InLimit,
		mtu:         opt.MTU,
		byteDelay:   opt.ByteDelay,
		latency:     opt.Latency,
		outTail:     sentinel,
		outHead:     sentinel,
		transitHead: sentinel,
		inHead:      sentinel,
	}
	s.notEmpty.L = &s.mu
	s.notFull.L = &s.mu
	return (*streamReader)(s), (*streamWriter)(s)
}

// Dodgy heuristic:
// If there's stuff in the transit queue that's ready to
// enter the input queue, but the input queue is full
// and it's been waiting for at least latency ns,
// then we block the output queue.
// TODO what do we do about latency for
// blocked packets - as it is a blocked packet
// will incur less latency.
func (s *stream) outBlocked(now int64) bool {
	return s.transitHead != s.outHead &&
		now >= s.transitHead.t+s.latency &&
		s.inAvail < len(s.transitHead.data)
}

func (s *stream) closeInput() os.Error {
	s.mu.Lock()
	s.inClosed = true
	s.notEmpty.Broadcast()
	s.notFull.Broadcast()
	s.mu.Unlock()
	return nil
}

func (s *stream) closeOutput() os.Error {
	s.mu.Lock()
	s.outClosed = true
	s.notEmpty.Broadcast()
	s.notFull.Broadcast()
	s.mu.Unlock()
	return nil
}

func (s *stream) pushLink(now int64) {
	if !s.outBlocked(now) {
		// move blocks from out queue to transit queue.
		for s.outTail != s.outHead && now >= s.outHead.t {
			s.outHead.t += s.latency
			s.outAvail += len(s.outHead.data)
			s.outHead = s.outHead.next
		}
	}
	// move blocks from transit queue to input queue
	for s.transitHead != s.outHead && now >= s.transitHead.t {
		if s.inAvail < len(s.transitHead.data) {
			break // or discard packet
		}
		s.inAvail -= len(s.transitHead.data)
		s.transitHead = s.transitHead.next
	}
}

func (s *stream) Write(data []byte) (int, os.Error) {
	// split the packet into MTU-sized portions if necessary.
	for len(data) > s.mtu {
		_, err := s.Write(data[0:s.mtu])
		if err != nil {
			return 0, err
		}
		data = data[s.mtu:]
	}
	s.mu.Lock()
	now := time.Nanoseconds()
	for {
		s.pushLink(now)
		if s.outAvail >= len(data) || s.outClosed {
			break
		}
		if s.outBlocked(time.Nanoseconds()) {
			if s.inClosed {
				s.mu.Unlock()
				return 0, os.EPIPE
			}
			s.notFull.Wait()
			continue
		}
		t := s.earliestWriteTime(len(data))
		now = s.sleepUntil(t)
	}
	if s.outClosed {
		s.mu.Unlock()
		return 0, os.EPIPE
	}
	delay := int64(len(data)) * s.byteDelay
	var t int64
	// If there's a block in the queue that's not yet due
	// for transit, then this block leaves delay ns after
	// that one.
	if s.outHead != s.outTail && now < s.outTail.prev.t {
		t = s.outTail.prev.t + delay
	} else {
		t = now + delay
	}
	s.addBlock(t, s.copy(data))
	s.outAvail -= len(data)

	s.notEmpty.Broadcast()
	s.mu.Unlock()
	// TODO runtime.Gosched() ?
	return len(data), nil
}

func (s *stream) Read(buf []byte) (int, os.Error) {
	s.mu.Lock()
	// Loop until there's something to read from the input queue.
	now := time.Nanoseconds()
	for {
		s.pushLink(now)
		if s.inHead != s.transitHead {
			break
		}
		if s.inHead == s.outTail {
			// No data at all in the queue.
			// If the queue is empty and the output queue is closed,
			// then we see EOF.
			if s.outClosed {
				s.mu.Unlock()
				return 0, os.EOF
			}
			s.notEmpty.Wait()
			continue
		}
		now = s.sleepUntil(s.earliestReadTime())
	}
	if s.inClosed {
		// input queue has been forcibly closed:
		// TODO is os.EOF the right error here?
		return 0, os.EOF
	}
	b := s.inHead
	n := copy(buf, b.data)
	b.data = b.data[n:]
	s.inAvail += n
	if len(b.data) == 0 {
		s.removeBlock()
	}
	// Wake up any writers blocked on a full queue.
	s.notFull.Broadcast()
	s.mu.Unlock()
	return n, nil
}

// earliestReadTime returns the earliest time that
// some data might arrive into the input queue.
// It assumes that there is some data in the system.
func (s *stream) earliestReadTime() int64 {
	if s.inAvail < s.inLimit {
		// data is available right now.
		return 0
	}
	if s.transitHead != s.outHead {
		return s.transitHead.t
	}
	if s.outHead != s.outTail {
		return s.outHead.t + s.latency
	}
	panic("no data")
}

// earliestWriteTime returns the earliest time that
// there may be space for n bytes of data to be
// placed into the output queue (it might be later
// if packets are dropped).
func (s *stream) earliestWriteTime(n int) int64 {
	if s.outAvail < s.outLimit {
		// space is available now.
		return 0
	}
	tot := s.outAvail
	for b := s.outHead; b != s.outTail; b = b.next {
		tot += len(b.data)
		if tot >= n {
			return b.t
		}
	}
	panic("write limit exceeded by block size")
}

// sleep until the absolute time t.
// Called with lock held.
func (s *stream) sleepUntil(t int64) int64 {
	now := time.Nanoseconds()
	if now >= t {
		return now
	}
	s.mu.Unlock()
	time.Sleep(t - now)
	s.mu.Lock()
	return time.Nanoseconds()
}

func (s *stream) copy(x []byte) []byte {
	y := make([]byte, len(x))
	copy(y, x)
	return y
}

// addBlock adds a block to the head of the queue.
// It does not adjust queue stats.
func (s *stream) addBlock(t int64, data []byte) {
	// If there are no items in output queue, replace sentinel block
	// so that other pointers into queue do not need
	// to change.
	if s.outHead == s.outTail {
		s.outHead.t = t
		s.outHead.data = data
		s.outHead.next = &block{prev: s.outHead} // new sentinel
		s.outTail = s.outHead.next
		return
	}

	// Add a new block just after the sentinel.
	b := &block{
		t:    t,
		data: data,
	}
	b.next = s.outTail
	b.prev = s.outTail.prev

	s.outTail.prev = b
	b.prev.next = b
}

// Remove the block from the front of the queue.
// (assumes that there is such a block to remove)
func (s *stream) removeBlock() {
	b := s.inHead
	s.inHead = b.next
	if s.inHead != nil {
		s.inHead.prev = nil
	}
	// help garbage collector
	b.next = nil
	b.prev = nil
}
<commit_msg>made loopback.Conn.Write return correct count for fragmented messages with error.<commit_after>package loopback

import (
	"sync"
	"io"
	"os"
	"time"
)

type block struct {
	t    int64
	data []byte
	prev *block
	next *block
}

type streamReader stream
type streamWriter stream

func (r *streamReader) Read(data []byte) (int, os.Error) {
	return (*stream)(r).Read(data)
}

func (r *streamReader) Close() os.Error {
	return (*stream)(r).closeInput()
}

func (w *streamWriter) Write(data []byte) (int, os.Error) {
	return (*stream)(w).Write(data)
}

func (w *streamWriter) Close() os.Error {
	return (*stream)(w).closeOutput()
}

type stream struct {
	mu sync.Mutex

	outClosed bool
	inClosed  bool

	outTail     *block // sentinel.
	outHead     *block // also transitTail.
	transitHead *block // also inTail.
	inHead      *block // overall head of list.

	outLimit int // total size of output queue.
	outAvail int // free bytes in output queue.

	inLimit int // total size of input queue.
	inAvail int // free bytes in input queue.

	byteDelay int64
	latency   int64
	mtu       int

	notEmpty sync.Cond
	notFull  sync.Cond
}

// Loopback options for use with Pipe.
type Options struct {
	// ByteDelay controls the time a packet takes in the link.  A packet
	// n bytes long takes ByteDelay * n nanoseconds to exit
	// the output queue and is available for reading Latency
	// nanoseconds later.
	ByteDelay int64
	Latency   int64

	// MTU gives the maximum packet size that can
	// be tranferred atomically across the link.
	// Larger packet will be split.
	// If this is zero, a default of 32768 is assumed
	MTU int

	// InLimit and OutLimit gives the size of the input and output queues.
	// If either is zero, a default of 10*MTU is assumed.
	InLimit  int
	OutLimit int
}

// Pipe creates an asynchronous in-memory pipe,
// Writes are divided into packets of at most opts.MTU bytes
// written to a flow-controlled output queue, transferred across the link,
// and put into an input queue where it is readable with the r.
// The options determine when and how the data will be transferred.
func Pipe(opt Options) (r io.ReadCloser, w io.WriteCloser) {
	if opt.MTU == 0 {
		opt.MTU = 32768
	}
	if opt.InLimit == 0 {
		opt.InLimit = 10 * opt.MTU
	}
	if opt.OutLimit == 0 {
		opt.OutLimit = 10 * opt.MTU
	}
	if opt.InLimit < opt.MTU {
		opt.InLimit = opt.MTU
	}
	if opt.OutLimit < opt.MTU {
		opt.OutLimit = opt.MTU
	}
	sentinel := &block{}
	s := &stream{
		outLimit:    opt.OutLimit,
		outAvail:    opt.OutLimit,
		inLimit:     opt.InLimit,
		inAvail:     opt.InLimit,
		mtu:         opt.MTU,
		byteDelay:   opt.ByteDelay,
		latency:     opt.Latency,
		outTail:     sentinel,
		outHead:     sentinel,
		transitHead: sentinel,
		inHead:      sentinel,
	}
	s.notEmpty.L = &s.mu
	s.notFull.L = &s.mu
	return (*streamReader)(s), (*streamWriter)(s)
}

// Dodgy heuristic:
// If there's stuff in the transit queue that's ready to
// enter the input queue, but the input queue is full
// and it's been waiting for at least latency ns,
// then we block the output queue.
// TODO what do we do about latency for
// blocked packets - as it is a blocked packet
// will incur less latency.
func (s *stream) outBlocked(now int64) bool {
	return s.transitHead != s.outHead &&
		now >= s.transitHead.t+s.latency &&
		s.inAvail < len(s.transitHead.data)
}

func (s *stream) closeInput() os.Error {
	s.mu.Lock()
	s.inClosed = true
	s.notEmpty.Broadcast()
	s.notFull.Broadcast()
	s.mu.Unlock()
	return nil
}

func (s *stream) closeOutput() os.Error {
	s.mu.Lock()
	s.outClosed = true
	s.notEmpty.Broadcast()
	s.notFull.Broadcast()
	s.mu.Unlock()
	return nil
}

func (s *stream) pushLink(now int64) {
	if !s.outBlocked(now) {
		// move blocks from out queue to transit queue.
		for s.outTail != s.outHead && now >= s.outHead.t {
			s.outHead.t += s.latency
			s.outAvail += len(s.outHead.data)
			s.outHead = s.outHead.next
		}
	}
	// move blocks from transit queue to input queue
	for s.transitHead != s.outHead && now >= s.transitHead.t {
		if s.inAvail < len(s.transitHead.data) {
			break // or discard packet
		}
		s.inAvail -= len(s.transitHead.data)
		s.transitHead = s.transitHead.next
	}
}

func (s *stream) Write(data []byte) (int, os.Error) {
	// split the packet into MTU-sized portions if necessary.
	tot := 0
	for len(data) > s.mtu {
		n, err := s.Write(data[0:s.mtu])
		tot += n
		if err != nil {
			return tot, err
		}
		data = data[s.mtu:]
	}
	s.mu.Lock()
	now := time.Nanoseconds()
	for {
		s.pushLink(now)
		if s.outAvail >= len(data) || s.outClosed {
			break
		}
		if s.outBlocked(time.Nanoseconds()) {
			if s.inClosed {
				s.mu.Unlock()
				return 0, os.EPIPE
			}
			s.notFull.Wait()
			continue
		}
		t := s.earliestWriteTime(len(data))
		now = s.sleepUntil(t)
	}
	if s.outClosed {
		s.mu.Unlock()
		return 0, os.EPIPE
	}
	delay := int64(len(data)) * s.byteDelay
	var t int64
	// If there's a block in the queue that's not yet due
	// for transit, then this block leaves delay ns after
	// that one.
	if s.outHead != s.outTail && now < s.outTail.prev.t {
		t = s.outTail.prev.t + delay
	} else {
		t = now + delay
	}
	s.addBlock(t, s.copy(data))
	s.outAvail -= len(data)

	s.notEmpty.Broadcast()
	s.mu.Unlock()
	// TODO runtime.Gosched() ?
	return len(data), nil
}

func (s *stream) Read(buf []byte) (int, os.Error) {
	s.mu.Lock()
	// Loop until there's something to read from the input queue.
	now := time.Nanoseconds()
	for {
		s.pushLink(now)
		if s.inHead != s.transitHead {
			break
		}
		if s.inHead == s.outTail {
			// No data at all in the queue.
			// If the queue is empty and the output queue is closed,
			// then we see EOF.
			if s.outClosed {
				s.mu.Unlock()
				return 0, os.EOF
			}
			s.notEmpty.Wait()
			continue
		}
		now = s.sleepUntil(s.earliestReadTime())
	}
	if s.inClosed {
		// input queue has been forcibly closed:
		// TODO is os.EOF the right error here?
		return 0, os.EOF
	}
	b := s.inHead
	n := copy(buf, b.data)
	b.data = b.data[n:]
	s.inAvail += n
	if len(b.data) == 0 {
		s.removeBlock()
	}
	// Wake up any writers blocked on a full queue.
	s.notFull.Broadcast()
	s.mu.Unlock()
	return n, nil
}

// earliestReadTime returns the earliest time that
// some data might arrive into the input queue.
// It assumes that there is some data in the system.
func (s *stream) earliestReadTime() int64 {
	if s.inAvail < s.inLimit {
		// data is available right now.
		return 0
	}
	if s.transitHead != s.outHead {
		return s.transitHead.t
	}
	if s.outHead != s.outTail {
		return s.outHead.t + s.latency
	}
	panic("no data")
}

// earliestWriteTime returns the earliest time that
// there may be space for n bytes of data to be
// placed into the output queue (it might be later
// if packets are dropped).
func (s *stream) earliestWriteTime(n int) int64 {
	if s.outAvail < s.outLimit {
		// space is available now.
		return 0
	}
	tot := s.outAvail
	for b := s.outHead; b != s.outTail; b = b.next {
		tot += len(b.data)
		if tot >= n {
			return b.t
		}
	}
	panic("write limit exceeded by block size")
}

// sleep until the absolute time t.
// Called with lock held.
func (s *stream) sleepUntil(t int64) int64 {
	now := time.Nanoseconds()
	if now >= t {
		return now
	}
	s.mu.Unlock()
	time.Sleep(t - now)
	s.mu.Lock()
	return time.Nanoseconds()
}

func (s *stream) copy(x []byte) []byte {
	y := make([]byte, len(x))
	copy(y, x)
	return y
}

// addBlock adds a block to the head of the queue.
// It does not adjust queue stats.
func (s *stream) addBlock(t int64, data []byte) {
	// If there are no items in output queue, replace sentinel block
	// so that other pointers into queue do not need
	// to change.
	if s.outHead == s.outTail {
		s.outHead.t = t
		s.outHead.data = data
		s.outHead.next = &block{prev: s.outHead} // new sentinel
		s.outTail = s.outHead.next
		return
	}

	// Add a new block just after the sentinel.
	b := &block{
		t:    t,
		data: data,
	}
	b.next = s.outTail
	b.prev = s.outTail.prev

	s.outTail.prev = b
	b.prev.next = b
}

// Remove the block from the front of the queue.
// (assumes that there is such a block to remove)
func (s *stream) removeBlock() {
	b := s.inHead
	s.inHead = b.next
	if s.inHead != nil {
		s.inHead.prev = nil
	}
	// help garbage collector
	b.next = nil
	b.prev = nil
}
<|endoftext|>


<commit_before>package query

import (
	"context"
	"database/sql"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"
)

// Dump returns a SQL text dump of all rows across all tables, similar to
// sqlite3's dump feature
func Dump(ctx context.Context, tx *sql.Tx, schema string, schemaOnly bool) (string, error) {
	schemas := dumpParseSchema(schema)

	// Begin
	dump := `PRAGMA foreign_keys=OFF;
BEGIN TRANSACTION;
`
	// Schema table
	tableDump, err := dumpTable(ctx, tx, "schema", dumpSchemaTable)
	if err != nil {
		return "", fmt.Errorf("failed to dump table schema: %w", err)
	}
	dump += tableDump

	// All other tables
	tables := make([]string, 0)
	for table := range schemas {
		tables = append(tables, table)
	}
	sort.Strings(tables)
	for _, table := range tables {
		if schemaOnly {
			// Dump only the schema.
			dump += schemas[table] + "\n"
			continue
		}
		tableDump, err := dumpTable(ctx, tx, table, schemas[table])
		if err != nil {
			return "", fmt.Errorf("failed to dump table %s: %w", table, err)
		}
		dump += tableDump
	}

	// Sequences (unless the schemaOnly flag is true)
	if !schemaOnly {
		tableDump, err = dumpTable(ctx, tx, "sqlite_sequence", "DELETE FROM sqlite_sequence;")
		if err != nil {
			return "", fmt.Errorf("failed to dump table sqlite_sequence: %w", err)
		}
		dump += tableDump
	}

	// Commit
	dump += "COMMIT;\n"

	return dump, nil
}

// Return a map from table names to their schema definition, taking a full
// schema SQL text generated with schema.Schema.Dump().
func dumpParseSchema(schema string) map[string]string {
	tables := map[string]string{}
	for _, statement := range strings.Split(schema, ";") {
		statement = strings.Trim(statement, " \n") + ";"
		if !strings.HasPrefix(statement, "CREATE TABLE") {
			continue
		}
		table := strings.Split(statement, " ")[2]
		tables[table] = statement
	}
	return tables
}

// Dump a single table, returning a SQL text containing statements for its
// schema and data.
func dumpTable(ctx context.Context, tx *sql.Tx, table, schema string) (string, error) {
	statements := []string{schema}

	// Query all rows.
	rows, err := tx.QueryContext(ctx, fmt.Sprintf("SELECT * FROM %s ORDER BY rowid", table))
	if err != nil {
		return "", fmt.Errorf("failed to fetch rows: %w", err)
	}
	defer rows.Close()

	// Figure column names
	columns, err := rows.Columns()
	if err != nil {
		return "", fmt.Errorf("failed to get columns: %w", err)
	}

	// Generate an INSERT statement for each row.
	for i := 0; rows.Next(); i++ {
		raw := make([]any, len(columns)) // Raw column values
		row := make([]any, len(columns))
		for i := range raw {
			row[i] = &raw[i]
		}
		err := rows.Scan(row...)
		if err != nil {
			return "", fmt.Errorf("failed to scan row %d: %w", i, err)
		}
		values := make([]string, len(columns))
		for j, v := range raw {
			switch v := v.(type) {
			case int64:
				values[j] = strconv.FormatInt(v, 10)
			case string:
				values[j] = fmt.Sprintf("'%s'", v)
			case []byte:
				values[j] = fmt.Sprintf("'%s'", string(v))
			case time.Time:
				values[j] = strconv.FormatInt(v.Unix(), 10)
			default:
				if v != nil {
					return "", fmt.Errorf("bad type in column %s of row %d", columns[j], i)
				}
				values[j] = "NULL"
			}
		}
		statement := fmt.Sprintf("INSERT INTO %s VALUES(%s);", table, strings.Join(values, ","))
		statements = append(statements, statement)
	}
	return strings.Join(statements, "\n") + "\n", nil
}

// Schema of the schema table.
const dumpSchemaTable = `CREATE TABLE schema (
    id         INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    version    INTEGER NOT NULL,
    updated_at DATETIME NOT NULL,
    UNIQUE (version)
);`
<commit_msg>lxd/db/query: Coding style<commit_after>package query

import (
	"context"
	"database/sql"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"
)

// Dump returns a SQL text dump of all rows across all tables, similar to
// sqlite3's dump feature
func Dump(ctx context.Context, tx *sql.Tx, schema string, schemaOnly bool) (string, error) {
	schemas := dumpParseSchema(schema)

	// Begin
	dump := `PRAGMA foreign_keys=OFF;
BEGIN TRANSACTION;
`
	// Schema table
	tableDump, err := dumpTable(ctx, tx, "schema", dumpSchemaTable)
	if err != nil {
		return "", fmt.Errorf("Failed to dump table schema: %w", err)
	}
	dump += tableDump

	// All other tables
	tables := make([]string, 0)
	for table := range schemas {
		tables = append(tables, table)
	}
	sort.Strings(tables)
	for _, table := range tables {
		if schemaOnly {
			// Dump only the schema.
			dump += schemas[table] + "\n"
			continue
		}

		tableDump, err := dumpTable(ctx, tx, table, schemas[table])
		if err != nil {
			return "", fmt.Errorf("Failed to dump table %s: %w", table, err)
		}

		dump += tableDump
	}

	// Sequences (unless the schemaOnly flag is true)
	if !schemaOnly {
		tableDump, err = dumpTable(ctx, tx, "sqlite_sequence", "DELETE FROM sqlite_sequence;")
		if err != nil {
			return "", fmt.Errorf("Failed to dump table sqlite_sequence: %w", err)
		}

		dump += tableDump
	}

	// Commit
	dump += "COMMIT;\n"

	return dump, nil
}

// Return a map from table names to their schema definition, taking a full
// schema SQL text generated with schema.Schema.Dump().
func dumpParseSchema(schema string) map[string]string {
	tables := map[string]string{}
	for _, statement := range strings.Split(schema, ";") {
		statement = strings.Trim(statement, " \n") + ";"
		if !strings.HasPrefix(statement, "CREATE TABLE") {
			continue
		}
		table := strings.Split(statement, " ")[2]
		tables[table] = statement
	}
	return tables
}

// Dump a single table, returning a SQL text containing statements for its
// schema and data.
func dumpTable(ctx context.Context, tx *sql.Tx, table, schema string) (string, error) {
	statements := []string{schema}

	// Query all rows.
	rows, err := tx.QueryContext(ctx, fmt.Sprintf("SELECT * FROM %s ORDER BY rowid", table))
	if err != nil {
		return "", fmt.Errorf("failed to fetch rows: %w", err)
	}
	defer rows.Close()

	// Figure column names
	columns, err := rows.Columns()
	if err != nil {
		return "", fmt.Errorf("failed to get columns: %w", err)
	}

	// Generate an INSERT statement for each row.
	for i := 0; rows.Next(); i++ {
		raw := make([]any, len(columns)) // Raw column values
		row := make([]any, len(columns))
		for i := range raw {
			row[i] = &raw[i]
		}
		err := rows.Scan(row...)
		if err != nil {
			return "", fmt.Errorf("failed to scan row %d: %w", i, err)
		}
		values := make([]string, len(columns))
		for j, v := range raw {
			switch v := v.(type) {
			case int64:
				values[j] = strconv.FormatInt(v, 10)
			case string:
				values[j] = fmt.Sprintf("'%s'", v)
			case []byte:
				values[j] = fmt.Sprintf("'%s'", string(v))
			case time.Time:
				values[j] = strconv.FormatInt(v.Unix(), 10)
			default:
				if v != nil {
					return "", fmt.Errorf("bad type in column %s of row %d", columns[j], i)
				}
				values[j] = "NULL"
			}
		}
		statement := fmt.Sprintf("INSERT INTO %s VALUES(%s);", table, strings.Join(values, ","))
		statements = append(statements, statement)
	}
	return strings.Join(statements, "\n") + "\n", nil
}

// Schema of the schema table.
const dumpSchemaTable = `CREATE TABLE schema (
    id         INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    version    INTEGER NOT NULL,
    updated_at DATETIME NOT NULL,
    UNIQUE (version)
);`
<|endoftext|>


<commit_before>package lzma

import (
	"errors"
	"fmt"
	"io"
)

// OpFinder enables the support of multiple different OpFinder
// algorithms.
type OpFinder interface {
	findOps(s *State, all bool) []operation
	fmt.Stringer
}

// Writer produces an LZMA stream. EOS requests Close to write an
// end-of-stream marker.
type Writer struct {
	state    *State
	eos      bool
	OpFinder OpFinder
	re       *rangeEncoder
	buf      *buffer
	closed   bool
}

// NewWriter creates a new writer instance.
func NewStreamWriter(pw io.Writer, p Parameters) (w *Writer, err error) {
	if err = p.Verify(); err != nil {
		return
	}
	buf, err := newBuffer(p.DictSize + p.ExtraBufSize)
	if err != nil {
		return nil, err
	}
	d, err := newHashDict(buf, buf.bottom, p.DictSize)
	if err != nil {
		return nil, err
	}
	d.sync()
	state := NewState(p.Properties(), d)
	w = &Writer{
		state:    state,
		eos:      !p.SizeInHeader || p.EOS,
		buf:      buf,
		re:       newRangeEncoder(pw),
		OpFinder: Greedy,
	}
	return w, nil
}

// writeLiteral writes a literal into the operation stream
func (w *Writer) writeLiteral(l lit) error {
	var err error
	state, state2, _ := w.state.states()
	if err = w.state.isMatch[state2].Encode(w.re, 0); err != nil {
		return err
	}
	litState := w.state.litState()
	match := w.state.dict.byteAt(int64(w.state.rep[0]) + 1)
	err = w.state.litCodec.Encode(w.re, l.b, state, match, litState)
	if err != nil {
		return err
	}
	w.state.updateStateLiteral()
	return nil
}

// iverson implements the Iverson operator as proposed by Donald Knuth in his
// book Concrete Mathematics.
func iverson(ok bool) uint32 {
	if ok {
		return 1
	}
	return 0
}

// writeMatch writes a repetition operation into the operation stream
func (w *Writer) writeMatch(m match) error {
	var err error
	if !(minDistance <= m.distance && m.distance <= maxDistance) {
		panic(errDistRange)
	}
	dist := uint32(m.distance - minDistance)
	if !(MinLength <= m.n && m.n <= MaxLength) &&
		!(dist == w.state.rep[0] && m.n == 1) {
		panic(errLenRange)
	}
	state, state2, posState := w.state.states()
	if err = w.state.isMatch[state2].Encode(w.re, 1); err != nil {
		return err
	}
	g := 0
	for ; g < 4; g++ {
		if w.state.rep[g] == dist {
			break
		}
	}
	b := iverson(g < 4)
	if err = w.state.isRep[state].Encode(w.re, b); err != nil {
		return err
	}
	n := uint32(m.n - MinLength)
	if b == 0 {
		// simple match
		w.state.rep[3], w.state.rep[2], w.state.rep[1], w.state.rep[0] =
			w.state.rep[2], w.state.rep[1], w.state.rep[0], dist
		w.state.updateStateMatch()
		if err = w.state.lenCodec.Encode(w.re, n, posState); err != nil {
			return err
		}
		return w.state.distCodec.Encode(w.re, dist, n)
	}
	b = iverson(g != 0)
	if err = w.state.isRepG0[state].Encode(w.re, b); err != nil {
		return err
	}
	if b == 0 {
		// g == 0
		b = iverson(m.n != 1)
		if err = w.state.isRepG0Long[state2].Encode(w.re, b); err != nil {
			return err
		}
		if b == 0 {
			w.state.updateStateShortRep()
			return nil
		}
	} else {
		// g in {1,2,3}
		b = iverson(g != 1)
		if err = w.state.isRepG1[state].Encode(w.re, b); err != nil {
			return err
		}
		if b == 1 {
			// g in {2,3}
			b = iverson(g != 2)
			err = w.state.isRepG2[state].Encode(w.re, b)
			if err != nil {
				return err
			}
			if b == 1 {
				w.state.rep[3] = w.state.rep[2]
			}
			w.state.rep[2] = w.state.rep[1]
		}
		w.state.rep[1] = w.state.rep[0]
		w.state.rep[0] = dist
	}
	w.state.updateStateRep()
	return w.state.repLenCodec.Encode(w.re, n, posState)
}

// writeOp writes an operation value into the stream. It checks whether there
// is still enough space available using an upper limit for the size required.
func (w *Writer) writeOp(op operation) error {
	var err error
	switch x := op.(type) {
	case match:
		err = w.writeMatch(x)
	case lit:
		err = w.writeLiteral(x)
	}
	if err != nil {
		return err
	}
	err = w.discard(op)
	return err
}

// discard processes an operation after it has been written into the
// compressed LZMA street by moving the dictionary head forward.
func (w *Writer) discard(op operation) error {
	k := op.Len()
	n, err := w.state.dict.(*hashDict).move(k)
	if err != nil {
		return fmt.Errorf("operation %s: move %d error %s", op, k, err)
	}
	if n < k {
		return fmt.Errorf("operation %s: move %d incomplete", op, k)
	}
	return nil
}

// compress does the actual compression. If all is set all data
// available will be compressed.
func (w *Writer) compress(all bool) error {
	ops := w.OpFinder.findOps(w.state, all)
	for _, op := range ops {
		if err := w.writeOp(op); err != nil {
			return err
		}
	}
	w.state.dict.(*hashDict).sync()
	return nil
}

// errWriterClosed indicates that a writer has been closed once before.
var errWriterClosed = errors.New("writer is closed")

// Write puts the provided data into the writer.
func (w *Writer) Write(p []byte) (n int, err error) {
	if w.closed {
		return 0, errWriterClosed
	}
	for len(p) > 0 {
		var k int
		k, err = w.buf.Write(p)
		n += k
		if err != errLimit {
			return
		}
		p = p[k:]
		if err = w.compress(false); err != nil {
			return
		}
	}
	return
}

// This operation will be encoded to indicate that the stream has ended.
var eosMatch = match{distance: maxDistance, n: MinLength}

// Close closes the writer.
func (w *Writer) Close() (err error) {
	if w.closed {
		return errWriterClosed
	}
	w.closed = true
	if err = w.compress(true); err != nil {
		return err
	}
	if w.eos {
		if err = w.writeMatch(eosMatch); err != nil {
			return err
		}
	}
	if err = w.re.Close(); err != nil {
		return err
	}
	return nil
}
<commit_msg>lzma: fixed documentation comment for NewStreamWriter<commit_after>package lzma

import (
	"errors"
	"fmt"
	"io"
)

// OpFinder enables the support of multiple different OpFinder
// algorithms.
type OpFinder interface {
	findOps(s *State, all bool) []operation
	fmt.Stringer
}

// Writer produces an LZMA stream. EOS requests Close to write an
// end-of-stream marker.
type Writer struct {
	state    *State
	eos      bool
	OpFinder OpFinder
	re       *rangeEncoder
	buf      *buffer
	closed   bool
}

// NewStreamWriter creates a new writer instance.
func NewStreamWriter(pw io.Writer, p Parameters) (w *Writer, err error) {
	if err = p.Verify(); err != nil {
		return
	}
	buf, err := newBuffer(p.DictSize + p.ExtraBufSize)
	if err != nil {
		return nil, err
	}
	d, err := newHashDict(buf, buf.bottom, p.DictSize)
	if err != nil {
		return nil, err
	}
	d.sync()
	state := NewState(p.Properties(), d)
	w = &Writer{
		state:    state,
		eos:      !p.SizeInHeader || p.EOS,
		buf:      buf,
		re:       newRangeEncoder(pw),
		OpFinder: Greedy,
	}
	return w, nil
}

// writeLiteral writes a literal into the operation stream
func (w *Writer) writeLiteral(l lit) error {
	var err error
	state, state2, _ := w.state.states()
	if err = w.state.isMatch[state2].Encode(w.re, 0); err != nil {
		return err
	}
	litState := w.state.litState()
	match := w.state.dict.byteAt(int64(w.state.rep[0]) + 1)
	err = w.state.litCodec.Encode(w.re, l.b, state, match, litState)
	if err != nil {
		return err
	}
	w.state.updateStateLiteral()
	return nil
}

// iverson implements the Iverson operator as proposed by Donald Knuth in his
// book Concrete Mathematics.
func iverson(ok bool) uint32 {
	if ok {
		return 1
	}
	return 0
}

// writeMatch writes a repetition operation into the operation stream
func (w *Writer) writeMatch(m match) error {
	var err error
	if !(minDistance <= m.distance && m.distance <= maxDistance) {
		panic(errDistRange)
	}
	dist := uint32(m.distance - minDistance)
	if !(MinLength <= m.n && m.n <= MaxLength) &&
		!(dist == w.state.rep[0] && m.n == 1) {
		panic(errLenRange)
	}
	state, state2, posState := w.state.states()
	if err = w.state.isMatch[state2].Encode(w.re, 1); err != nil {
		return err
	}
	g := 0
	for ; g < 4; g++ {
		if w.state.rep[g] == dist {
			break
		}
	}
	b := iverson(g < 4)
	if err = w.state.isRep[state].Encode(w.re, b); err != nil {
		return err
	}
	n := uint32(m.n - MinLength)
	if b == 0 {
		// simple match
		w.state.rep[3], w.state.rep[2], w.state.rep[1], w.state.rep[0] =
			w.state.rep[2], w.state.rep[1], w.state.rep[0], dist
		w.state.updateStateMatch()
		if err = w.state.lenCodec.Encode(w.re, n, posState); err != nil {
			return err
		}
		return w.state.distCodec.Encode(w.re, dist, n)
	}
	b = iverson(g != 0)
	if err = w.state.isRepG0[state].Encode(w.re, b); err != nil {
		return err
	}
	if b == 0 {
		// g == 0
		b = iverson(m.n != 1)
		if err = w.state.isRepG0Long[state2].Encode(w.re, b); err != nil {
			return err
		}
		if b == 0 {
			w.state.updateStateShortRep()
			return nil
		}
	} else {
		// g in {1,2,3}
		b = iverson(g != 1)
		if err = w.state.isRepG1[state].Encode(w.re, b); err != nil {
			return err
		}
		if b == 1 {
			// g in {2,3}
			b = iverson(g != 2)
			err = w.state.isRepG2[state].Encode(w.re, b)
			if err != nil {
				return err
			}
			if b == 1 {
				w.state.rep[3] = w.state.rep[2]
			}
			w.state.rep[2] = w.state.rep[1]
		}
		w.state.rep[1] = w.state.rep[0]
		w.state.rep[0] = dist
	}
	w.state.updateStateRep()
	return w.state.repLenCodec.Encode(w.re, n, posState)
}

// writeOp writes an operation value into the stream. It checks whether there
// is still enough space available using an upper limit for the size required.
func (w *Writer) writeOp(op operation) error {
	var err error
	switch x := op.(type) {
	case match:
		err = w.writeMatch(x)
	case lit:
		err = w.writeLiteral(x)
	}
	if err != nil {
		return err
	}
	err = w.discard(op)
	return err
}

// discard processes an operation after it has been written into the
// compressed LZMA street by moving the dictionary head forward.
func (w *Writer) discard(op operation) error {
	k := op.Len()
	n, err := w.state.dict.(*hashDict).move(k)
	if err != nil {
		return fmt.Errorf("operation %s: move %d error %s", op, k, err)
	}
	if n < k {
		return fmt.Errorf("operation %s: move %d incomplete", op, k)
	}
	return nil
}

// compress does the actual compression. If all is set all data
// available will be compressed.
func (w *Writer) compress(all bool) error {
	ops := w.OpFinder.findOps(w.state, all)
	for _, op := range ops {
		if err := w.writeOp(op); err != nil {
			return err
		}
	}
	w.state.dict.(*hashDict).sync()
	return nil
}

// errWriterClosed indicates that a writer has been closed once before.
var errWriterClosed = errors.New("writer is closed")

// Write puts the provided data into the writer.
func (w *Writer) Write(p []byte) (n int, err error) {
	if w.closed {
		return 0, errWriterClosed
	}
	for len(p) > 0 {
		var k int
		k, err = w.buf.Write(p)
		n += k
		if err != errLimit {
			return
		}
		p = p[k:]
		if err = w.compress(false); err != nil {
			return
		}
	}
	return
}

// This operation will be encoded to indicate that the stream has ended.
var eosMatch = match{distance: maxDistance, n: MinLength}

// Close closes the writer.
func (w *Writer) Close() (err error) {
	if w.closed {
		return errWriterClosed
	}
	w.closed = true
	if err = w.compress(true); err != nil {
		return err
	}
	if w.eos {
		if err = w.writeMatch(eosMatch); err != nil {
			return err
		}
	}
	if err = w.re.Close(); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package jsonstore

import (
	"io/ioutil"
	"os"
	"regexp"
	"strconv"
	"testing"
)

func testFile() *os.File {
	f, err := ioutil.TempFile(".", "jsonstore")
	if err != nil {
		panic(err)
	}
	return f
}

func TestOpen(t *testing.T) {
	f := testFile()
	defer os.Remove(f.Name())
	ioutil.WriteFile(f.Name(), []byte(`{"hello":"world"}`), 0644)
	ks, err := Open(f.Name())
	if err != nil {
		t.Error(err)
	}
	if len(ks.Data) != 1 {
		t.Errorf("expected %d got %d", 1, len(ks.Data))
	}
	if world, ok := ks.Data["hello"]; !ok || string(world) != `world` {
		t.Errorf("expected %s got %s", "world", world)
	}
}

func TestGeneral(t *testing.T) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	err := ks.Set("hello", "world")
	if err != nil {
		t.Error(err)
	}
	if err = Save(ks, f.Name()); err != nil {
		t.Error(err)
	}

	ks2, _ := Open(f.Name())
	var a string
	var b string
	ks.Get("hello", &a)
	ks2.Get("hello", &b)
	if a != b {
		t.Errorf("expected '%s' got '%s'", a, b)
	}

	// Set a object, using a Gzipped JSON
	type Human struct {
		Name   string
		Height float64
	}
	ks.Set("human:1", Human{"Dante", 5.4})
	Save(ks, "test2.json.gz")
	ks2, _ = Open("test2.json.gz")
	var human Human
	ks2.Get("human:1", &human)
	if human.Height != 5.4 {
		t.Errorf("expected '%v', got '%v'", Human{"Dante", 5.4}, human)
	}
}

func TestRegex(t *testing.T) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	ks.Set("hello:1", "world1")
	ks.Set("hello:2", "world2")
	ks.Set("hello:3", "world3")
	ks.Set("world:1", "hello1")
	if len(ks.GetAll(regexp.MustCompile(`hello`))) != len(ks.Keys())-1 {
		t.Errorf("Problem getting all")
	}
}

func BenchmarkOpenBig(b *testing.B) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	for i := 1; i < 1000; i++ {
		ks.Set("hello:"+strconv.Itoa(i), "world"+strconv.Itoa(i))
	}
	Save(ks, f.Name())

	var err error
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ks, err = Open(f.Name())
		if err != nil {
			panic(err)
		}
	}
	Save(ks, f.Name())
}

func BenchmarkOpenOldBig(b *testing.B) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	for i := 1; i < 1000; i++ {
		ks.Set("hello:"+strconv.Itoa(i), "world"+strconv.Itoa(i))
	}
	Save(ks, f.Name())

	var err error
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ks, err = OpenOld(f.Name())
		if err != nil {
			panic(err)
		}
	}
	Save(ks, f.Name())
}

func BenchmarkOpenSmall(b *testing.B) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	for i := 1; i < 10; i++ {
		ks.Set("hello:"+strconv.Itoa(i), "world"+strconv.Itoa(i))
	}
	Save(ks, f.Name())

	var err error
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ks, err = Open(f.Name())
		if err != nil {
			panic(err)
		}
	}
	Save(ks, f.Name())
}

func BenchmarkOpenOldSmall(b *testing.B) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	for i := 1; i < 10; i++ {
		ks.Set("hello:"+strconv.Itoa(i), "world"+strconv.Itoa(i))
	}
	Save(ks, f.Name())

	var err error
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ks, err = OpenOld(f.Name())
		if err != nil {
			panic(err)
		}
	}
	Save(ks, f.Name())
}

func BenchmarkGet(b *testing.B) {
	ks := new(JSONStore)
	err := ks.Set("human:1", Human{"Dante", 5.4})
	if err != nil {
		panic(err)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		var human Human
		ks.Get("human:1", &human)
	}
}

type Human struct {
	Name   string
	Height float64
}

func BenchmarkSet(b *testing.B) {
	ks := new(JSONStore)
	b.ResetTimer()
	// set a key to any object you want
	for i := 0; i < b.N; i++ {
		err := ks.Set("human:1", Human{"Dante", 5.4})
		if err != nil {
			panic(err)
		}
	}
}

func BenchmarkSave(b *testing.B) {
	ks := new(JSONStore)
	ks.Set("data", 1234)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		Save(ks, "benchmark.json.gz")
	}
}
<commit_msg>Better tests<commit_after>package jsonstore

import (
	"io/ioutil"
	"os"
	"regexp"
	"strconv"
	"testing"
)

type Human struct {
	Name   string
	Height float64
}

func testFile() *os.File {
	f, err := ioutil.TempFile(".", "jsonstore")
	if err != nil {
		panic(err)
	}
	return f
}

func TestOpen(t *testing.T) {
	f := testFile()
	defer os.Remove(f.Name())
	ioutil.WriteFile(f.Name(), []byte(`{"hello":"world"}`), 0644)
	ks, err := Open(f.Name())
	if err != nil {
		t.Error(err)
	}
	if len(ks.Data) != 1 {
		t.Errorf("expected %d got %d", 1, len(ks.Data))
	}
	if world, ok := ks.Data["hello"]; !ok || string(world) != `world` {
		t.Errorf("expected %s got %s", "world", world)
	}
}

func TestGeneral(t *testing.T) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	err := ks.Set("hello", "world")
	if err != nil {
		t.Error(err)
	}
	if err = Save(ks, f.Name()); err != nil {
		t.Error(err)
	}

	ks2, _ := Open(f.Name())
	var a string
	var b string
	ks.Get("hello", &a)
	ks2.Get("hello", &b)
	if a != b {
		t.Errorf("expected '%s' got '%s'", a, b)
	}

	// Set a object, using a Gzipped JSON
	type Human struct {
		Name   string
		Height float64
	}
	ks.Set("human:1", Human{"Dante", 5.4})
	Save(ks, "test2.json.gz")
	Save(ks, "test2.json")
	var human Human

	ks2, err = Open("test2.json")
	if err != nil {
		t.Errorf(err.Error())
	}
	ks2.Get("human:1", &human)
	if human.Height != 5.4 {
		t.Errorf("expected '%v', got '%v'", Human{"Dante", 5.4}, human)
	}

	ks2, err = Open("test2.json.gz")
	if err != nil {
		t.Errorf(err.Error())
	}
	ks2.Get("human:1", &human)
	if human.Height != 5.4 {
		t.Errorf("expected '%v', got '%v'", Human{"Dante", 5.4}, human)
	}

}

func TestRegex(t *testing.T) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	ks.Set("hello:1", "world1")
	ks.Set("hello:2", "world2")
	ks.Set("hello:3", "world3")
	ks.Set("world:1", "hello1")
	if len(ks.GetAll(regexp.MustCompile(`hello`))) != len(ks.Keys())-1 {
		t.Errorf("Problem getting all")
	}
}

func BenchmarkOpenBig(b *testing.B) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	for i := 1; i < 1000; i++ {
		ks.Set("hello:"+strconv.Itoa(i), "world"+strconv.Itoa(i))
	}
	Save(ks, f.Name())

	var err error
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ks, err = Open(f.Name())
		if err != nil {
			panic(err)
		}
	}
	Save(ks, f.Name())
}

func BenchmarkOpenOldBig(b *testing.B) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	for i := 1; i < 1000; i++ {
		ks.Set("hello:"+strconv.Itoa(i), "world"+strconv.Itoa(i))
	}
	Save(ks, f.Name())

	var err error
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ks, err = OpenOld(f.Name())
		if err != nil {
			panic(err)
		}
	}
	Save(ks, f.Name())
}

func BenchmarkOpenSmall(b *testing.B) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	for i := 1; i < 10; i++ {
		ks.Set("hello:"+strconv.Itoa(i), "world"+strconv.Itoa(i))
	}
	Save(ks, f.Name())

	var err error
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ks, err = Open(f.Name())
		if err != nil {
			panic(err)
		}
	}
	Save(ks, f.Name())
}

func BenchmarkOpenOldSmall(b *testing.B) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	for i := 1; i < 10; i++ {
		ks.Set("hello:"+strconv.Itoa(i), "world"+strconv.Itoa(i))
	}
	Save(ks, f.Name())

	var err error
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ks, err = OpenOld(f.Name())
		if err != nil {
			panic(err)
		}
	}
	Save(ks, f.Name())
}

func BenchmarkOpen(b *testing.B) {
	f := testFile()
	defer os.Remove(f.Name())
	ks := new(JSONStore)
	for i := 1; i < 100; i++ {
		ks.Set("hello:"+strconv.Itoa(i), "world"+strconv.Itoa(i))
	}
	Save(ks, f.Name())

	var err error
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ks, err = Open(f.Name())
		if err != nil {
			panic(err)
		}
	}
	Save(ks, f.Name())
}

func BenchmarkGet(b *testing.B) {
	ks := new(JSONStore)
	err := ks.Set("human:1", Human{"Dante", 5.4})
	if err != nil {
		panic(err)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		var human Human
		ks.Get("human:1", &human)
	}
}

func BenchmarkSet(b *testing.B) {
	ks := new(JSONStore)
	b.ResetTimer()
	// set a key to any object you want
	for i := 0; i < b.N; i++ {
		err := ks.Set("human:1", Human{"Dante", 5.4})
		if err != nil {
			panic(err)
		}
	}
}

func BenchmarkSave(b *testing.B) {
	ks := new(JSONStore)
	ks.Set("data", 1234)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		Save(ks, "benchmark.json.gz")
	}
}
<|endoftext|>