<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/ipdcode/containerfs/logger"
	dp "github.com/ipdcode/containerfs/proto/dp"
	vp "github.com/ipdcode/containerfs/proto/vp"
	"github.com/ipdcode/containerfs/utils"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
	"net"
	"os"
	"io"
	"runtime"
	"runtime/debug"
	"strconv"
	"sync"
	"time"
)

// DataNodeServer ...
type DataNodeServer struct {
	Mutex sync.Mutex
}

type addr struct {
	Ipnr  net.IP
	IPInt int32
	IPStr string
	Port  int32
	Path  string
	Flag  string
	Log   string

	VolMgrHost string
}

// DataNodeServerAddr ...
var DataNodeServerAddr addr

func startDataService() {

	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", DataNodeServerAddr.Port))
	if err != nil {
		panic(fmt.Sprintf("Failed to listen on:%v", DataNodeServerAddr.Port))
	}
	s := grpc.NewServer()
	dp.RegisterDataNodeServer(s, &DataNodeServer{})
	reflection.Register(s)
	if err := s.Serve(lis); err != nil {
		panic("Failed to serve")
	}
}

func registryToVolMgr() {
	conn, err := grpc.Dial(DataNodeServerAddr.VolMgrHost, grpc.WithInsecure())
	if err != nil {
		logger.Debug("data node statup failed : Dial to volmgr failed !")
		os.Exit(1)
	}
	defer conn.Close()
	c := vp.NewVolMgrClient(conn)

	var datanodeRegistryReq vp.DatanodeRegistryReq
	datanodeRegistryReq.Ip = DataNodeServerAddr.IPInt
	datanodeRegistryReq.Port = DataNodeServerAddr.Port
	diskInfo := utils.DiskUsage(DataNodeServerAddr.Path)
	capacity := int32(float64(diskInfo.All) / float64(1024*1024*1024))
	datanodeRegistryReq.Capacity = capacity
	datanodeRegistryReq.MountPoint = DataNodeServerAddr.Path

	_, err = os.Stat(datanodeRegistryReq.MountPoint)
	if err != nil {
		logger.Error("data node statup failed : DataNodeServerAddr.Path not exist !")
		os.Exit(1)
	}

	pDatanodeRegistryAck, _ := c.DatanodeRegistry(context.Background(), &datanodeRegistryReq)
	if pDatanodeRegistryAck.Ret == 0 {
		logger.Debug("registry success!")
		os.Create(DataNodeServerAddr.Flag)
	} else {
		logger.Debug("data node statup failed : registry to volmgr failed !")
		os.Exit(1)
	}

	return

}

func heartbeatToVolMgr() {

	conn, err := grpc.Dial(DataNodeServerAddr.VolMgrHost, grpc.WithInsecure())
	if err != nil {
		logger.Debug("HearBeat failed : Dial to volmgr failed !")
	}
	defer conn.Close()
	c := vp.NewVolMgrClient(conn)

	diskInfo := utils.DiskUsage(DataNodeServerAddr.Path)
	free := int32(float64(diskInfo.Free) / float64(1024*1024*1024))
	used := int32(float64(diskInfo.Used) / float64(1024*1024*1024))

	var datanodeHeartbeatReq vp.DatanodeHeartbeatReq
	datanodeHeartbeatReq.Ip = DataNodeServerAddr.IPInt
	datanodeHeartbeatReq.Port = DataNodeServerAddr.Port
	datanodeHeartbeatReq.Free = free
	datanodeHeartbeatReq.Used = used
	//datanodeHeartbeatReq.Status = 0

	f, err := os.OpenFile(DataNodeServerAddr.Path+"/health", os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0666)
	if err != nil {
		logger.Error("Open datanode check health file error:%v", err)
		datanodeHeartbeatReq.Status = 2
	}
	defer f.Close()

	_, err = f.WriteString("ok")
	if err != nil {
		logger.Error("Write datanode check health file error:%v", err)
		datanodeHeartbeatReq.Status = 2
	} else {
		datanodeHeartbeatReq.Status = 0
	}

	c.DatanodeHeartbeat(context.Background(), &datanodeHeartbeatReq)
}

// DatanodeHealthCheck rpc GetChunks(GetChunksReq) returns (GetChunksAck){};
func (s *DataNodeServer) DatanodeHealthCheck(ctx context.Context, in *dp.DatanodeHealthCheckReq) (*dp.DatanodeHealthCheckAck, error) {
	ack := dp.DatanodeHealthCheckAck{}
	ack.Ret = 1
	return &ack, nil
}

// WriteChunk ...
func (s *DataNodeServer) WriteChunk(ctx context.Context, in *dp.WriteChunkReq) (*dp.WriteChunkAck, error) {
	var f *os.File
	var err error

	ack := dp.WriteChunkAck{}
	chunkID := in.ChunkID
	blockID := in.BlockID

	path := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID))
	if ok, err := utils.LocalPathExists(path); !ok && err == nil {
		os.MkdirAll(path, 0777)
	}

	chunkFileName := path + "/chunk-" + strconv.Itoa(int(chunkID))

	f, err = os.OpenFile(chunkFileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0660)
	defer f.Close()
	if err != nil {
		ack.Ret = -1
		return &ack, nil
	}
	w := bufio.NewWriter(f)
	w.Write(in.Databuf)
	w.Flush()

	ack.Ret = 0
	return &ack, nil
}

/*rpc WriteChunkStream(stream WriteChunkReq) returns (WriteChunkAck){}; */
/*
func (s *DataNodeServer) WriteChunkStream(stream dp.DataNode_WriteChunkStreamServer) error {

	var f *os.File
	var err error
	//fmt.Println("writechunking in datanode ...")
	ack := dp.WriteChunkAck{}
	if err != nil {
		ack.Ret = -1
		return stream.SendAndClose(&ack)
	}
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			ack.Ret = 0
			return stream.SendAndClose(&ack)
		}
		if err != nil {
			ack.Ret = 1
			return stream.SendAndClose(&ack)
		}
		chunkID := in.ChunkID
		blockID := in.BlockID
		chunkFileName := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID)) + "/chunk-" + strconv.Itoa(int(chunkID))
		f, err = os.OpenFile(chunkFileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		f.WriteString(in.Databuf)
		f.Close()
	}
}
*/
/*
func (s *DataNodeServer) ReadChunk(ctx context.Context, in *dp.ReadChunkReq) (*dp.ReadChunkAck, error) {
	ack := dp.ReadChunkAck{}
	chunkID := in.ChunkID
	blockID := in.BlockID
	readsize := in.Readlen
	offset := in.Offset

	chunkFileName := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID)) + "/chunk-" + strconv.Itoa(int(chunkID))
	f, err := os.Open(chunkFileName)
	defer f.Close()
	if err != nil {
		ack.Ret = -1
		return &ack, nil
	}
	_, err = f.Seek(offset, 0)
	if err != nil {
		ack.Ret = -1
		return &ack, nil
	}

	buf := make([]byte, readsize)
	bfRd := bufio.NewReader(f)
	for {
		n, err := bfRd.Read(buf)
		if err != nil {
			ack.Ret = -1
			return &ack, nil
		}
		ack.Databuf = utils.B2S(buf)
		ack.Ret = 1
		ack.Readsize = int64(n)
		return &ack, nil
	}
}
*/

// StreamReadChunk ...
func (s *DataNodeServer) StreamReadChunk(in *dp.StreamReadChunkReq, stream dp.DataNode_StreamReadChunkServer) error {
	chunkID := in.ChunkID
	blockID := in.BlockID
	//offset := in.Offset
	readsize := in.Readsize

	chunkFileName := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID)) + "/chunk-" + strconv.Itoa(int(chunkID))
	f, err := os.Open(chunkFileName)
	defer f.Close()
	if err != nil {
		return err
	}

	var ack dp.StreamReadChunkAck
	var totalsize int64
	buf := make([]byte, 2*1024*1024)

	bfRd := bufio.NewReader(f)
	for {
		n, err := bfRd.Read(buf)
		if err != nil && err != io.EOF {
			logger.Error("read chunkfile:%v error:%v",chunkFileName,err)
			return err
		}

		if n == 0 {
            logger.Debug("read chunkfile:%v endsize:%v",chunkFileName,totalsize)
            break
        }
		totalsize += int64(n)

		ack.Databuf = buf[:n]
		if err := stream.Send(&ack); err != nil {
			logger.Error("Send stream data to fuse error:%v", err)
			return err
		}

		if totalsize == readsize {
            logger.Debug("Request chunkFileName:%v size:%v equal real ret size:%v", chunkFileName, readsize, totalsize)
            break
        }
	}

	return nil
}

//DeleteChunk rpc DeleteChunks(eleteChunksReq) returns (eleteChunksAck){};
func (s *DataNodeServer) DeleteChunk(ctx context.Context, in *dp.DeleteChunkReq) (*dp.DeleteChunkAck, error) {
	var err error

	ack := dp.DeleteChunkAck{}
	chunkID := in.ChunkID
	blockID := in.BlockID

	chunkFileName := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID)) + "/chunk-" + strconv.Itoa(int(chunkID))

	err = os.Remove(chunkFileName)
	if err != nil {
		ack.Ret = 0
	} else {
		ack.Ret = 0
	}
	ack.Ret = 0
	return &ack, nil
}

func init() {

	var loglevel string
	var port int

	flag.StringVar(&DataNodeServerAddr.IPStr, "host", "127.0.0.1", "ContainerFS DataNode Host")
	flag.IntVar(&port, "port", 8000, "ContainerFS DataNode Port")
	flag.StringVar(&DataNodeServerAddr.Path, "datapath", "", "ContainerFS DataNode Data Path")
	flag.StringVar(&DataNodeServerAddr.VolMgrHost, "volmgr", "127.0.0.1:7000", "ContainerFS VolMgr Host")
	flag.StringVar(&DataNodeServerAddr.Log, "logpath", "/export/Logs/containerfs/logs/", "ContainerFS Log Path")
	flag.StringVar(&loglevel, "loglevel", "error", "ContainerFS Log Level")

	flag.Parse()

	DataNodeServerAddr.Port = int32(port)
	ipnr := net.ParseIP(DataNodeServerAddr.IPStr)
	DataNodeServerAddr.Ipnr = ipnr
	ipint := utils.InetAton(ipnr)
	DataNodeServerAddr.IPInt = ipint
	DataNodeServerAddr.Flag = DataNodeServerAddr.Path + "/.registryflag"

	logger.SetConsole(true)
	logger.SetRollingFile(DataNodeServerAddr.Log, "datanode.log", 10, 100, logger.MB) //each 100M rolling

	switch loglevel {
	case "error":
		logger.SetLevel(logger.ERROR)
	case "debug":
		logger.SetLevel(logger.DEBUG)
	case "info":
		logger.SetLevel(logger.INFO)
	default:
		logger.SetLevel(logger.ERROR)
	}

	if ok, _ := utils.LocalPathExists(DataNodeServerAddr.Flag); !ok {
		logger.Debug("Start registry to volmgr ...")
		registryToVolMgr()
		logger.Debug("registry to volmgr success")
	} else {
		logger.Debug("already registied")
	}

}

func main() {

	numCPU := runtime.NumCPU()
	runtime.GOMAXPROCS(numCPU)

	defer func() {
		if err := recover(); err != nil {
			logger.Error("panic !!! :%v", err)
			logger.Error("stacks:%v", string(debug.Stack()))
		}
	}()

	heartbeatToVolMgr()
	ticker := time.NewTicker(time.Second * 60)
	go func() {
		for range ticker.C {
			heartbeatToVolMgr()
		}
	}()
	startDataService()
}
<commit_msg>fix bug<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/ipdcode/containerfs/logger"
	dp "github.com/ipdcode/containerfs/proto/dp"
	vp "github.com/ipdcode/containerfs/proto/vp"
	"github.com/ipdcode/containerfs/utils"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
	"net"
	"os"
	"io"
	"runtime"
	"runtime/debug"
	"strconv"
	"sync"
	"time"
)

// DataNodeServer ...
type DataNodeServer struct {
	Mutex sync.Mutex
}

type addr struct {
	Ipnr  net.IP
	IPInt int32
	IPStr string
	Port  int32
	Path  string
	Flag  string
	Log   string

	VolMgrHost string
}

// DataNodeServerAddr ...
var DataNodeServerAddr addr

func startDataService() {

	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", DataNodeServerAddr.Port))
	if err != nil {
		panic(fmt.Sprintf("Failed to listen on:%v", DataNodeServerAddr.Port))
	}
	s := grpc.NewServer()
	dp.RegisterDataNodeServer(s, &DataNodeServer{})
	reflection.Register(s)
	if err := s.Serve(lis); err != nil {
		panic("Failed to serve")
	}
}

func registryToVolMgr() {
	conn, err := grpc.Dial(DataNodeServerAddr.VolMgrHost, grpc.WithInsecure())
	if err != nil {
		logger.Debug("data node statup failed : Dial to volmgr failed !")
		os.Exit(1)
	}
	defer conn.Close()
	c := vp.NewVolMgrClient(conn)

	var datanodeRegistryReq vp.DatanodeRegistryReq
	datanodeRegistryReq.Ip = DataNodeServerAddr.IPInt
	datanodeRegistryReq.Port = DataNodeServerAddr.Port
	diskInfo := utils.DiskUsage(DataNodeServerAddr.Path)
	capacity := int32(float64(diskInfo.All) / float64(1024*1024*1024))
	datanodeRegistryReq.Capacity = capacity
	datanodeRegistryReq.MountPoint = DataNodeServerAddr.Path

	_, err = os.Stat(datanodeRegistryReq.MountPoint)
	if err != nil {
		logger.Error("data node statup failed : DataNodeServerAddr.Path not exist !")
		os.Exit(1)
	}

	pDatanodeRegistryAck, _ := c.DatanodeRegistry(context.Background(), &datanodeRegistryReq)
	if pDatanodeRegistryAck.Ret == 0 {
		logger.Debug("registry success!")
		os.Create(DataNodeServerAddr.Flag)
	} else {
		logger.Debug("data node statup failed : registry to volmgr failed !")
		os.Exit(1)
	}

	return

}

func heartbeatToVolMgr() {

	conn, err := grpc.Dial(DataNodeServerAddr.VolMgrHost, grpc.WithInsecure())
	if err != nil {
		logger.Debug("HearBeat failed : Dial to volmgr failed !")
	}
	defer conn.Close()
	c := vp.NewVolMgrClient(conn)

	diskInfo := utils.DiskUsage(DataNodeServerAddr.Path)
	free := int32(float64(diskInfo.Free) / float64(1024*1024*1024))
	used := int32(float64(diskInfo.Used) / float64(1024*1024*1024))

	var datanodeHeartbeatReq vp.DatanodeHeartbeatReq
	datanodeHeartbeatReq.Ip = DataNodeServerAddr.IPInt
	datanodeHeartbeatReq.Port = DataNodeServerAddr.Port
	datanodeHeartbeatReq.Free = free
	datanodeHeartbeatReq.Used = used
	//datanodeHeartbeatReq.Status = 0

	f, err := os.OpenFile(DataNodeServerAddr.Path+"/health", os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0666)
	if err != nil {
		logger.Error("Open datanode check health file error:%v", err)
		datanodeHeartbeatReq.Status = 2
	}
	defer f.Close()

	_, err = f.WriteString("ok")
	if err != nil {
		logger.Error("Write datanode check health file error:%v", err)
		datanodeHeartbeatReq.Status = 2
	} else {
		datanodeHeartbeatReq.Status = 0
	}

	c.DatanodeHeartbeat(context.Background(), &datanodeHeartbeatReq)
}

// DatanodeHealthCheck rpc GetChunks(GetChunksReq) returns (GetChunksAck){};
func (s *DataNodeServer) DatanodeHealthCheck(ctx context.Context, in *dp.DatanodeHealthCheckReq) (*dp.DatanodeHealthCheckAck, error) {
	ack := dp.DatanodeHealthCheckAck{}
	ack.Ret = 1
	return &ack, nil
}

// WriteChunk ...
func (s *DataNodeServer) WriteChunk(ctx context.Context, in *dp.WriteChunkReq) (*dp.WriteChunkAck, error) {
	var f *os.File
	var err error

	ack := dp.WriteChunkAck{}
	chunkID := in.ChunkID
	blockID := in.BlockID

	path := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID))
	if ok, err := utils.LocalPathExists(path); !ok && err == nil {
		os.MkdirAll(path, 0777)
	}

	chunkFileName := path + "/chunk-" + strconv.Itoa(int(chunkID))

	f, err = os.OpenFile(chunkFileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0660)
	defer f.Close()
	if err != nil {
		ack.Ret = -1
		return &ack, nil
	}
	w := bufio.NewWriter(f)
	w.Write(in.Databuf)
	w.Flush()

	ack.Ret = 0
	return &ack, nil
}

/*rpc WriteChunkStream(stream WriteChunkReq) returns (WriteChunkAck){}; */
/*
func (s *DataNodeServer) WriteChunkStream(stream dp.DataNode_WriteChunkStreamServer) error {

	var f *os.File
	var err error
	//fmt.Println("writechunking in datanode ...")
	ack := dp.WriteChunkAck{}
	if err != nil {
		ack.Ret = -1
		return stream.SendAndClose(&ack)
	}
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			ack.Ret = 0
			return stream.SendAndClose(&ack)
		}
		if err != nil {
			ack.Ret = 1
			return stream.SendAndClose(&ack)
		}
		chunkID := in.ChunkID
		blockID := in.BlockID
		chunkFileName := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID)) + "/chunk-" + strconv.Itoa(int(chunkID))
		f, err = os.OpenFile(chunkFileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		f.WriteString(in.Databuf)
		f.Close()
	}
}
*/
/*
func (s *DataNodeServer) ReadChunk(ctx context.Context, in *dp.ReadChunkReq) (*dp.ReadChunkAck, error) {
	ack := dp.ReadChunkAck{}
	chunkID := in.ChunkID
	blockID := in.BlockID
	readsize := in.Readlen
	offset := in.Offset

	chunkFileName := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID)) + "/chunk-" + strconv.Itoa(int(chunkID))
	f, err := os.Open(chunkFileName)
	defer f.Close()
	if err != nil {
		ack.Ret = -1
		return &ack, nil
	}
	_, err = f.Seek(offset, 0)
	if err != nil {
		ack.Ret = -1
		return &ack, nil
	}

	buf := make([]byte, readsize)
	bfRd := bufio.NewReader(f)
	for {
		n, err := bfRd.Read(buf)
		if err != nil {
			ack.Ret = -1
			return &ack, nil
		}
		ack.Databuf = utils.B2S(buf)
		ack.Ret = 1
		ack.Readsize = int64(n)
		return &ack, nil
	}
}
*/

// StreamReadChunk ...
func (s *DataNodeServer) StreamReadChunk(in *dp.StreamReadChunkReq, stream dp.DataNode_StreamReadChunkServer) error {
	chunkID := in.ChunkID
	blockID := in.BlockID
	//offset := in.Offset
	readsize := in.Readsize

	chunkFileName := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID)) + "/chunk-" + strconv.Itoa(int(chunkID))
	f, err := os.Open(chunkFileName)
	defer f.Close()
	if err != nil {
		return err
	}

	var ack dp.StreamReadChunkAck
	var totalsize int64
	buf := make([]byte, 2*1024*1024)

	bfRd := bufio.NewReader(f)
	for {
		n, err := bfRd.Read(buf)
		if err != nil && err != io.EOF {
			logger.Error("read chunkfile:%v error:%v",chunkFileName,err)
			return err
		}

		if n == 0 {
			logger.Debug("read chunkfile:%v endsize:%v",chunkFileName,totalsize)
			break
		}

		totalsize += int64(n)

		ack.Databuf = buf[:n]
		if err := stream.Send(&ack); err != nil {
			logger.Error("Send stream data to fuse error:%v", err)
			return err
		}

		if totalsize == readsize {
			logger.Debug("Request chunkFileName:%v size:%v equal real ret size:%v", chunkFileName, readsize, totalsize)
			break
		}
	}

	return nil
}

//DeleteChunk rpc DeleteChunks(eleteChunksReq) returns (eleteChunksAck){};
func (s *DataNodeServer) DeleteChunk(ctx context.Context, in *dp.DeleteChunkReq) (*dp.DeleteChunkAck, error) {
	var err error

	ack := dp.DeleteChunkAck{}
	chunkID := in.ChunkID
	blockID := in.BlockID

	chunkFileName := DataNodeServerAddr.Path + "/block-" + strconv.Itoa(int(blockID)) + "/chunk-" + strconv.Itoa(int(chunkID))

	err = os.Remove(chunkFileName)
	if err != nil {
		ack.Ret = 0
	} else {
		ack.Ret = 0
	}
	ack.Ret = 0
	return &ack, nil
}

func init() {

	var loglevel string
	var port int

	flag.StringVar(&DataNodeServerAddr.IPStr, "host", "127.0.0.1", "ContainerFS DataNode Host")
	flag.IntVar(&port, "port", 8000, "ContainerFS DataNode Port")
	flag.StringVar(&DataNodeServerAddr.Path, "datapath", "", "ContainerFS DataNode Data Path")
	flag.StringVar(&DataNodeServerAddr.VolMgrHost, "volmgr", "127.0.0.1:7000", "ContainerFS VolMgr Host")
	flag.StringVar(&DataNodeServerAddr.Log, "logpath", "/export/Logs/containerfs/logs/", "ContainerFS Log Path")
	flag.StringVar(&loglevel, "loglevel", "error", "ContainerFS Log Level")

	flag.Parse()

	DataNodeServerAddr.Port = int32(port)
	ipnr := net.ParseIP(DataNodeServerAddr.IPStr)
	DataNodeServerAddr.Ipnr = ipnr
	ipint := utils.InetAton(ipnr)
	DataNodeServerAddr.IPInt = ipint
	DataNodeServerAddr.Flag = DataNodeServerAddr.Path + "/.registryflag"

	logger.SetConsole(true)
	logger.SetRollingFile(DataNodeServerAddr.Log, "datanode.log", 10, 100, logger.MB) //each 100M rolling

	switch loglevel {
	case "error":
		logger.SetLevel(logger.ERROR)
	case "debug":
		logger.SetLevel(logger.DEBUG)
	case "info":
		logger.SetLevel(logger.INFO)
	default:
		logger.SetLevel(logger.ERROR)
	}

	if ok, _ := utils.LocalPathExists(DataNodeServerAddr.Flag); !ok {
		logger.Debug("Start registry to volmgr ...")
		registryToVolMgr()
		logger.Debug("registry to volmgr success")
	} else {
		logger.Debug("already registied")
	}

}

func main() {

	numCPU := runtime.NumCPU()
	runtime.GOMAXPROCS(numCPU)

	defer func() {
		if err := recover(); err != nil {
			logger.Error("panic !!! :%v", err)
			logger.Error("stacks:%v", string(debug.Stack()))
		}
	}()

	heartbeatToVolMgr()
	ticker := time.NewTicker(time.Second * 60)
	go func() {
		for range ticker.C {
			heartbeatToVolMgr()
		}
	}()
	startDataService()
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 exec

import (
	"context"
	"fmt"
	"io"
	"reflect"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/graph"
	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
	"github.com/apache/beam/sdks/go/pkg/beam/core/typex"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/ioutilx"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/reflectx"
)

// StreamID represents the information needed to identify a data stream.
type StreamID struct {
	Port   graph.Port
	Target graph.Target
	InstID string
}

func (id StreamID) String() string {
	return fmt.Sprintf("S:%v:[%v:%v]:%v", id.Port.URL, id.Target.ID, id.Target.Name, id.InstID)
}

// DataReader is the interface for reading data elements from a particular stream.
type DataReader interface {
	OpenRead(ctx context.Context, id StreamID) (io.ReadCloser, error)
}

// DataWriter is the interface for writing data elements to a particular stream.
type DataWriter interface {
	OpenWrite(ctx context.Context, id StreamID) (io.WriteCloser, error)
}

// DataManager manages external data byte streams. Each data stream can be
// opened by one consumer only.
type DataManager interface {
	DataReader
	DataWriter
}

// NOTE(herohde) 4/30/2017: The main complication is CoGBK results, which have
// nested streams. Hence, a simple read-one-element-at-a-time approach doesn't
// work, because each "element" can be too large to fit into memory. Instead,
// we handle the top GBK/CoGBK layer in the processing node directly.

// ElementEncoder handles FullValue serialization to a byte stream. The encoder
// can be reused, even if an error is encountered.
type ElementEncoder interface {
	// Encode serializes the given value to the writer.
	Encode(FullValue, io.Writer) error
}

// ElementDecoder handles FullValue deserialization from a byte stream. The decoder
// can be reused, even if an error is encountered.
type ElementDecoder interface {
	// Decode deserializes a value from the given reader.
	Decode(io.Reader) (FullValue, error)
}

// MakeElementEncoder returns a ElementCoder for the given coder. It panics
// if given a coder with stream types, such as GBK.
func MakeElementEncoder(c *coder.Coder) ElementEncoder {
	switch c.Kind {
	case coder.Bytes:
		return &bytesEncoder{}

	case coder.VarInt:
		return &varIntEncoder{}

	case coder.Custom:
		return &customEncoder{
			t:   c.Custom.Type,
			enc: makeEncoder(c.Custom.Enc.Fn),
		}

	case coder.KV:
		return &kvEncoder{
			fst: MakeElementEncoder(c.Components[0]),
			snd: MakeElementEncoder(c.Components[1]),
		}

	default:
		panic(fmt.Sprintf("Unexpected coder: %v", c))
	}
}

// MakeElementDecoder returns a ElementDecoder for the given coder. It panics
// if given a coder with stream types, such as GBK.
func MakeElementDecoder(c *coder.Coder) ElementDecoder {
	switch c.Kind {
	case coder.Bytes:
		return &bytesDecoder{}

	case coder.VarInt:
		return &varIntDecoder{}

	case coder.Custom:
		return &customDecoder{
			t:   c.Custom.Type,
			dec: makeDecoder(c.Custom.Dec.Fn),
		}

	case coder.KV:
		return &kvDecoder{
			fst: MakeElementDecoder(c.Components[0]),
			snd: MakeElementDecoder(c.Components[1]),
		}

	default:
		panic(fmt.Sprintf("Unexpected coder: %v", c))
	}
}

type bytesEncoder struct{}

func (*bytesEncoder) Encode(val FullValue, w io.Writer) error {
	// Encoding: size (varint) + raw data

	data := reflect.ValueOf(val.Elm).Convert(reflectx.ByteSlice).Interface().([]byte)
	size := len(data)

	if err := coder.EncodeVarInt((int32)(size), w); err != nil {
		return err
	}
	_, err := w.Write(data)
	return err
}

type bytesDecoder struct{}

func (*bytesDecoder) Decode(r io.Reader) (FullValue, error) {
	// Encoding: size (varint) + raw data

	size, err := coder.DecodeVarInt(r)
	if err != nil {
		return FullValue{}, err
	}
	data, err := ioutilx.ReadN(r, (int)(size))
	if err != nil {
		return FullValue{}, err
	}
	return FullValue{Elm: data}, nil
}

type varIntEncoder struct{}

func (*varIntEncoder) Encode(val FullValue, w io.Writer) error {
	// Encoding: beam varint

	n := Convert(val.Elm, reflectx.Int32).(int32) // Convert needed?
	return coder.EncodeVarInt(n, w)
}

type varIntDecoder struct{}

func (*varIntDecoder) Decode(r io.Reader) (FullValue, error) {
	// Encoding: beam varint

	n, err := coder.DecodeVarInt(r)
	if err != nil {
		return FullValue{}, err
	}
	return FullValue{Elm: n}, nil
}

type customEncoder struct {
	t   reflect.Type
	enc Encoder
}

func (c *customEncoder) Encode(val FullValue, w io.Writer) error {
	// (1) Call encode

	data, err := c.enc.Encode(c.t, val.Elm)
	if err != nil {
		return err
	}

	// (2) Add length prefix

	size := len(data)
	if err := coder.EncodeVarInt((int32)(size), w); err != nil {
		return err
	}
	_, err = w.Write(data)
	return err
}

type customDecoder struct {
	t   reflect.Type
	dec Decoder
}

func (c *customDecoder) Decode(r io.Reader) (FullValue, error) {
	// (1) Read length-prefixed encoded data

	size, err := coder.DecodeVarInt(r)
	if err != nil {
		return FullValue{}, err
	}
	data, err := ioutilx.ReadN(r, (int)(size))
	if err != nil {
		return FullValue{}, err
	}

	// (2) Call decode

	val, err := c.dec.Decode(c.t, data)
	if err != nil {
		return FullValue{}, err
	}
	return FullValue{Elm: val}, err
}

type kvEncoder struct {
	fst, snd ElementEncoder
}

func (c *kvEncoder) Encode(val FullValue, w io.Writer) error {
	if err := c.fst.Encode(FullValue{Elm: val.Elm}, w); err != nil {
		return err
	}
	return c.snd.Encode(FullValue{Elm: val.Elm2}, w)

}

type kvDecoder struct {
	fst, snd ElementDecoder
}

func (c *kvDecoder) Decode(r io.Reader) (FullValue, error) {
	key, err := c.fst.Decode(r)
	if err != nil {
		return FullValue{}, err
	}
	value, err := c.snd.Decode(r)
	if err != nil {
		return FullValue{}, err
	}
	return FullValue{Elm: key.Elm, Elm2: value.Elm}, nil

}

// TODO(herohde) 4/7/2017: actually handle windows.

// EncodeWindowedValueHeader uses the supplied coder to serialize a windowed value header.
func EncodeWindowedValueHeader(c *coder.Coder, t typex.EventTime, w io.Writer) error {
	// Encoding: Timestamp, Window, Pane (header) + Element

	if err := coder.EncodeEventTime(t, w); err != nil {
		return err
	}
	if err := coder.EncodeInt32(1, w); err != nil { // #windows
		return err
	}
	// Ignore GlobalWindow, for now. It encoded into the empty string.

	_, err := w.Write([]byte{0xf}) // NO_FIRING pane
	return err
}

// DecodeWindowedValueHeader uses the supplied coder to deserialize a windowed value header.
func DecodeWindowedValueHeader(c *coder.Coder, r io.Reader) (typex.EventTime, error) {
	// Encoding: Timestamp, Window, Pane (header) + Element

	t, err := coder.DecodeEventTime(r)
	if err != nil {
		return typex.EventTime(time.Time{}), err
	}
	if _, err := coder.DecodeInt32(r); err != nil { // #windows
		return typex.EventTime(time.Time{}), err
	}
	if _, err := ioutilx.ReadN(r, 1); err != nil { // NO_FIRING pane
		return typex.EventTime(time.Time{}), err
	}
	return t, nil
}
<commit_msg>Use a typeswitch instead of reflect.Convert when encoding strings or bytes.<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 exec

import (
	"context"
	"fmt"
	"io"
	"reflect"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/graph"
	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
	"github.com/apache/beam/sdks/go/pkg/beam/core/typex"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/ioutilx"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/reflectx"
)

// StreamID represents the information needed to identify a data stream.
type StreamID struct {
	Port   graph.Port
	Target graph.Target
	InstID string
}

func (id StreamID) String() string {
	return fmt.Sprintf("S:%v:[%v:%v]:%v", id.Port.URL, id.Target.ID, id.Target.Name, id.InstID)
}

// DataReader is the interface for reading data elements from a particular stream.
type DataReader interface {
	OpenRead(ctx context.Context, id StreamID) (io.ReadCloser, error)
}

// DataWriter is the interface for writing data elements to a particular stream.
type DataWriter interface {
	OpenWrite(ctx context.Context, id StreamID) (io.WriteCloser, error)
}

// DataManager manages external data byte streams. Each data stream can be
// opened by one consumer only.
type DataManager interface {
	DataReader
	DataWriter
}

// NOTE(herohde) 4/30/2017: The main complication is CoGBK results, which have
// nested streams. Hence, a simple read-one-element-at-a-time approach doesn't
// work, because each "element" can be too large to fit into memory. Instead,
// we handle the top GBK/CoGBK layer in the processing node directly.

// ElementEncoder handles FullValue serialization to a byte stream. The encoder
// can be reused, even if an error is encountered.
type ElementEncoder interface {
	// Encode serializes the given value to the writer.
	Encode(FullValue, io.Writer) error
}

// ElementDecoder handles FullValue deserialization from a byte stream. The decoder
// can be reused, even if an error is encountered.
type ElementDecoder interface {
	// Decode deserializes a value from the given reader.
	Decode(io.Reader) (FullValue, error)
}

// MakeElementEncoder returns a ElementCoder for the given coder. It panics
// if given a coder with stream types, such as GBK.
func MakeElementEncoder(c *coder.Coder) ElementEncoder {
	switch c.Kind {
	case coder.Bytes:
		return &bytesEncoder{}

	case coder.VarInt:
		return &varIntEncoder{}

	case coder.Custom:
		return &customEncoder{
			t:   c.Custom.Type,
			enc: makeEncoder(c.Custom.Enc.Fn),
		}

	case coder.KV:
		return &kvEncoder{
			fst: MakeElementEncoder(c.Components[0]),
			snd: MakeElementEncoder(c.Components[1]),
		}

	default:
		panic(fmt.Sprintf("Unexpected coder: %v", c))
	}
}

// MakeElementDecoder returns a ElementDecoder for the given coder. It panics
// if given a coder with stream types, such as GBK.
func MakeElementDecoder(c *coder.Coder) ElementDecoder {
	switch c.Kind {
	case coder.Bytes:
		return &bytesDecoder{}

	case coder.VarInt:
		return &varIntDecoder{}

	case coder.Custom:
		return &customDecoder{
			t:   c.Custom.Type,
			dec: makeDecoder(c.Custom.Dec.Fn),
		}

	case coder.KV:
		return &kvDecoder{
			fst: MakeElementDecoder(c.Components[0]),
			snd: MakeElementDecoder(c.Components[1]),
		}

	default:
		panic(fmt.Sprintf("Unexpected coder: %v", c))
	}
}

type bytesEncoder struct{}

func (*bytesEncoder) Encode(val FullValue, w io.Writer) error {
	// Encoding: size (varint) + raw data
	var data []byte
	switch v := val.Elm.(type) {
	case []byte:
		data = v
	case string:
		data = []byte(v)
	default:
		return fmt.Errorf("received unknown value type: want []byte or string, got %T", v)
	}

	if err := coder.EncodeVarInt((int32)(size), w); err != nil {
		return err
	}
	_, err := w.Write(data)
	return err
}

type bytesDecoder struct{}

func (*bytesDecoder) Decode(r io.Reader) (FullValue, error) {
	// Encoding: size (varint) + raw data

	size, err := coder.DecodeVarInt(r)
	if err != nil {
		return FullValue{}, err
	}
	data, err := ioutilx.ReadN(r, (int)(size))
	if err != nil {
		return FullValue{}, err
	}
	return FullValue{Elm: data}, nil
}

type varIntEncoder struct{}

func (*varIntEncoder) Encode(val FullValue, w io.Writer) error {
	// Encoding: beam varint

	n := Convert(val.Elm, reflectx.Int32).(int32) // Convert needed?
	return coder.EncodeVarInt(n, w)
}

type varIntDecoder struct{}

func (*varIntDecoder) Decode(r io.Reader) (FullValue, error) {
	// Encoding: beam varint

	n, err := coder.DecodeVarInt(r)
	if err != nil {
		return FullValue{}, err
	}
	return FullValue{Elm: n}, nil
}

type customEncoder struct {
	t   reflect.Type
	enc Encoder
}

func (c *customEncoder) Encode(val FullValue, w io.Writer) error {
	// (1) Call encode

	data, err := c.enc.Encode(c.t, val.Elm)
	if err != nil {
		return err
	}

	// (2) Add length prefix

	size := len(data)
	if err := coder.EncodeVarInt((int32)(size), w); err != nil {
		return err
	}
	_, err = w.Write(data)
	return err
}

type customDecoder struct {
	t   reflect.Type
	dec Decoder
}

func (c *customDecoder) Decode(r io.Reader) (FullValue, error) {
	// (1) Read length-prefixed encoded data

	size, err := coder.DecodeVarInt(r)
	if err != nil {
		return FullValue{}, err
	}
	data, err := ioutilx.ReadN(r, (int)(size))
	if err != nil {
		return FullValue{}, err
	}

	// (2) Call decode

	val, err := c.dec.Decode(c.t, data)
	if err != nil {
		return FullValue{}, err
	}
	return FullValue{Elm: val}, err
}

type kvEncoder struct {
	fst, snd ElementEncoder
}

func (c *kvEncoder) Encode(val FullValue, w io.Writer) error {
	if err := c.fst.Encode(FullValue{Elm: val.Elm}, w); err != nil {
		return err
	}
	return c.snd.Encode(FullValue{Elm: val.Elm2}, w)

}

type kvDecoder struct {
	fst, snd ElementDecoder
}

func (c *kvDecoder) Decode(r io.Reader) (FullValue, error) {
	key, err := c.fst.Decode(r)
	if err != nil {
		return FullValue{}, err
	}
	value, err := c.snd.Decode(r)
	if err != nil {
		return FullValue{}, err
	}
	return FullValue{Elm: key.Elm, Elm2: value.Elm}, nil

}

// TODO(herohde) 4/7/2017: actually handle windows.

// EncodeWindowedValueHeader uses the supplied coder to serialize a windowed value header.
func EncodeWindowedValueHeader(c *coder.Coder, t typex.EventTime, w io.Writer) error {
	// Encoding: Timestamp, Window, Pane (header) + Element

	if err := coder.EncodeEventTime(t, w); err != nil {
		return err
	}
	if err := coder.EncodeInt32(1, w); err != nil { // #windows
		return err
	}
	// Ignore GlobalWindow, for now. It encoded into the empty string.

	_, err := w.Write([]byte{0xf}) // NO_FIRING pane
	return err
}

// DecodeWindowedValueHeader uses the supplied coder to deserialize a windowed value header.
func DecodeWindowedValueHeader(c *coder.Coder, r io.Reader) (typex.EventTime, error) {
	// Encoding: Timestamp, Window, Pane (header) + Element

	t, err := coder.DecodeEventTime(r)
	if err != nil {
		return typex.EventTime(time.Time{}), err
	}
	if _, err := coder.DecodeInt32(r); err != nil { // #windows
		return typex.EventTime(time.Time{}), err
	}
	if _, err := ioutilx.ReadN(r, 1); err != nil { // NO_FIRING pane
		return typex.EventTime(time.Time{}), err
	}
	return t, nil
}
<|endoftext|>


<commit_before>package watchers

import (
	"fmt"
	"sync"
	"time"

	"github.com/funkygao/dbus/pkg/cluster"
	czk "github.com/funkygao/dbus/pkg/cluster/zk"
	"github.com/funkygao/gafka/cmd/kguard/monitor"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	log "github.com/funkygao/log4go"
)

type dbusWatcher struct {
	ident string

	zkzone  *zk.ZkZone
	stopper <-chan struct{}
	wg      *sync.WaitGroup
}

func (this *dbusWatcher) Init(ctx monitor.Context) {
	this.zkzone = ctx.ZkZone()
	this.stopper = ctx.StopChan()
	this.wg = ctx.Inflight()
}

func (this *dbusWatcher) Run() {
	defer this.wg.Done()

	resourcesGauge := metrics.NewRegisteredGauge("dbus.resources", nil)
	orphanResourcesGauge := metrics.NewRegisteredGauge("dbus.resources.orphan", nil)
	participantsGauge := metrics.NewRegisteredGauge("dbus.participants", nil)

	mgr := czk.NewManager(this.zkzone.ZkAddrs())
	if err := mgr.Open(); err != nil {
		log.Warn("%s quit: %v", this.ident, err)
		return
	}

	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-this.stopper:
			log.Info("%s stopped", this.ident)
			return

		case <-ticker.C:
			orphanN := 0
			if resources, err := mgr.RegisteredResources(); err != nil {
				log.Error("%s %v", this.ident, err)
			} else {
				for _, r := range resources {
					if r.IsOrphan() {
						orphanN++
					}
				}
			}
			resourcesGauge.Update(int64(len(resources)))
			orphanResourcesGauge.Update(int64(orphanN))

			if liveParticipants, err := this.mgr.LiveParticipants(); err != nil {
				log.Error("%s %v", this.ident, err)
			} else {
				participantsGauge.Update(int64(len(liveParticipants)))
			}
		}
	}
}

func init() {
	monitor.RegisterWatcher("dbus.dbus", func() monitor.Watcher {
		return &dbusWatcher{ident: "dbus.dbus"}
	})
}
<commit_msg>monitor reelection of cluster leader<commit_after>package watchers

import (
	"sync"
	"time"

	"github.com/funkygao/dbus/pkg/cluster"
	czk "github.com/funkygao/dbus/pkg/cluster/zk"
	"github.com/funkygao/gafka/cmd/kguard/monitor"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	log "github.com/funkygao/log4go"
)

type dbusWatcher struct {
	ident string

	zkzone  *zk.ZkZone
	stopper <-chan struct{}
	wg      *sync.WaitGroup
}

func (this *dbusWatcher) Init(ctx monitor.Context) {
	this.zkzone = ctx.ZkZone()
	this.stopper = ctx.StopChan()
	this.wg = ctx.Inflight()
}

func (this *dbusWatcher) Run() {
	defer this.wg.Done()

	resourcesGauge := metrics.NewRegisteredGauge("dbus.resources", nil)
	orphanResourcesGauge := metrics.NewRegisteredGauge("dbus.resources.orphan", nil)
	participantsGauge := metrics.NewRegisteredGauge("dbus.participants", nil)
	reelectGauge := metrics.NewRegisteredGauge("dbus.reelect", nil)

	mgr := czk.NewManager(this.zkzone.ZkAddrs())
	if err := mgr.Open(); err != nil {
		log.Warn("%s quit: %v", this.ident, err)
		return
	}

	lastLeader := cluster.Participant{}
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-this.stopper:
			log.Info("%s stopped", this.ident)
			return

		case <-ticker.C:
			orphanN := 0
			if resources, err := mgr.RegisteredResources(); err != nil {
				log.Error("%s %v", this.ident, err)
			} else {
				for _, r := range resources {
					if r.IsOrphan() {
						orphanN++
					}
				}

				resourcesGauge.Update(int64(len(resources)))
				orphanResourcesGauge.Update(int64(orphanN))
			}

			if liveParticipants, err := mgr.LiveParticipants(); err != nil {
				log.Error("%s %v", this.ident, err)
			} else {
				participantsGauge.Update(int64(len(liveParticipants)))
			}

			if leader, err := mgr.Leader(); err != nil {
				log.Error("%s %v", this.ident, err)
			} else {
				if lastLeader.Valid() && !leader.Equals(lastLeader) {
					log.Trace("%s new leader: %s", this.ident, leader)
					reelectGauge.Update(1)
				} else {
					reelectGauge.Update(0)
				}

				lastLeader = leader
			}
		}
	}
}

func init() {
	monitor.RegisterWatcher("dbus.dbus", func() monitor.Watcher {
		return &dbusWatcher{ident: "dbus.dbus"}
	})
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 bigquery

import (
	"fmt"

	"golang.org/x/net/context"
	bq "google.golang.org/api/bigquery/v2"
)

type queryOption interface {
	customizeQuery(conf *bq.JobConfigurationQuery, projectID string)
}

// DisableQueryCache returns an Option that prevents results being fetched from the query cache.
// If this Option is not used, results are fetched from the cache if they are available.
// The query cache is a best-effort cache that is flushed whenever tables in the query are modified.
// Cached results are only available when TableID is unspecified in the query's destination Table.
// For more information, see https://cloud.google.com/bigquery/querying-data#querycaching
func DisableQueryCache() Option { return disableQueryCache{} }

type disableQueryCache struct{}

func (opt disableQueryCache) implementsOption() {}

func (opt disableQueryCache) customizeQuery(conf *bq.JobConfigurationQuery, projectID string) {
	f := false
	conf.UseQueryCache = &f
}

// JobPriority returns an Option that causes a query to be scheduled with the specified priority.
// The default priority is InteractivePriority.
// For more information, see https://cloud.google.com/bigquery/querying-data#batchqueries
func JobPriority(priority string) Option { return jobPriority(priority) }

type jobPriority string

func (opt jobPriority) implementsOption() {}

func (opt jobPriority) customizeQuery(conf *bq.JobConfigurationQuery, projectID string) {
	conf.Priority = string(opt)
}

const (
	BatchPriority       = "BATCH"
	InteractivePriority = "INTERACTIVE"
)

// TODO(mcgreevy): support large results.
// TODO(mcgreevy): support non-flattened results.

func (c *Client) query(ctx context.Context, dst *Table, src *Query, options []Option) (*Job, error) {
	job, options := initJobProto(c.projectID, options)
	payload := &bq.JobConfigurationQuery{}

	dst.customizeQueryDst(payload, c.projectID)
	src.customizeQuerySrc(payload, c.projectID)

	for _, opt := range options {
		o, ok := opt.(queryOption)
		if !ok {
			return nil, fmt.Errorf("option (%#v) not applicable to dst/src pair: dst: %T ; src: %T", opt, dst, src)
		}
		o.customizeQuery(payload, c.projectID)
	}

	job.Configuration = &bq.JobConfiguration{
		Query: payload,
	}
	j, err := c.service.insertJob(ctx, job, c.projectID)
	if err != nil {
		j.isQuery = true
	}
	return j, err
}
<commit_msg>cloud/bigquery: Set isQuery bit correctly when inserting a query job.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 bigquery

import (
	"fmt"

	"golang.org/x/net/context"
	bq "google.golang.org/api/bigquery/v2"
)

type queryOption interface {
	customizeQuery(conf *bq.JobConfigurationQuery, projectID string)
}

// DisableQueryCache returns an Option that prevents results being fetched from the query cache.
// If this Option is not used, results are fetched from the cache if they are available.
// The query cache is a best-effort cache that is flushed whenever tables in the query are modified.
// Cached results are only available when TableID is unspecified in the query's destination Table.
// For more information, see https://cloud.google.com/bigquery/querying-data#querycaching
func DisableQueryCache() Option { return disableQueryCache{} }

type disableQueryCache struct{}

func (opt disableQueryCache) implementsOption() {}

func (opt disableQueryCache) customizeQuery(conf *bq.JobConfigurationQuery, projectID string) {
	f := false
	conf.UseQueryCache = &f
}

// JobPriority returns an Option that causes a query to be scheduled with the specified priority.
// The default priority is InteractivePriority.
// For more information, see https://cloud.google.com/bigquery/querying-data#batchqueries
func JobPriority(priority string) Option { return jobPriority(priority) }

type jobPriority string

func (opt jobPriority) implementsOption() {}

func (opt jobPriority) customizeQuery(conf *bq.JobConfigurationQuery, projectID string) {
	conf.Priority = string(opt)
}

const (
	BatchPriority       = "BATCH"
	InteractivePriority = "INTERACTIVE"
)

// TODO(mcgreevy): support large results.
// TODO(mcgreevy): support non-flattened results.

func (c *Client) query(ctx context.Context, dst *Table, src *Query, options []Option) (*Job, error) {
	job, options := initJobProto(c.projectID, options)
	payload := &bq.JobConfigurationQuery{}

	dst.customizeQueryDst(payload, c.projectID)
	src.customizeQuerySrc(payload, c.projectID)

	for _, opt := range options {
		o, ok := opt.(queryOption)
		if !ok {
			return nil, fmt.Errorf("option (%#v) not applicable to dst/src pair: dst: %T ; src: %T", opt, dst, src)
		}
		o.customizeQuery(payload, c.projectID)
	}

	job.Configuration = &bq.JobConfiguration{
		Query: payload,
	}
	j, err := c.service.insertJob(ctx, job, c.projectID)
	if err != nil {
		return nil, err
	}
	j.isQuery = true
	return j, nil
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 session

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/harness/session"
	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	fnpb "github.com/apache/beam/sdks/go/pkg/beam/model/fnexecution_v1"
	pipepb "github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
)

const (
	// The maximum length of an encoded varint. We can Peek this much data
	// and find a value to decode.
	peekLen = 9
)

func init() {
	beam.RegisterRunner("session", Execute)
}

var sessionFile = flag.String("session_file", "", "Session file for the runner")

// controlServer manages the FnAPI control channel.
type controlServer struct {
	filename   string
	wg         *sync.WaitGroup // used to signal when the session is completed
	ctrlStream fnpb.BeamFnControl_ControlServer
	dataServer *grpc.Server
	dataStream fnpb.BeamFnData_DataServer
	dwg        *sync.WaitGroup
}

func (c *controlServer) Control(stream fnpb.BeamFnControl_ControlServer) error {
	fmt.Println("Go SDK connected")
	c.ctrlStream = stream
	// We have a connected worker. Start reading the session file and issuing
	// commands.

	c.readSession(c.filename)
	c.wg.Done()
	fmt.Println("session replay complete")
	return nil
}

func (c *controlServer) GetProcessBundleDescriptor(ctx context.Context, r *fnpb.GetProcessBundleDescriptorRequest) (*fnpb.ProcessBundleDescriptor, error) {
	return nil, nil
}

func (c *controlServer) establishDataChannel(beamPort, tcpPort string) {
	if c.dataServer != nil {
		// Already a data server, we're done
		return
	}

	// grpc can allow a grpc service running on two different ports, but there's
	// no way (in Go at least) to differentiate the two of them to identify the
	// source of the incoming data. So we don't even try. Session files that
	// specify data ports have the content rewritten to use the port that
	// the data server is listening on.

	c.dataServer = grpc.NewServer()
	fnpb.RegisterBeamFnDataServer(c.dataServer, &dataServer{ctrl: c})
	dp, err := net.Listen("tcp", tcpPort)
	if err != nil {
		panic(err)
	}
	c.dwg = &sync.WaitGroup{}
	c.dwg.Add(1)
	go c.dataServer.Serve(dp)
}

func (c *controlServer) registerStream(stream fnpb.BeamFnData_DataServer) {
	c.dataStream = stream
	c.dwg.Done()
}

// TODO(wcn): move this code to a session file framework. I imagine this will
// take an additional function argument that performs the handleEntry() work.
func (c *controlServer) readSession(filename string) {
	// Keep the reading simple by ensuring the buffer is large enough
	// to hold any single recorded message. Since grpc has a message
	// cap of 4 megs, we make our buffer larger. Future versions of the
	// header will include this constant, so we can read the header
	// unbuffered, then move to the appropriately sized buffer reader.
	f, err := os.Open(filename)
	if err != nil {
		panic(err)
	}

	br := bufio.NewReaderSize(f, 5000000)
	for {
		b, err := br.Peek(peekLen)
		if err != nil && err != io.EOF {
			panic(errors.Wrap(err, "Problem peeking length value"))
		}
		if err == io.EOF {
			break
		}
		l, inc := proto.DecodeVarint(b)
		br.Discard(inc)

		// Read out the entry header message.
		b, err = br.Peek(int(l))
		var hMsg session.EntryHeader
		if err := proto.Unmarshal(b, &hMsg); err != nil {
			panic(errors.Wrap(err, "Error decoding entry header"))
		}
		br.Discard(int(l))

		msgBytes, err := br.Peek(int(hMsg.Len))
		if err != nil {
			panic(errors.Wrap(err, "Couldn't peek message"))
		}

		var bMsg session.Entry
		if err := proto.Unmarshal(msgBytes, &bMsg); err != nil {
			panic(errors.Wrap(err, "Error decoding message"))
		}
		c.handleEntry(&bMsg)
		br.Discard(int(hMsg.Len))
	}
}

func (c *controlServer) handleEntry(msg *session.Entry) {
	/*
		if msg.Kind != session.Entry_LOG_ENTRIES {
			fmt.Printf("handleEntry: %v\n", msg.Kind.String())
		}
	*/
	switch msg.Msg.(type) {
	case *session.Entry_Elems:
		if msg.GetKind() == session.Kind_DATA_RECEIVED {
			c.dwg.Wait()
			c.dataStream.Send(msg.GetElems())
		}
	case *session.Entry_InstResp:
		_, err := c.ctrlStream.Recv()
		if err == io.EOF {
			panic("SDK closed connection but work remaining")
		}

		if err != nil {
			return
		}

	case *session.Entry_InstReq:
		// Look for the register requests and extract the port information.
		ir := msg.GetInstReq()
		c.ctrlStream.Send(ir)

		if rr := ir.GetRegister(); rr != nil {
			for _, desc := range rr.GetProcessBundleDescriptor() {
				for beamPort, t := range desc.GetTransforms() {
					s := t.GetSpec()
					if s.GetUrn() == "beam:runner:source:v1" {
						tcpPort := extractPortSpec(s)
						c.establishDataChannel(beamPort, tcpPort)
					}
					if s.GetUrn() == "beam:runner:sink:v1" {
						tcpPort := extractPortSpec(s)
						c.establishDataChannel(beamPort, tcpPort)
					}
				}
			}
		}
	}
}

func extractPortSpec(spec *pipepb.FunctionSpec) string {
	var port fnpb.RemoteGrpcPort
	if err := proto.Unmarshal(spec.GetPayload(), &port); err != nil {
		panic(err)
	}
	lp := port.ApiServiceDescriptor.Url
	// Leave the colon, so as to match the form net.Listen uses.
	bp := strings.Replace(lp, "localhost", "", 1)
	if bp != lp {
		return bp
	}
	panic("unable to extract port")
}

// dataServer manages the FnAPI data channel.
type dataServer struct {
	ctrl *controlServer
}

func (d *dataServer) Data(stream fnpb.BeamFnData_DataServer) error {
	// This goroutine is only used for reading data. The stream object
	// is passed to the control server so that all data is sent from
	// a single goroutine to ensure proper ordering.

	d.ctrl.registerStream(stream)

	// Consume data messages that are received
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			return nil
		}

		if err != nil {
			return err
		}

		_ = in
		//log.Printf("Data received: %v", in)
	}
}

// loggingServer manages the FnAPI logging channel.
type loggingServer struct{} // no data content

func (l *loggingServer) Logging(stream fnpb.BeamFnLogging_LoggingServer) error {
	// This stream object is only used here. The stream is used for receiving, and
	// no sends happen on it.
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			return nil
		}

		if err != nil {
			return err
		}

		for _, e := range in.GetLogEntries() {
			log.Info(stream.Context(), e.GetMessage())
		}
	}
}

// Execute launches the supplied pipeline using a session file as the source of inputs.
func Execute(ctx context.Context, p *beam.Pipeline) error {
	worker, err := buildLocalBinary(ctx)
	if err != nil {
		return errors.WithContext(err, "building worker binary")
	}

	log.Infof(ctx, "built worker binary at %s\n", worker)

	// Start up the grpc logging service.
	ls := grpc.NewServer()
	fnpb.RegisterBeamFnLoggingServer(ls, &loggingServer{})
	logPort, err := net.Listen("tcp", ":0")
	if err != nil {
		panic("No logging port")
	}
	go ls.Serve(logPort)

	// The wait group is used by the control service goroutine to signal
	// completion.
	var wg sync.WaitGroup
	wg.Add(1)

	cs := grpc.NewServer()
	fnpb.RegisterBeamFnControlServer(cs, &controlServer{
		filename: *sessionFile,
		wg:       &wg,
	})

	ctrlPort, err := net.Listen("tcp", ":0")
	if err != nil {
		panic("No control port")
	}
	go cs.Serve(ctrlPort)

	fmt.Println("fake harness initialized")
	cmd := exec.Command(
		worker,
		"--worker",
		fmt.Sprintf("--logging_endpoint=%s", logPort.Addr().String()),
		fmt.Sprintf("--control_endpoint=%s", ctrlPort.Addr().String()),
		"--persist_dir=/tmp/worker")
	go cmd.Start()

	wg.Wait()
	return nil
}

// buildLocalBinary is cribbed from the Dataflow runner, but doesn't force the
// Linux architecture, since the worker runs in the pipeline launch
// environment.
func buildLocalBinary(ctx context.Context) (string, error) {
	ret := filepath.Join(os.TempDir(), fmt.Sprintf("session-runner-%v", time.Now().UnixNano()))

	program := ""
	for i := 3; ; i++ {
		_, file, _, ok := runtime.Caller(i)
		if !ok || strings.HasSuffix(file, "runtime/proc.go") {
			break
		}
		program = file
	}
	if program == "" {
		return "", errors.New("could not detect user main")
	}

	log.Infof(ctx, "Compiling %v as %v", program, ret)

	// Cross-compile given go program. Not awesome.
	build := []string{"go", "build", "-o", ret, program}

	cmd := exec.Command(build[0], build[1:]...)
	if out, err := cmd.CombinedOutput(); err != nil {
		log.Info(ctx, string(out))
		return "", errors.Wrapf(err, "failed to compile %v", program)
	}
	return ret, nil
}
<commit_msg>[BEAM-11207] Adding PipelineResult to session runner. (#13371)<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 session

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/harness/session"
	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	fnpb "github.com/apache/beam/sdks/go/pkg/beam/model/fnexecution_v1"
	pipepb "github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
)

const (
	// The maximum length of an encoded varint. We can Peek this much data
	// and find a value to decode.
	peekLen = 9
)

func init() {
	beam.RegisterRunner("session", Execute)
}

var sessionFile = flag.String("session_file", "", "Session file for the runner")

// controlServer manages the FnAPI control channel.
type controlServer struct {
	filename   string
	wg         *sync.WaitGroup // used to signal when the session is completed
	ctrlStream fnpb.BeamFnControl_ControlServer
	dataServer *grpc.Server
	dataStream fnpb.BeamFnData_DataServer
	dwg        *sync.WaitGroup
}

func (c *controlServer) Control(stream fnpb.BeamFnControl_ControlServer) error {
	fmt.Println("Go SDK connected")
	c.ctrlStream = stream
	// We have a connected worker. Start reading the session file and issuing
	// commands.

	c.readSession(c.filename)
	c.wg.Done()
	fmt.Println("session replay complete")
	return nil
}

func (c *controlServer) GetProcessBundleDescriptor(ctx context.Context, r *fnpb.GetProcessBundleDescriptorRequest) (*fnpb.ProcessBundleDescriptor, error) {
	return nil, nil
}

func (c *controlServer) establishDataChannel(beamPort, tcpPort string) {
	if c.dataServer != nil {
		// Already a data server, we're done
		return
	}

	// grpc can allow a grpc service running on two different ports, but there's
	// no way (in Go at least) to differentiate the two of them to identify the
	// source of the incoming data. So we don't even try. Session files that
	// specify data ports have the content rewritten to use the port that
	// the data server is listening on.

	c.dataServer = grpc.NewServer()
	fnpb.RegisterBeamFnDataServer(c.dataServer, &dataServer{ctrl: c})
	dp, err := net.Listen("tcp", tcpPort)
	if err != nil {
		panic(err)
	}
	c.dwg = &sync.WaitGroup{}
	c.dwg.Add(1)
	go c.dataServer.Serve(dp)
}

func (c *controlServer) registerStream(stream fnpb.BeamFnData_DataServer) {
	c.dataStream = stream
	c.dwg.Done()
}

// TODO(wcn): move this code to a session file framework. I imagine this will
// take an additional function argument that performs the handleEntry() work.
func (c *controlServer) readSession(filename string) {
	// Keep the reading simple by ensuring the buffer is large enough
	// to hold any single recorded message. Since grpc has a message
	// cap of 4 megs, we make our buffer larger. Future versions of the
	// header will include this constant, so we can read the header
	// unbuffered, then move to the appropriately sized buffer reader.
	f, err := os.Open(filename)
	if err != nil {
		panic(err)
	}

	br := bufio.NewReaderSize(f, 5000000)
	for {
		b, err := br.Peek(peekLen)
		if err != nil && err != io.EOF {
			panic(errors.Wrap(err, "Problem peeking length value"))
		}
		if err == io.EOF {
			break
		}
		l, inc := proto.DecodeVarint(b)
		br.Discard(inc)

		// Read out the entry header message.
		b, err = br.Peek(int(l))
		var hMsg session.EntryHeader
		if err := proto.Unmarshal(b, &hMsg); err != nil {
			panic(errors.Wrap(err, "Error decoding entry header"))
		}
		br.Discard(int(l))

		msgBytes, err := br.Peek(int(hMsg.Len))
		if err != nil {
			panic(errors.Wrap(err, "Couldn't peek message"))
		}

		var bMsg session.Entry
		if err := proto.Unmarshal(msgBytes, &bMsg); err != nil {
			panic(errors.Wrap(err, "Error decoding message"))
		}
		c.handleEntry(&bMsg)
		br.Discard(int(hMsg.Len))
	}
}

func (c *controlServer) handleEntry(msg *session.Entry) {
	/*
		if msg.Kind != session.Entry_LOG_ENTRIES {
			fmt.Printf("handleEntry: %v\n", msg.Kind.String())
		}
	*/
	switch msg.Msg.(type) {
	case *session.Entry_Elems:
		if msg.GetKind() == session.Kind_DATA_RECEIVED {
			c.dwg.Wait()
			c.dataStream.Send(msg.GetElems())
		}
	case *session.Entry_InstResp:
		_, err := c.ctrlStream.Recv()
		if err == io.EOF {
			panic("SDK closed connection but work remaining")
		}

		if err != nil {
			return
		}

	case *session.Entry_InstReq:
		// Look for the register requests and extract the port information.
		ir := msg.GetInstReq()
		c.ctrlStream.Send(ir)

		if rr := ir.GetRegister(); rr != nil {
			for _, desc := range rr.GetProcessBundleDescriptor() {
				for beamPort, t := range desc.GetTransforms() {
					s := t.GetSpec()
					if s.GetUrn() == "beam:runner:source:v1" {
						tcpPort := extractPortSpec(s)
						c.establishDataChannel(beamPort, tcpPort)
					}
					if s.GetUrn() == "beam:runner:sink:v1" {
						tcpPort := extractPortSpec(s)
						c.establishDataChannel(beamPort, tcpPort)
					}
				}
			}
		}
	}
}

func extractPortSpec(spec *pipepb.FunctionSpec) string {
	var port fnpb.RemoteGrpcPort
	if err := proto.Unmarshal(spec.GetPayload(), &port); err != nil {
		panic(err)
	}
	lp := port.ApiServiceDescriptor.Url
	// Leave the colon, so as to match the form net.Listen uses.
	bp := strings.Replace(lp, "localhost", "", 1)
	if bp != lp {
		return bp
	}
	panic("unable to extract port")
}

// dataServer manages the FnAPI data channel.
type dataServer struct {
	ctrl *controlServer
}

func (d *dataServer) Data(stream fnpb.BeamFnData_DataServer) error {
	// This goroutine is only used for reading data. The stream object
	// is passed to the control server so that all data is sent from
	// a single goroutine to ensure proper ordering.

	d.ctrl.registerStream(stream)

	// Consume data messages that are received
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			return nil
		}

		if err != nil {
			return err
		}

		_ = in
		//log.Printf("Data received: %v", in)
	}
}

// loggingServer manages the FnAPI logging channel.
type loggingServer struct{} // no data content

func (l *loggingServer) Logging(stream fnpb.BeamFnLogging_LoggingServer) error {
	// This stream object is only used here. The stream is used for receiving, and
	// no sends happen on it.
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			return nil
		}

		if err != nil {
			return err
		}

		for _, e := range in.GetLogEntries() {
			log.Info(stream.Context(), e.GetMessage())
		}
	}
}

// Execute launches the supplied pipeline using a session file as the source of inputs.
func Execute(ctx context.Context, p *beam.Pipeline) (beam.PipelineResult, error) {
	worker, err := buildLocalBinary(ctx)
	if err != nil {
		return nil, errors.WithContext(err, "building worker binary")
	}

	log.Infof(ctx, "built worker binary at %s\n", worker)

	// Start up the grpc logging service.
	ls := grpc.NewServer()
	fnpb.RegisterBeamFnLoggingServer(ls, &loggingServer{})
	logPort, err := net.Listen("tcp", ":0")
	if err != nil {
		panic("No logging port")
	}
	go ls.Serve(logPort)

	// The wait group is used by the control service goroutine to signal
	// completion.
	var wg sync.WaitGroup
	wg.Add(1)

	cs := grpc.NewServer()
	fnpb.RegisterBeamFnControlServer(cs, &controlServer{
		filename: *sessionFile,
		wg:       &wg,
	})

	ctrlPort, err := net.Listen("tcp", ":0")
	if err != nil {
		panic("No control port")
	}
	go cs.Serve(ctrlPort)

	fmt.Println("fake harness initialized")
	cmd := exec.Command(
		worker,
		"--worker",
		fmt.Sprintf("--logging_endpoint=%s", logPort.Addr().String()),
		fmt.Sprintf("--control_endpoint=%s", ctrlPort.Addr().String()),
		"--persist_dir=/tmp/worker")
	go cmd.Start()

	wg.Wait()
	return nil, nil
}

// buildLocalBinary is cribbed from the Dataflow runner, but doesn't force the
// Linux architecture, since the worker runs in the pipeline launch
// environment.
func buildLocalBinary(ctx context.Context) (string, error) {
	ret := filepath.Join(os.TempDir(), fmt.Sprintf("session-runner-%v", time.Now().UnixNano()))

	program := ""
	for i := 3; ; i++ {
		_, file, _, ok := runtime.Caller(i)
		if !ok || strings.HasSuffix(file, "runtime/proc.go") {
			break
		}
		program = file
	}
	if program == "" {
		return "", errors.New("could not detect user main")
	}

	log.Infof(ctx, "Compiling %v as %v", program, ret)

	// Cross-compile given go program. Not awesome.
	build := []string{"go", "build", "-o", ret, program}

	cmd := exec.Command(build[0], build[1:]...)
	if out, err := cmd.CombinedOutput(); err != nil {
		log.Info(ctx, string(out))
		return "", errors.Wrapf(err, "failed to compile %v", program)
	}
	return ret, nil
}
<|endoftext|>


<commit_before>package ari

import (
	"bytes"
	"encoding/json"
	"strings"
	"testing"
	"time"
)

type dtTest struct {
	Date DateTime `json:"dt"`
}

func (dt *dtTest) Equal(o *dtTest) bool {
	return time.Time(dt.Date).Equal(time.Time(o.Date))
}

var dtMarshalTests = []struct {
	Input    dtTest
	Output   string
	HasError bool
}{
	{dtTest{DateTime(time.Date(2005, 02, 04, 13, 12, 6, 0, time.UTC))}, `{"dt":"2005-02-04T13:12:06.000+0000"}`, false},
}

func TestDateTimeMarshal(t *testing.T) {
	for _, tx := range dtMarshalTests {
		var buf bytes.Buffer
		err := json.NewEncoder(&buf).Encode(tx.Input)
		out := strings.TrimSpace(buf.String())

		failed := false
		failed = failed || (err == nil && tx.HasError)
		failed = failed || (out != tx.Output)

		if failed {
			t.Errorf("Marshal(%s) => '%s', 'err != nil => %v'; expected '%s', 'err != nil => %v'.", tx.Input, out, err != nil, tx.Output, tx.HasError)
		}
	}
}

var dtUnmarshalTests = []struct {
	Input    string
	Output   dtTest
	HasError bool
}{
	{`{"dt":"2005-02-04T13:12:06.000+0000"}`, dtTest{DateTime(time.Date(2005, 02, 04, 13, 12, 6, 0, time.UTC))}, false},
	{`{"dt":"2x05-02-04T13:12:06.000+0000"}`, dtTest{}, true},
	{`{"dt": 0 }`, dtTest{}, true},
}

func TestDateTimeUnmarshal(t *testing.T) {
	for _, tx := range dtUnmarshalTests {
		var out dtTest
		err := json.NewDecoder(strings.NewReader(tx.Input)).Decode(&out)

		failed := false
		failed = failed || (err == nil && tx.HasError)
		failed = failed || (!out.Equal(&tx.Output))

		if failed {
			t.Errorf("Unmarshal(%s) => '%s', 'err != nil => %v'; expected '%s', 'err != nil => %v'.", tx.Input, out, err != nil, tx.Output, tx.HasError)
		}
	}
}

type dsTest struct {
	Duration DurationSec `json:"ds"`
}

func (ds *dsTest) Equal(o *dsTest) bool {
	return time.Duration(ds.Duration) == time.Duration(o.Duration)
}

var dsMarshalTests = []struct {
	Input    dsTest
	Output   string
	HasError bool
}{
	{dsTest{DurationSec(4 * time.Second)}, `{"ds":4}`, false},
	{dsTest{DurationSec(40 * time.Second)}, `{"ds":40}`, false},
}

func TestDurationSecsMarshal(t *testing.T) {
	for _, tx := range dsMarshalTests {
		var buf bytes.Buffer
		err := json.NewEncoder(&buf).Encode(tx.Input)
		out := strings.TrimSpace(buf.String())

		failed := false
		failed = failed || (err == nil && tx.HasError)
		failed = failed || (out != tx.Output)

		if failed {
			t.Errorf("Marshal(%v) => '%s', 'err != nil => %v'; expected '%s', 'err != nil => %v'.", tx.Input, out, err != nil, tx.Output, tx.HasError)
		}
	}
}

var dsUnmarshalTests = []struct {
	Input    string
	Output   dsTest
	HasError bool
}{
	{`{"ds":4}`, dsTest{DurationSec(4 * time.Second)}, false},
	{`{"ds":40}`, dsTest{DurationSec(40 * time.Second)}, false},
	{`{"ds":"4"}`, dsTest{}, true},
	{`{"ds":""}`, dsTest{}, true},
	{`{"ds":"xzsad"}`, dsTest{}, true},
}

func TestDurationSecsUnmarshal(t *testing.T) {
	for _, tx := range dsUnmarshalTests {
		var out dsTest
		err := json.NewDecoder(strings.NewReader(tx.Input)).Decode(&out)

		failed := false
		failed = failed || (err == nil && tx.HasError)
		failed = failed || (!out.Equal(&tx.Output))

		if failed {
			t.Errorf("Unmarshal(%s) => '%v', 'err != nil => %v'; expected '%v', 'err != nil => %v'.", tx.Input, out, err != nil, tx.Output, tx.HasError)
		}
	}
}
<commit_msg>v3 - datetime tests - Refactor, Reorder, and remove test duplication<commit_after>package ari

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"testing"
	"time"
)

// test data

var dtMarshalTests = []struct {
	Input    dtTest
	Output   string
	HasError bool
}{
	{dtTest{DateTime(time.Date(2005, 02, 04, 13, 12, 6, 0, time.UTC))}, `{"dt":"2005-02-04T13:12:06.000+0000"}`, false},
}

var dtUnmarshalTests = []struct {
	Input    string
	Output   dtTest
	HasError bool
}{
	{`{"dt":"2005-02-04T13:12:06.000+0000"}`, dtTest{DateTime(time.Date(2005, 02, 04, 13, 12, 6, 0, time.UTC))}, false},
	{`{"dt":"2x05-02-04T13:12:06.000+0000"}`, dtTest{}, true},
	{`{"dt": 0 }`, dtTest{}, true},
}

var dsUnmarshalTests = []struct {
	Input    string
	Output   dsTest
	HasError bool
}{
	{`{"ds":4}`, dsTest{DurationSec(4 * time.Second)}, false},
	{`{"ds":40}`, dsTest{DurationSec(40 * time.Second)}, false},
	{`{"ds":"4"}`, dsTest{}, true},
	{`{"ds":""}`, dsTest{}, true},
	{`{"ds":"xzsad"}`, dsTest{}, true},
}

var dsMarshalTests = []struct {
	Input    dsTest
	Output   string
	HasError bool
}{
	{dsTest{DurationSec(4 * time.Second)}, `{"ds":4}`, false},
	{dsTest{DurationSec(40 * time.Second)}, `{"ds":40}`, false},
}

// test runners
func TestDateTimeMarshal(t *testing.T) {
	for _, tx := range dtMarshalTests {
		ret := runTestMarshal(tx.Input, tx.Output, tx.HasError)
		if ret != "" {
			t.Errorf(ret)
		}
	}
}
func TestDateTimeUnmarshal(t *testing.T) {
	for _, tx := range dtUnmarshalTests {
		var out dtTest
		ret := runTestUnmarshal(&out, tx.Input, &tx.Output, tx.HasError)
		if ret != "" {
			t.Errorf(ret)
		}
	}
}

func TestDurationSecsMarshal(t *testing.T) {
	for _, tx := range dsMarshalTests {
		ret := runTestMarshal(tx.Input, tx.Output, tx.HasError)
		if ret != "" {
			t.Errorf(ret)
		}
	}
}
func TestDurationSecsUnmarshal(t *testing.T) {
	for _, tx := range dsUnmarshalTests {
		var out dsTest
		ret := runTestUnmarshal(&out, tx.Input, &tx.Output, tx.HasError)
		if ret != "" {
			t.Errorf(ret)
		}
	}
}

// generalized test functions

func runTestMarshal(input interface{}, output string, hasError bool) (ret string) {
	var buf bytes.Buffer
	err := json.NewEncoder(&buf).Encode(input)
	out := strings.TrimSpace(buf.String())

	failed := false
	failed = failed || (err == nil && hasError)
	failed = failed || (out != output)

	if failed {
		ret = fmt.Sprintf("Marshal(%s) => '%s', 'err != nil => %v'; expected '%s', 'err != nil => %v'.", input, out, err != nil, output, hasError)
	}

	return
}

func runTestUnmarshal(out eq, input string, output eq, hasError bool) (ret string) {
	err := json.NewDecoder(strings.NewReader(input)).Decode(&out)

	failed := false
	failed = failed || (err == nil && hasError)
	failed = failed || (!out.Equal(output))

	if failed {
		ret = fmt.Sprintf("Unmarshal(%s) => '%s', 'err != nil => %v'; expected '%s', 'err != nil => %v'.", input, out, err != nil, output, hasError)
	}

	return
}

// test structures

type dtTest struct {
	Date DateTime `json:"dt"`
}

func (dt *dtTest) Equal(i interface{}) bool {
	o, ok := i.(*dtTest)
	return ok && time.Time(dt.Date).Equal(time.Time(o.Date))
}

type eq interface {
	Equal(i interface{}) bool
}

type dsTest struct {
	Duration DurationSec `json:"ds"`
}

func (ds *dsTest) Equal(i interface{}) bool {
	o, ok := i.(*dsTest)
	return ok && time.Duration(ds.Duration) == time.Duration(o.Duration)
}
<|endoftext|>


<commit_before>package types

import (
	"encoding/json"
	"fmt"
	"testing"
	"time"
)

/*
	Проверяет функции ToDateTime(), String() и StringToDateTime()
*/
func TestDateTimeStringConversion(t *testing.T) {
	timeMonthAgo := time.Now().AddDate(0, -1, 0)
	dateTime := ToDateTime(timeMonthAgo)
	dateTimeString := dateTime.String()

	dateTimeFromString, err := StringToDateTime(dateTimeString)
	if err != nil {
		t.Fatal(err)
	}

	expectedString := dateTimeFromString.String()
	if expectedString != dateTimeString {
		t.Fatalf("Ошибка %v. Ожидалось %v, получено %v", err, expectedString, dateTimeString)
	}
}

/*
	Проверяет функции ToDate(), String() и StringToDate()
*/
func TestDateStringConversion(t *testing.T) {
	timeMonthAgo := time.Now().AddDate(0, -1, 0)
	date := ToDate(timeMonthAgo)
	dateString := date.String()

	dateFromString, err := StringToDate(dateString)
	if err != nil {
		t.Fatal(err)
	}

	expectedString := dateFromString.String()
	if expectedString != dateString {
		t.Fatalf("Ошибка %v. Ожидалось %v, получено %v", err, expectedString, dateString)
	}
}

func TestDaysBefore(t *testing.T) {
	dateToday := DateNow()
	dateYesterday := ToDate(time.Now().AddDate(0, 0, -1))
	dateTomorrow := ToDate(time.Now().AddDate(0, 0, 1))

	if d := dateToday.DaysBefore(dateToday); d != 0 {
		t.Fatalf("Ожидалось 0, получено %d дней", d)
	}

	if d := dateToday.DaysBefore(dateYesterday); d != -1 {
		t.Fatalf("Ожидалось -1, получено %d дней", d)
	}

	if d := dateToday.DaysBefore(dateTomorrow); d != 1 {
		t.Fatalf("Ожидалось 1, получено %d дней", d)
	}

	if d := dateYesterday.DaysBefore(dateTomorrow); d != 2 {
		t.Fatalf("Ожидалось 2, получено %d дней", d)
	}
}

func TestStringToDateTimeHMS(t *testing.T) {
	s := "2015-07-30"
	dt, err := StringDateToDateTimeHMS(s, 5, 6, 7)
	if err != nil {
		t.Fatal(err)
	}
	sExpected := "2015-07-30 05:06:07"
	sReceived := dt.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestNeverDate(t *testing.T) {
	nd := NeverDate()
	sExpected := "1990-01-01"
	sReceived := nd.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestNeverTime(t *testing.T) {
	ndt := NeverTime()
	sExpected := "1990-01-01 00:00:00"
	sReceived := ndt.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

/*
	Проверяет возврат текущих даты, даты-времени и
	преобразование Date -> DateTime с учётом HMS
*/
func TestNow(t *testing.T) {
	dateTimeNow := DateTimeNow()
	dateNow := dateTimeNow.ConvertToDate()
	dateTimeAfterConversion := dateNow.ConvertToDateTimeHMS(
		dateTimeNow.Hour(), dateTimeNow.Minute(), dateTimeNow.Second())
	dateTimeAfterConversion.Layout = DateTimeLayout

	sExpected := dateTimeNow.String()
	sReceived := dateTimeAfterConversion.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestDateTimeTodayHMS(t *testing.T) {
	dt := DateTimeTodayHMS(23, 59, 59)
	sExpected := DateNow().String() + " 23:59:59"
	sReceived := dt.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestDateSetDefaultLayoutIfEmpty(t *testing.T) {
	d := DateNow()
	d.Layout = ""
	d.setDefaultLayoutIfEmpty()
	if d.Layout != DateLayout {
		t.Fatal("Установлен неправильный Layout")
	}
}

func TestDateTimeSetDefaultLayoutIfEmpty(t *testing.T) {
	dt := DateTimeNow()
	dt.Layout = ""
	dt.setDefaultLayoutIfEmpty()
	if dt.Layout != DateTimeLayout {
		t.Fatal("Установлен неправильный Layout")
	}
}

func TestDateBeforeAfterBetween(t *testing.T) {
	dateToday := DateNow()
	dateYesterday := ToDate(time.Now().AddDate(0, 0, -1))
	dateTomorrow := ToDate(time.Now().AddDate(0, 0, 1))

	if dateToday.Before(dateToday) {
		t.Fatal("Сегодня не может быть раньше чем сегодня")
	}
	if dateToday.After(dateToday) {
		t.Fatal("Сегодня не может быть позже чем сегодня")
	}
	if dateToday.Between(dateToday, dateTomorrow) {
		t.Fatal("Сегодня не может быть между сегодня и завтра")
	}
	if dateToday.Between(dateYesterday, dateToday) {
		t.Fatal("Сегодня не может быть между вчера и сегодня")
	}

	if dateToday.Before(dateYesterday) {
		t.Fatal("Сегодня не может быть раньше чем вчера")
	}
	if dateYesterday.After(dateToday) {
		t.Fatal("Вчера не может быть позже чем сегодня")
	}
	if dateYesterday.Between(dateToday, dateTomorrow) {
		t.Fatal("Вчера не может быть между сегодня и завтра")
	}

	if !dateToday.Before(dateTomorrow) {
		t.Fatal("Сегодня должно быть раньше чем завтра")
	}
	if !dateToday.After(dateYesterday) {
		t.Fatal("Сегодня должно быть позже чем вчера")
	}
	if !dateToday.Between(dateYesterday, dateTomorrow) {
		t.Fatal("Сегодня должно быть между вчера и завтра")
	}
}

func TestDateTimeBeforeAfterBetween(t *testing.T) {
	dateNow := DateTimeNow()
	dateMinBefore := ToDateTime(time.Now().Add(-time.Minute))
	dateMinAfter := ToDateTime(time.Now().Add(time.Minute))

	if dateNow.Before(dateNow) {
		t.Fatal("Сейчас не может быть раньше чем сейчас")
	}
	if dateNow.After(dateNow) {
		t.Fatal("Сейчас не может быть позже чем сейчас")
	}
	if dateNow.Between(dateNow, dateMinAfter) {
		t.Fatal("Сейчас не может быть между сейчас и через минуту")
	}
	if dateNow.Between(dateMinBefore, dateNow) {
		t.Fatal("Сейчас не может быть между минуту назад и сейчас")
	}

	if dateNow.Before(dateMinBefore) {
		t.Fatal("Сейчас не может быть раньше чем минуту назад")
	}
	if dateMinBefore.After(dateNow) {
		t.Fatal("Минуту назад не может быть позже чем сейчас")
	}
	if dateMinBefore.Between(dateNow, dateMinAfter) {
		t.Fatal("Минуту назад не может быть между сейчас и через минуту")
	}

	if !dateNow.Before(dateMinAfter) {
		t.Fatal("Сейчас должно быть раньше чем через минуту")
	}
	if !dateNow.After(dateMinBefore) {
		t.Fatal("Сейчас должно быть позже чем минуту назад")
	}
	if !dateNow.Between(dateMinBefore, dateMinAfter) {
		t.Fatal("Сейчас должно быть между минуту назад и через минуту")
	}
}

func TestDateTimeJSON(t *testing.T) {
	dtString := "2015-07-30 20:58:59"
	jsonExpected := fmt.Sprintf(`"%s"`, dtString)
	dt, err := StringToDateTime(dtString)
	if err != nil {
		t.Fatal(err)
	}
	b, err := json.Marshal(dt)
	if err != nil {
		t.Fatal(err)
	}
	jsonReceived := string(b)
	if jsonExpected != jsonReceived {
		t.Fatalf("Ошибка Marshal. Ожидалось получить JSON %s, получен JSON %s", jsonExpected, jsonReceived)
	}

	var dtFromJson DateTime
	if err := json.Unmarshal(b, &dtFromJson); err != nil {
		t.Fatal(err)
	}
	dtFromJsonString := dtFromJson.String()
	if dtFromJsonString != dtString {
		t.Fatalf("Ошибка Unmarshal. Ожидалось получить JSON %s, получен JSON %s", dtString, dtFromJsonString)
	}
}

func TestDateTimeUnmarshalBadJSON(t *testing.T) {
	b := []byte(`"wrong"`)
	var dtFromJson DateTime
	if err := json.Unmarshal(b, &dtFromJson); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}

func TestDateJSON(t *testing.T) {
	dString := "2015-07-30"
	jsonExpected := fmt.Sprintf(`"%s"`, dString)
	d, err := StringToDate(dString)
	if err != nil {
		t.Fatal(err)
	}
	b, err := json.Marshal(d)
	if err != nil {
		t.Fatal(err)
	}
	jsonReceived := string(b)
	if jsonExpected != jsonReceived {
		t.Fatalf("Ошибка Marshal. Ожидалось получить JSON %s, получен JSON %s", jsonExpected, jsonReceived)
	}

	var dFromJson Date
	if err := json.Unmarshal(b, &dFromJson); err != nil {
		t.Fatal(err)
	}
	dFromJsonString := dFromJson.String()
	if dFromJsonString != dString {
		t.Fatalf("Ошибка Unmarshal. Ожидалось получить JSON %s, получен JSON %s", dString, dFromJsonString)
	}
}

func TestDateUnmarshalBadJSON(t *testing.T) {
	b := []byte(`"wrong"`)
	var dFromJson Date
	if err := json.Unmarshal(b, &dFromJson); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}

func TestDateTimeScanValueForDB(t *testing.T) {
	dt := DateTimeNow()
	v, err := dt.Value()
	if err != nil {
		t.Fatal(err)
	}
	vString, ok := v.(string)
	if !ok {
		t.Fatal("Ошибка преобразования интерфейса")
	}
	vdt, err := StringToDateTime(vString)
	if err != nil {
		t.Fatal(err)
	}
	var dtFromScan DateTime
	if err := dtFromScan.Scan(vdt.Time); err != nil {
		t.Fatal(err)
	}
	sExpected := dt.String()
	sReceived := dtFromScan.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestDateTimeScanIfBadValue(t *testing.T) {
	var dtFromScan DateTime

	if err := dtFromScan.Scan(nil); err != nil {
		t.Fatal(err)
	}

	if err := dtFromScan.Scan("wrong"); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}

func TestDateScanValueForDB(t *testing.T) {
	d := DateNow()
	v, err := d.Value()
	if err != nil {
		t.Fatal(err)
	}
	vString, ok := v.(string)
	if !ok {
		t.Fatal("Ошибка преобразования интерфейса")
	}
	vd, err := StringToDate(vString)
	if err != nil {
		t.Fatal(err)
	}
	var dFromScan Date
	if err := dFromScan.Scan(vd.Time); err != nil {
		t.Fatal(err)
	}
	sExpected := d.String()
	sReceived := dFromScan.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestDateScanIfBadValue(t *testing.T) {
	var dFromScan Date

	if err := dFromScan.Scan(nil); err != nil {
		t.Fatal(err)
	}

	if err := dFromScan.Scan("wrong"); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}
<commit_msg>Ещё парочка тестов<commit_after>package types

import (
	"encoding/json"
	"fmt"
	"testing"
	"time"
)

/*
	Проверяет функции ToDateTime(), String() и StringToDateTime()
*/
func TestDateTimeStringConversion(t *testing.T) {
	timeMonthAgo := time.Now().AddDate(0, -1, 0)
	dateTime := ToDateTime(timeMonthAgo)
	dateTimeString := dateTime.String()

	dateTimeFromString, err := StringToDateTime(dateTimeString)
	if err != nil {
		t.Fatal(err)
	}

	expectedString := dateTimeFromString.String()
	if expectedString != dateTimeString {
		t.Fatalf("Ошибка %v. Ожидалось %v, получено %v", err, expectedString, dateTimeString)
	}
}

func TestBadStringToDateTime(t *testing.T) {
	if _, err := StringToDateTime("wrong"); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}

func TestBadStringToDateTimeHMS(t *testing.T) {
	if _, err := StringDateToDateTimeHMS("wrong", 23, 59, 59); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}

/*
	Проверяет функции ToDate(), String() и StringToDate()
*/
func TestDateStringConversion(t *testing.T) {
	timeMonthAgo := time.Now().AddDate(0, -1, 0)
	date := ToDate(timeMonthAgo)
	dateString := date.String()

	dateFromString, err := StringToDate(dateString)
	if err != nil {
		t.Fatal(err)
	}

	expectedString := dateFromString.String()
	if expectedString != dateString {
		t.Fatalf("Ошибка %v. Ожидалось %v, получено %v", err, expectedString, dateString)
	}
}

func TestBadStringToDate(t *testing.T) {
	if _, err := StringToDate("wrong"); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}

func TestDaysBefore(t *testing.T) {
	dateToday := DateNow()
	dateYesterday := ToDate(time.Now().AddDate(0, 0, -1))
	dateTomorrow := ToDate(time.Now().AddDate(0, 0, 1))

	if d := dateToday.DaysBefore(dateToday); d != 0 {
		t.Fatalf("Ожидалось 0, получено %d дней", d)
	}

	if d := dateToday.DaysBefore(dateYesterday); d != -1 {
		t.Fatalf("Ожидалось -1, получено %d дней", d)
	}

	if d := dateToday.DaysBefore(dateTomorrow); d != 1 {
		t.Fatalf("Ожидалось 1, получено %d дней", d)
	}

	if d := dateYesterday.DaysBefore(dateTomorrow); d != 2 {
		t.Fatalf("Ожидалось 2, получено %d дней", d)
	}
}

func TestStringToDateTimeHMS(t *testing.T) {
	s := "2015-07-30"
	dt, err := StringDateToDateTimeHMS(s, 5, 6, 7)
	if err != nil {
		t.Fatal(err)
	}
	sExpected := "2015-07-30 05:06:07"
	sReceived := dt.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestNeverDate(t *testing.T) {
	nd := NeverDate()
	sExpected := "1990-01-01"
	sReceived := nd.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestNeverTime(t *testing.T) {
	ndt := NeverTime()
	sExpected := "1990-01-01 00:00:00"
	sReceived := ndt.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

/*
	Проверяет возврат текущих даты, даты-времени и
	преобразование Date -> DateTime с учётом HMS
*/
func TestNow(t *testing.T) {
	dateTimeNow := DateTimeNow()
	dateNow := dateTimeNow.ConvertToDate()
	dateTimeAfterConversion := dateNow.ConvertToDateTimeHMS(
		dateTimeNow.Hour(), dateTimeNow.Minute(), dateTimeNow.Second())
	dateTimeAfterConversion.Layout = DateTimeLayout

	sExpected := dateTimeNow.String()
	sReceived := dateTimeAfterConversion.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestDateTimeTodayHMS(t *testing.T) {
	dt := DateTimeTodayHMS(23, 59, 59)
	sExpected := DateNow().String() + " 23:59:59"
	sReceived := dt.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestDateSetDefaultLayoutIfEmpty(t *testing.T) {
	d := DateNow()
	d.Layout = ""
	d.setDefaultLayoutIfEmpty()
	if d.Layout != DateLayout {
		t.Fatal("Установлен неправильный Layout")
	}
}

func TestDateTimeSetDefaultLayoutIfEmpty(t *testing.T) {
	dt := DateTimeNow()
	dt.Layout = ""
	dt.setDefaultLayoutIfEmpty()
	if dt.Layout != DateTimeLayout {
		t.Fatal("Установлен неправильный Layout")
	}
}

func TestDateBeforeAfterBetween(t *testing.T) {
	dateToday := DateNow()
	dateYesterday := ToDate(time.Now().AddDate(0, 0, -1))
	dateTomorrow := ToDate(time.Now().AddDate(0, 0, 1))

	if dateToday.Before(dateToday) {
		t.Fatal("Сегодня не может быть раньше чем сегодня")
	}
	if dateToday.After(dateToday) {
		t.Fatal("Сегодня не может быть позже чем сегодня")
	}
	if dateToday.Between(dateToday, dateTomorrow) {
		t.Fatal("Сегодня не может быть между сегодня и завтра")
	}
	if dateToday.Between(dateYesterday, dateToday) {
		t.Fatal("Сегодня не может быть между вчера и сегодня")
	}

	if dateToday.Before(dateYesterday) {
		t.Fatal("Сегодня не может быть раньше чем вчера")
	}
	if dateYesterday.After(dateToday) {
		t.Fatal("Вчера не может быть позже чем сегодня")
	}
	if dateYesterday.Between(dateToday, dateTomorrow) {
		t.Fatal("Вчера не может быть между сегодня и завтра")
	}

	if !dateToday.Before(dateTomorrow) {
		t.Fatal("Сегодня должно быть раньше чем завтра")
	}
	if !dateToday.After(dateYesterday) {
		t.Fatal("Сегодня должно быть позже чем вчера")
	}
	if !dateToday.Between(dateYesterday, dateTomorrow) {
		t.Fatal("Сегодня должно быть между вчера и завтра")
	}
}

func TestDateTimeBeforeAfterBetween(t *testing.T) {
	dateNow := DateTimeNow()
	dateMinBefore := ToDateTime(time.Now().Add(-time.Minute))
	dateMinAfter := ToDateTime(time.Now().Add(time.Minute))

	if dateNow.Before(dateNow) {
		t.Fatal("Сейчас не может быть раньше чем сейчас")
	}
	if dateNow.After(dateNow) {
		t.Fatal("Сейчас не может быть позже чем сейчас")
	}
	if dateNow.Between(dateNow, dateMinAfter) {
		t.Fatal("Сейчас не может быть между сейчас и через минуту")
	}
	if dateNow.Between(dateMinBefore, dateNow) {
		t.Fatal("Сейчас не может быть между минуту назад и сейчас")
	}

	if dateNow.Before(dateMinBefore) {
		t.Fatal("Сейчас не может быть раньше чем минуту назад")
	}
	if dateMinBefore.After(dateNow) {
		t.Fatal("Минуту назад не может быть позже чем сейчас")
	}
	if dateMinBefore.Between(dateNow, dateMinAfter) {
		t.Fatal("Минуту назад не может быть между сейчас и через минуту")
	}

	if !dateNow.Before(dateMinAfter) {
		t.Fatal("Сейчас должно быть раньше чем через минуту")
	}
	if !dateNow.After(dateMinBefore) {
		t.Fatal("Сейчас должно быть позже чем минуту назад")
	}
	if !dateNow.Between(dateMinBefore, dateMinAfter) {
		t.Fatal("Сейчас должно быть между минуту назад и через минуту")
	}
}

func TestDateTimeJSON(t *testing.T) {
	dtString := "2015-07-30 20:58:59"
	jsonExpected := fmt.Sprintf(`"%s"`, dtString)
	dt, err := StringToDateTime(dtString)
	if err != nil {
		t.Fatal(err)
	}
	b, err := json.Marshal(dt)
	if err != nil {
		t.Fatal(err)
	}
	jsonReceived := string(b)
	if jsonExpected != jsonReceived {
		t.Fatalf("Ошибка Marshal. Ожидалось получить JSON %s, получен JSON %s", jsonExpected, jsonReceived)
	}

	var dtFromJson DateTime
	if err := json.Unmarshal(b, &dtFromJson); err != nil {
		t.Fatal(err)
	}
	dtFromJsonString := dtFromJson.String()
	if dtFromJsonString != dtString {
		t.Fatalf("Ошибка Unmarshal. Ожидалось получить JSON %s, получен JSON %s", dtString, dtFromJsonString)
	}
}

func TestDateTimeUnmarshalBadJSON(t *testing.T) {
	b := []byte(`"wrong"`)
	var dtFromJson DateTime
	if err := json.Unmarshal(b, &dtFromJson); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}

func TestDateJSON(t *testing.T) {
	dString := "2015-07-30"
	jsonExpected := fmt.Sprintf(`"%s"`, dString)
	d, err := StringToDate(dString)
	if err != nil {
		t.Fatal(err)
	}
	b, err := json.Marshal(d)
	if err != nil {
		t.Fatal(err)
	}
	jsonReceived := string(b)
	if jsonExpected != jsonReceived {
		t.Fatalf("Ошибка Marshal. Ожидалось получить JSON %s, получен JSON %s", jsonExpected, jsonReceived)
	}

	var dFromJson Date
	if err := json.Unmarshal(b, &dFromJson); err != nil {
		t.Fatal(err)
	}
	dFromJsonString := dFromJson.String()
	if dFromJsonString != dString {
		t.Fatalf("Ошибка Unmarshal. Ожидалось получить JSON %s, получен JSON %s", dString, dFromJsonString)
	}
}

func TestDateUnmarshalBadJSON(t *testing.T) {
	b := []byte(`"wrong"`)
	var dFromJson Date
	if err := json.Unmarshal(b, &dFromJson); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}

func TestDateTimeScanValueForDB(t *testing.T) {
	dt := DateTimeNow()
	v, err := dt.Value()
	if err != nil {
		t.Fatal(err)
	}
	vString, ok := v.(string)
	if !ok {
		t.Fatal("Ошибка преобразования интерфейса")
	}
	vdt, err := StringToDateTime(vString)
	if err != nil {
		t.Fatal(err)
	}
	var dtFromScan DateTime
	if err := dtFromScan.Scan(vdt.Time); err != nil {
		t.Fatal(err)
	}
	sExpected := dt.String()
	sReceived := dtFromScan.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestDateTimeScanIfBadValue(t *testing.T) {
	var dtFromScan DateTime

	if err := dtFromScan.Scan(nil); err != nil {
		t.Fatal(err)
	}

	if err := dtFromScan.Scan("wrong"); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}

func TestDateScanValueForDB(t *testing.T) {
	d := DateNow()
	v, err := d.Value()
	if err != nil {
		t.Fatal(err)
	}
	vString, ok := v.(string)
	if !ok {
		t.Fatal("Ошибка преобразования интерфейса")
	}
	vd, err := StringToDate(vString)
	if err != nil {
		t.Fatal(err)
	}
	var dFromScan Date
	if err := dFromScan.Scan(vd.Time); err != nil {
		t.Fatal(err)
	}
	sExpected := d.String()
	sReceived := dFromScan.String()
	if sExpected != sReceived {
		t.Fatalf("Ожидалось получить строку %s, получена строка %s", sExpected, sReceived)
	}
}

func TestDateScanIfBadValue(t *testing.T) {
	var dFromScan Date

	if err := dFromScan.Scan(nil); err != nil {
		t.Fatal(err)
	}

	if err := dFromScan.Scan("wrong"); err == nil {
		t.Fatal("Ожидалась ошибка")
	}
}
<|endoftext|>


<commit_before>package db

import (
	"github.com/deiwin/luncher-api/db/model"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type Restaurants interface {
	Insert(...*model.Restaurant) ([]*model.Restaurant, error)
	Get() ([]*model.Restaurant, error)
	GetByID(bson.ObjectId) (*model.Restaurant, error)
	Exists(name string) (bool, error)
}

type restaurantsCollection struct {
	c *mgo.Collection
}

func NewRestaurants(client *Client) Restaurants {
	collection := client.database.C(model.RestaurantCollectionName)
	return &restaurantsCollection{collection}
}

func (c restaurantsCollection) Insert(restaurantsToInsert ...*model.Restaurant) ([]*model.Restaurant, error) {
	for _, restaurant := range restaurantsToInsert {
		if restaurant.ID == "" {
			restaurant.ID = bson.NewObjectId()
		}
	}
	docs := make([]interface{}, len(restaurantsToInsert))
	for i, restaurant := range restaurantsToInsert {
		docs[i] = restaurant
	}
	return restaurantsToInsert, c.c.Insert(docs...)
}

func (c restaurantsCollection) Get() (restaurants []*model.Restaurant, err error) {
	err = c.c.Find(bson.M{}).All(&restaurants)
	return
}

func (c restaurantsCollection) GetByID(id bson.ObjectId) (*model.Restaurant, error) {
	var restaurant *model.Restaurant
	err := c.c.FindId(id).One(&restaurant)
	return restaurant, err
}

func (c restaurantsCollection) Exists(name string) (bool, error) {
	count, err := c.c.Find(bson.M{"name": name}).Count()
	if err != nil {
		return false, err
	}
	return count != 0, nil
}
<commit_msg>fix another nil pointer issue<commit_after>package db

import (
	"github.com/deiwin/luncher-api/db/model"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type Restaurants interface {
	Insert(...*model.Restaurant) ([]*model.Restaurant, error)
	Get() ([]*model.Restaurant, error)
	GetByID(bson.ObjectId) (*model.Restaurant, error)
	Exists(name string) (bool, error)
}

type restaurantsCollection struct {
	c *mgo.Collection
}

func NewRestaurants(client *Client) Restaurants {
	collection := client.database.C(model.RestaurantCollectionName)
	return &restaurantsCollection{collection}
}

func (c restaurantsCollection) Insert(restaurantsToInsert ...*model.Restaurant) ([]*model.Restaurant, error) {
	for _, restaurant := range restaurantsToInsert {
		if restaurant.ID == "" {
			restaurant.ID = bson.NewObjectId()
		}
	}
	docs := make([]interface{}, len(restaurantsToInsert))
	for i, restaurant := range restaurantsToInsert {
		docs[i] = restaurant
	}
	return restaurantsToInsert, c.c.Insert(docs...)
}

func (c restaurantsCollection) Get() (restaurants []*model.Restaurant, err error) {
	err = c.c.Find(bson.M{}).All(&restaurants)
	return
}

func (c restaurantsCollection) GetByID(id bson.ObjectId) (*model.Restaurant, error) {
	var restaurant model.Restaurant
	err := c.c.FindId(id).One(&restaurant)
	return &restaurant, err
}

func (c restaurantsCollection) Exists(name string) (bool, error) {
	count, err := c.c.Find(bson.M{"name": name}).Count()
	if err != nil {
		return false, err
	}
	return count != 0, nil
}
<|endoftext|>


<commit_before>package typhon

import (
	"testing"

	"github.com/monzo/terrors"
	"github.com/stretchr/testify/assert"
)

func TestRouter(t *testing.T) {
	t.Parallel()

	router := NewRouter()
	router.GET("/foo", func(req Request) Response {
		rsp := NewResponse(req)
		rsp.Write([]byte("abcdef"))
		return rsp
	})

	// Matching path
	req := NewRequest(nil, "GET", "/foo", nil)
	rsp := router.Serve()(req)
	assert.NoError(t, rsp.Error)
	b, _ := rsp.BodyBytes(true)
	assert.Equal(t, "abcdef", string(b))

	// Wrong method should result in not found
	// @TODO: This should really be HTTP Method Not Found
	req = NewRequest(nil, "POST", "/foo", nil)
	rsp = router.Serve()(req)
	assert.Error(t, rsp.Error)
	err := terrors.Wrap(rsp.Error, nil).(*terrors.Error)
	assert.True(t, err.Matches(terrors.ErrNotFound))

	// Wrong path should result in not found
	req = NewRequest(nil, "GET", "/", nil)
	rsp = router.Serve()(req)
	assert.Error(t, rsp.Error)
	err = terrors.Wrap(rsp.Error, nil).(*terrors.Error)
	assert.True(t, err.Matches(terrors.ErrNotFound))
}

func TestRouter_CatchallPath(t *testing.T) {
	t.Parallel()

	// Registering a global handler should apply to all paths
	router := NewRouter()
	router.GET("/*residual", func(req Request) Response {
		rsp := NewResponse(req)
		rsp.Write([]byte("catchall"))
		return rsp
	})
	req := NewRequest(nil, "GET", "/bar/baz/doodad/123/abc", nil)
	rsp := router.Serve()(req)
	assert.NoError(t, rsp.Error)
	b, _ := rsp.BodyBytes(true)
	assert.Equal(t, "catchall", string(b))
	// …but not on another method
	req = NewRequest(nil, "POST", "/foo", nil)
	rsp = router.Serve()(req)
	assert.Error(t, rsp.Error)
	err := terrors.Wrap(rsp.Error, nil).(*terrors.Error)
	assert.True(t, err.Matches(terrors.ErrNotFound))
}

// Validate that partial matches are resolved correctly
func TestRouterPartials(t *testing.T) {
	router := NewRouter()
	router.GET("/foo/", func(req Request) Response {
		return req.Response("/foo/")
	})
	router.GET("/:param/bar", func(req Request) Response {
		return req.Response("/:param/bar")
	})

	req := NewRequest(nil, "GET", "/foo/", nil)
	rsp := router.Serve()(req)
	s := ""
	rsp.Decode(&s)
	assert.Equal(t, "/foo/", s)

	req = NewRequest(nil, "GET", "/abc/bar", nil)
	rsp = router.Serve()(req)
	rsp.Decode(&s)
	assert.Equal(t, "/:param/bar", s)
}
<commit_msg>Parallel test<commit_after>package typhon

import (
	"testing"

	"github.com/monzo/terrors"
	"github.com/stretchr/testify/assert"
)

func TestRouter(t *testing.T) {
	t.Parallel()

	router := NewRouter()
	router.GET("/foo", func(req Request) Response {
		rsp := NewResponse(req)
		rsp.Write([]byte("abcdef"))
		return rsp
	})

	// Matching path
	req := NewRequest(nil, "GET", "/foo", nil)
	rsp := router.Serve()(req)
	assert.NoError(t, rsp.Error)
	b, _ := rsp.BodyBytes(true)
	assert.Equal(t, "abcdef", string(b))

	// Wrong method should result in not found
	// @TODO: This should really be HTTP Method Not Found
	req = NewRequest(nil, "POST", "/foo", nil)
	rsp = router.Serve()(req)
	assert.Error(t, rsp.Error)
	err := terrors.Wrap(rsp.Error, nil).(*terrors.Error)
	assert.True(t, err.Matches(terrors.ErrNotFound))

	// Wrong path should result in not found
	req = NewRequest(nil, "GET", "/", nil)
	rsp = router.Serve()(req)
	assert.Error(t, rsp.Error)
	err = terrors.Wrap(rsp.Error, nil).(*terrors.Error)
	assert.True(t, err.Matches(terrors.ErrNotFound))
}

func TestRouter_CatchallPath(t *testing.T) {
	t.Parallel()

	// Registering a global handler should apply to all paths
	router := NewRouter()
	router.GET("/*residual", func(req Request) Response {
		rsp := NewResponse(req)
		rsp.Write([]byte("catchall"))
		return rsp
	})
	req := NewRequest(nil, "GET", "/bar/baz/doodad/123/abc", nil)
	rsp := router.Serve()(req)
	assert.NoError(t, rsp.Error)
	b, _ := rsp.BodyBytes(true)
	assert.Equal(t, "catchall", string(b))
	// …but not on another method
	req = NewRequest(nil, "POST", "/foo", nil)
	rsp = router.Serve()(req)
	assert.Error(t, rsp.Error)
	err := terrors.Wrap(rsp.Error, nil).(*terrors.Error)
	assert.True(t, err.Matches(terrors.ErrNotFound))
}

// Validate that partial matches are resolved correctly
func TestRouterPartials(t *testing.T) {
	t.Parallel()

	router := NewRouter()
	router.GET("/foo/", func(req Request) Response {
		return req.Response("/foo/")
	})
	router.GET("/:param/bar", func(req Request) Response {
		return req.Response("/:param/bar")
	})

	req := NewRequest(nil, "GET", "/foo/", nil)
	rsp := router.Serve()(req)
	s := ""
	rsp.Decode(&s)
	assert.Equal(t, "/foo/", s)

	req = NewRequest(nil, "GET", "/abc/bar", nil)
	rsp = router.Serve()(req)
	rsp.Decode(&s)
	assert.Equal(t, "/:param/bar", s)
}
<|endoftext|>