<commit_before>package telebot

import (
	"encoding/json"
)

// Album lets you group multiple media (so-called InputMedia)
// into a single messsage.
//
// On older clients albums look like N regular messages.
type Album []InputMedia

// InputMedia is a generic type for all kinds of media you
// can put into an album.
type InputMedia interface {
	// As some files must be uploaded (instead of referencing)
	// outer layers of Telebot require it.
	MediaFile() *File
}

// Photo object represents a single photo file.
type Photo struct {
	File

	Width  int `json:"width"`
	Height int `json:"height"`

	// (Optional)
	Caption   string    `json:"caption,omitempty"`
	ParseMode ParseMode `json:"parse_mode,omitempty"`
}

type photoSize struct {
	File
	Width   int    `json:"width"`
	Height  int    `json:"height"`
	Caption string `json:"caption,omitempty"`
}

// MediaFile returns &Photo.File
func (p *Photo) MediaFile() *File {
	return &p.File
}

// UnmarshalJSON is custom unmarshaller required to abstract
// away the hassle of treating different thumbnail sizes.
// Instead, Telebot chooses the hi-res one and just sticks to
// it.
//
// I really do find it a beautiful solution.
func (p *Photo) UnmarshalJSON(jsonStr []byte) error {
	var hq photoSize

	if jsonStr[0] == '{' {
		if err := json.Unmarshal(jsonStr, &hq); err != nil {
			return err
		}
	} else {
		var sizes []photoSize

		if err := json.Unmarshal(jsonStr, &sizes); err != nil {
			return err
		}

		hq = sizes[len(sizes)-1]
	}

	p.File = hq.File
	p.Width = hq.Width
	p.Height = hq.Height

	return nil
}

// Audio object represents an audio file.
type Audio struct {
	File

	// Duration of the recording in seconds as defined by sender.
	Duration int `json:"duration,omitempty"`

	// (Optional)
	Caption   string `json:"caption,omitempty"`
	Thumbnail *Photo `json:"thumb,omitempty"`
	Title     string `json:"title,omitempty"`
	Performer string `json:"performer,omitempty"`
	MIME      string `json:"mime_type,omitempty"`
	FileName  string `json:"file_name,omitempty"`
}

// MediaFile returns &Audio.File
func (a *Audio) MediaFile() *File {
	return &a.File
}

// Document object represents a general file (as opposed to Photo or Audio).
// Telegram users can send files of any type of up to 1.5 GB in size.
type Document struct {
	File

	// (Optional)
	Thumbnail *Photo `json:"thumb,omitempty"`
	Caption   string `json:"caption,omitempty"`
	MIME      string `json:"mime_type"`
	FileName  string `json:"file_name,omitempty"`
}

// MediaFile returns &Document.File
func (d *Document) MediaFile() *File {
	return &d.File
}

// Video object represents a video file.
type Video struct {
	File

	Width  int `json:"width"`
	Height int `json:"height"`

	Duration int `json:"duration,omitempty"`

	// (Optional)
	Caption           string `json:"caption,omitempty"`
	Thumbnail         *Photo `json:"thumb,omitempty"`
	SupportsStreaming bool   `json:"supports_streaming,omitempty"`
	MIME              string `json:"mime_type,omitempty"`
	FileName          string `json:"file_name,omitempty"`
}

// MediaFile returns &Video.File
func (v *Video) MediaFile() *File {
	return &v.File
}

// Animation object represents a animation file.
type Animation struct {
	File

	Width  int `json:"width"`
	Height int `json:"height"`

	Duration int `json:"duration,omitempty"`

	// (Optional)
	Caption   string `json:"caption,omitempty"`
	Thumbnail *Photo `json:"thumb,omitempty"`
	MIME      string `json:"mime_type,omitempty"`
	FileName  string `json:"file_name,omitempty"`
}

// MediaFile returns &Animation.File
func (a *Animation) MediaFile() *File {
	return &a.File
}

// Voice object represents a voice note.
type Voice struct {
	File

	// Duration of the recording in seconds as defined by sender.
	Duration int `json:"duration"`

	// (Optional)
	MIME string `json:"mime_type,omitempty"`
}

// VideoNote represents a video message (available in Telegram apps
// as of v.4.0).
type VideoNote struct {
	File

	// Duration of the recording in seconds as defined by sender.
	Duration int `json:"duration"`

	// (Optional)
	Thumbnail *Photo `json:"thumb,omitempty"`
	Length    int    `json:"length,omitempty"`
}

// Contact object represents a contact to Telegram user
type Contact struct {
	PhoneNumber string `json:"phone_number"`
	FirstName   string `json:"first_name"`

	// (Optional)
	LastName string `json:"last_name"`
	UserID   int    `json:"user_id,omitempty"`
}

// Location object represents geographic position.
type Location struct {
	// Latitude
	Lat float32 `json:"latitude"`
	// Longitude
	Lng float32 `json:"longitude"`

	// Period in seconds for which the location will be updated
	// (see Live Locations, should be between 60 and 86400.)
	LivePeriod int `json:"live_period,omitempty"`
}

// Venue object represents a venue location with name, address and
// optional foursquare ID.
type Venue struct {
	Location Location `json:"location"`
	Title    string   `json:"title"`
	Address  string   `json:"address"`

	// (Optional)
	FoursquareID   string `json:"foursquare_id,omitempty"`
	FoursquareType string `json:"foursquare_type,omitempty"`
}

type Dice struct {
	Type  DiceType `json:"emoji"`
	Value int      `json:"value"`
}
<commit_msg>docs: add missing Dice documentation<commit_after>package telebot

import (
	"encoding/json"
)

// Album lets you group multiple media (so-called InputMedia)
// into a single messsage.
//
// On older clients albums look like N regular messages.
type Album []InputMedia

// InputMedia is a generic type for all kinds of media you
// can put into an album.
type InputMedia interface {
	// As some files must be uploaded (instead of referencing)
	// outer layers of Telebot require it.
	MediaFile() *File
}

// Photo object represents a single photo file.
type Photo struct {
	File

	Width  int `json:"width"`
	Height int `json:"height"`

	// (Optional)
	Caption   string    `json:"caption,omitempty"`
	ParseMode ParseMode `json:"parse_mode,omitempty"`
}

type photoSize struct {
	File
	Width   int    `json:"width"`
	Height  int    `json:"height"`
	Caption string `json:"caption,omitempty"`
}

// MediaFile returns &Photo.File
func (p *Photo) MediaFile() *File {
	return &p.File
}

// UnmarshalJSON is custom unmarshaller required to abstract
// away the hassle of treating different thumbnail sizes.
// Instead, Telebot chooses the hi-res one and just sticks to
// it.
//
// I really do find it a beautiful solution.
func (p *Photo) UnmarshalJSON(jsonStr []byte) error {
	var hq photoSize

	if jsonStr[0] == '{' {
		if err := json.Unmarshal(jsonStr, &hq); err != nil {
			return err
		}
	} else {
		var sizes []photoSize

		if err := json.Unmarshal(jsonStr, &sizes); err != nil {
			return err
		}

		hq = sizes[len(sizes)-1]
	}

	p.File = hq.File
	p.Width = hq.Width
	p.Height = hq.Height

	return nil
}

// Audio object represents an audio file.
type Audio struct {
	File

	// Duration of the recording in seconds as defined by sender.
	Duration int `json:"duration,omitempty"`

	// (Optional)
	Caption   string `json:"caption,omitempty"`
	Thumbnail *Photo `json:"thumb,omitempty"`
	Title     string `json:"title,omitempty"`
	Performer string `json:"performer,omitempty"`
	MIME      string `json:"mime_type,omitempty"`
	FileName  string `json:"file_name,omitempty"`
}

// MediaFile returns &Audio.File
func (a *Audio) MediaFile() *File {
	return &a.File
}

// Document object represents a general file (as opposed to Photo or Audio).
// Telegram users can send files of any type of up to 1.5 GB in size.
type Document struct {
	File

	// (Optional)
	Thumbnail *Photo `json:"thumb,omitempty"`
	Caption   string `json:"caption,omitempty"`
	MIME      string `json:"mime_type"`
	FileName  string `json:"file_name,omitempty"`
}

// MediaFile returns &Document.File
func (d *Document) MediaFile() *File {
	return &d.File
}

// Video object represents a video file.
type Video struct {
	File

	Width  int `json:"width"`
	Height int `json:"height"`

	Duration int `json:"duration,omitempty"`

	// (Optional)
	Caption           string `json:"caption,omitempty"`
	Thumbnail         *Photo `json:"thumb,omitempty"`
	SupportsStreaming bool   `json:"supports_streaming,omitempty"`
	MIME              string `json:"mime_type,omitempty"`
	FileName          string `json:"file_name,omitempty"`
}

// MediaFile returns &Video.File
func (v *Video) MediaFile() *File {
	return &v.File
}

// Animation object represents a animation file.
type Animation struct {
	File

	Width  int `json:"width"`
	Height int `json:"height"`

	Duration int `json:"duration,omitempty"`

	// (Optional)
	Caption   string `json:"caption,omitempty"`
	Thumbnail *Photo `json:"thumb,omitempty"`
	MIME      string `json:"mime_type,omitempty"`
	FileName  string `json:"file_name,omitempty"`
}

// MediaFile returns &Animation.File
func (a *Animation) MediaFile() *File {
	return &a.File
}

// Voice object represents a voice note.
type Voice struct {
	File

	// Duration of the recording in seconds as defined by sender.
	Duration int `json:"duration"`

	// (Optional)
	MIME string `json:"mime_type,omitempty"`
}

// VideoNote represents a video message (available in Telegram apps
// as of v.4.0).
type VideoNote struct {
	File

	// Duration of the recording in seconds as defined by sender.
	Duration int `json:"duration"`

	// (Optional)
	Thumbnail *Photo `json:"thumb,omitempty"`
	Length    int    `json:"length,omitempty"`
}

// Contact object represents a contact to Telegram user
type Contact struct {
	PhoneNumber string `json:"phone_number"`
	FirstName   string `json:"first_name"`

	// (Optional)
	LastName string `json:"last_name"`
	UserID   int    `json:"user_id,omitempty"`
}

// Location object represents geographic position.
type Location struct {
	// Latitude
	Lat float32 `json:"latitude"`
	// Longitude
	Lng float32 `json:"longitude"`

	// Period in seconds for which the location will be updated
	// (see Live Locations, should be between 60 and 86400.)
	LivePeriod int `json:"live_period,omitempty"`
}

// Venue object represents a venue location with name, address and
// optional foursquare ID.
type Venue struct {
	Location Location `json:"location"`
	Title    string   `json:"title"`
	Address  string   `json:"address"`

	// (Optional)
	FoursquareID   string `json:"foursquare_id,omitempty"`
	FoursquareType string `json:"foursquare_type,omitempty"`
}

// Dice object represents a dice with a random value
// from 1 to 6 for currently supported base emoji.
type Dice struct {
	Type  DiceType `json:"emoji"`
	Value int      `json:"value"`
}
<|endoftext|>


<commit_before>// Copyright 2018 Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package api

import (
	"errors"
	"fmt"
	"net/http"
	"strconv"

	common_model "github.com/goharbor/harbor/src/common/models"
	"github.com/goharbor/harbor/src/replication"
	"github.com/goharbor/harbor/src/replication/dao/models"
	"github.com/goharbor/harbor/src/replication/event"
	"github.com/goharbor/harbor/src/replication/model"
	"github.com/goharbor/harbor/src/replication/registry"
)

// TODO rename the file to "replication.go"

// ReplicationPolicyAPI handles the replication policy requests
type ReplicationPolicyAPI struct {
	BaseController
}

// Prepare ...
func (r *ReplicationPolicyAPI) Prepare() {
	r.BaseController.Prepare()
	if !r.SecurityCtx.IsSysAdmin() {
		if !r.SecurityCtx.IsAuthenticated() {
			r.SendUnAuthorizedError(errors.New("UnAuthorized"))
			return
		}
		r.SendForbiddenError(errors.New(r.SecurityCtx.GetUsername()))
		return
	}
}

// List the replication policies
func (r *ReplicationPolicyAPI) List() {
	page, size, err := r.GetPaginationParams()
	if err != nil {
		r.SendInternalServerError(err)
		return
	}
	// TODO: support more query
	query := &model.PolicyQuery{
		Name: r.GetString("name"),
		Pagination: common_model.Pagination{
			Page: page,
			Size: size,
		},
	}

	total, policies, err := replication.PolicyCtl.List(query)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to list policies: %v", err))
		return
	}
	for _, policy := range policies {
		if err = populateRegistries(replication.RegistryMgr, policy); err != nil {
			r.SendInternalServerError(fmt.Errorf("failed to populate registries for policy %d: %v", policy.ID, err))
			return
		}
	}
	r.SetPaginationHeader(total, query.Page, query.Size)
	r.WriteJSONData(policies)
}

// Create the replication policy
func (r *ReplicationPolicyAPI) Create() {
	policy := &model.Policy{}
	isValid, err := r.DecodeJSONReqAndValidate(policy)
	if !isValid {
		r.SendBadRequestError(err)
		return
	}

	if !r.validateName(policy) {
		return
	}
	if !r.validateRegistry(policy) {
		return
	}

	id, err := replication.PolicyCtl.Create(policy)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to create the policy: %v", err))
		return
	}
	r.Redirect(http.StatusCreated, strconv.FormatInt(id, 10))
}

// make sure the policy name doesn't exist
func (r *ReplicationPolicyAPI) validateName(policy *model.Policy) bool {
	p, err := replication.PolicyCtl.GetByName(policy.Name)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to get policy %s: %v", policy.Name, err))
		return false
	}
	if p != nil {
		r.SendConflictError(fmt.Errorf("policy %s already exists", policy.Name))
		return false
	}
	return true
}

// make sure the registry referred exists
func (r *ReplicationPolicyAPI) validateRegistry(policy *model.Policy) bool {
	var registryID int64
	if policy.SrcRegistry != nil && policy.SrcRegistry.ID > 0 {
		registryID = policy.SrcRegistry.ID
	} else {
		registryID = policy.DestRegistry.ID
	}
	registry, err := replication.RegistryMgr.Get(registryID)
	if err != nil {
		r.SendConflictError(fmt.Errorf("failed to get registry %d: %v", registryID, err))
		return false
	}
	if registry == nil {
		r.SendNotFoundError(fmt.Errorf("registry %d not found", registryID))
		return false
	}
	return true
}

// Get the specified replication policy
func (r *ReplicationPolicyAPI) Get() {
	id, err := r.GetInt64FromPath(":id")
	if id <= 0 || err != nil {
		r.SendBadRequestError(errors.New("invalid policy ID"))
		return
	}

	policy, err := replication.PolicyCtl.Get(id)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to get the policy %d: %v", id, err))
		return
	}
	if policy == nil {
		r.SendNotFoundError(fmt.Errorf("policy %d not found", id))
		return
	}
	if err = populateRegistries(replication.RegistryMgr, policy); err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to populate registries for policy %d: %v", policy.ID, err))
		return
	}

	r.WriteJSONData(policy)
}

// Update the replication policy
func (r *ReplicationPolicyAPI) Update() {
	id, err := r.GetInt64FromPath(":id")
	if id <= 0 || err != nil {
		r.SendBadRequestError(errors.New("invalid policy ID"))
		return
	}

	originalPolicy, err := replication.PolicyCtl.Get(id)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to get the policy %d: %v", id, err))
		return
	}
	if originalPolicy == nil {
		r.SendNotFoundError(fmt.Errorf("policy %d not found", id))
		return
	}

	policy := &model.Policy{}
	isValid, err := r.DecodeJSONReqAndValidate(policy)
	if !isValid {
		r.SendBadRequestError(err)
		return
	}

	if policy.Name != originalPolicy.Name &&
		!r.validateName(policy) {
		return
	}

	if !r.validateRegistry(policy) {
		return
	}

	policy.ID = id
	if err := replication.PolicyCtl.Update(policy); err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to update the policy %d: %v", id, err))
		return
	}
}

// Delete the replication policy
func (r *ReplicationPolicyAPI) Delete() {
	id, err := r.GetInt64FromPath(":id")
	if id <= 0 || err != nil {
		r.SendBadRequestError(errors.New("invalid policy ID"))
		return
	}

	policy, err := replication.PolicyCtl.Get(id)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to get the policy %d: %v", id, err))
		return
	}
	if policy == nil {
		r.SendNotFoundError(fmt.Errorf("policy %d not found", id))
		return
	}

	_, executions, err := replication.OperationCtl.ListExecutions(&models.ExecutionQuery{
		PolicyID: id,
	})
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to get the executions of policy %d: %v", id, err))
		return
	}

	for _, execution := range executions {
		if execution.Status == models.ExecutionStatusInProgress {
			r.SendInternalServerError(fmt.Errorf("the policy %d has running executions, can not be deleted", id))
			return
		}
	}

	if err := replication.PolicyCtl.Remove(id); err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to delete the policy %d: %v", id, err))
		return
	}
}

// ignore the credential for the registries
func populateRegistries(registryMgr registry.Manager, policy *model.Policy) error {
	if err := event.PopulateRegistries(registryMgr, policy); err != nil {
		return err
	}
	if policy.SrcRegistry != nil {
		hideAccessSecret(policy.SrcRegistry.Credential)
	}
	if policy.DestRegistry != nil {
		hideAccessSecret(policy.DestRegistry.Credential)
	}
	return nil
}
<commit_msg>Check the task status of execution whose status is running when deleting the policy<commit_after>// Copyright 2018 Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package api

import (
	"errors"
	"fmt"
	"net/http"
	"strconv"

	common_model "github.com/goharbor/harbor/src/common/models"
	"github.com/goharbor/harbor/src/replication"
	"github.com/goharbor/harbor/src/replication/dao/models"
	"github.com/goharbor/harbor/src/replication/event"
	"github.com/goharbor/harbor/src/replication/model"
	"github.com/goharbor/harbor/src/replication/registry"
)

// TODO rename the file to "replication.go"

// ReplicationPolicyAPI handles the replication policy requests
type ReplicationPolicyAPI struct {
	BaseController
}

// Prepare ...
func (r *ReplicationPolicyAPI) Prepare() {
	r.BaseController.Prepare()
	if !r.SecurityCtx.IsSysAdmin() {
		if !r.SecurityCtx.IsAuthenticated() {
			r.SendUnAuthorizedError(errors.New("UnAuthorized"))
			return
		}
		r.SendForbiddenError(errors.New(r.SecurityCtx.GetUsername()))
		return
	}
}

// List the replication policies
func (r *ReplicationPolicyAPI) List() {
	page, size, err := r.GetPaginationParams()
	if err != nil {
		r.SendInternalServerError(err)
		return
	}
	// TODO: support more query
	query := &model.PolicyQuery{
		Name: r.GetString("name"),
		Pagination: common_model.Pagination{
			Page: page,
			Size: size,
		},
	}

	total, policies, err := replication.PolicyCtl.List(query)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to list policies: %v", err))
		return
	}
	for _, policy := range policies {
		if err = populateRegistries(replication.RegistryMgr, policy); err != nil {
			r.SendInternalServerError(fmt.Errorf("failed to populate registries for policy %d: %v", policy.ID, err))
			return
		}
	}
	r.SetPaginationHeader(total, query.Page, query.Size)
	r.WriteJSONData(policies)
}

// Create the replication policy
func (r *ReplicationPolicyAPI) Create() {
	policy := &model.Policy{}
	isValid, err := r.DecodeJSONReqAndValidate(policy)
	if !isValid {
		r.SendBadRequestError(err)
		return
	}

	if !r.validateName(policy) {
		return
	}
	if !r.validateRegistry(policy) {
		return
	}

	id, err := replication.PolicyCtl.Create(policy)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to create the policy: %v", err))
		return
	}
	r.Redirect(http.StatusCreated, strconv.FormatInt(id, 10))
}

// make sure the policy name doesn't exist
func (r *ReplicationPolicyAPI) validateName(policy *model.Policy) bool {
	p, err := replication.PolicyCtl.GetByName(policy.Name)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to get policy %s: %v", policy.Name, err))
		return false
	}
	if p != nil {
		r.SendConflictError(fmt.Errorf("policy %s already exists", policy.Name))
		return false
	}
	return true
}

// make sure the registry referred exists
func (r *ReplicationPolicyAPI) validateRegistry(policy *model.Policy) bool {
	var registryID int64
	if policy.SrcRegistry != nil && policy.SrcRegistry.ID > 0 {
		registryID = policy.SrcRegistry.ID
	} else {
		registryID = policy.DestRegistry.ID
	}
	registry, err := replication.RegistryMgr.Get(registryID)
	if err != nil {
		r.SendConflictError(fmt.Errorf("failed to get registry %d: %v", registryID, err))
		return false
	}
	if registry == nil {
		r.SendNotFoundError(fmt.Errorf("registry %d not found", registryID))
		return false
	}
	return true
}

// Get the specified replication policy
func (r *ReplicationPolicyAPI) Get() {
	id, err := r.GetInt64FromPath(":id")
	if id <= 0 || err != nil {
		r.SendBadRequestError(errors.New("invalid policy ID"))
		return
	}

	policy, err := replication.PolicyCtl.Get(id)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to get the policy %d: %v", id, err))
		return
	}
	if policy == nil {
		r.SendNotFoundError(fmt.Errorf("policy %d not found", id))
		return
	}
	if err = populateRegistries(replication.RegistryMgr, policy); err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to populate registries for policy %d: %v", policy.ID, err))
		return
	}

	r.WriteJSONData(policy)
}

// Update the replication policy
func (r *ReplicationPolicyAPI) Update() {
	id, err := r.GetInt64FromPath(":id")
	if id <= 0 || err != nil {
		r.SendBadRequestError(errors.New("invalid policy ID"))
		return
	}

	originalPolicy, err := replication.PolicyCtl.Get(id)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to get the policy %d: %v", id, err))
		return
	}
	if originalPolicy == nil {
		r.SendNotFoundError(fmt.Errorf("policy %d not found", id))
		return
	}

	policy := &model.Policy{}
	isValid, err := r.DecodeJSONReqAndValidate(policy)
	if !isValid {
		r.SendBadRequestError(err)
		return
	}

	if policy.Name != originalPolicy.Name &&
		!r.validateName(policy) {
		return
	}

	if !r.validateRegistry(policy) {
		return
	}

	policy.ID = id
	if err := replication.PolicyCtl.Update(policy); err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to update the policy %d: %v", id, err))
		return
	}
}

// Delete the replication policy
func (r *ReplicationPolicyAPI) Delete() {
	id, err := r.GetInt64FromPath(":id")
	if id <= 0 || err != nil {
		r.SendBadRequestError(errors.New("invalid policy ID"))
		return
	}

	policy, err := replication.PolicyCtl.Get(id)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to get the policy %d: %v", id, err))
		return
	}
	if policy == nil {
		r.SendNotFoundError(fmt.Errorf("policy %d not found", id))
		return
	}

	isRunning, err := hasRunningExecutions(id)
	if err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to check the execution status of policy %d: %v", id, err))
		return
	}

	if isRunning {
		r.SendPreconditionFailedError(fmt.Errorf("the policy %d has running executions, can not be deleted", id))
		return
	}

	if err := replication.PolicyCtl.Remove(id); err != nil {
		r.SendInternalServerError(fmt.Errorf("failed to delete the policy %d: %v", id, err))
		return
	}
}

// the execution's status will not be updated if it is not queried
// so need to check the status of tasks to determine the status of
// the execution
func hasRunningExecutions(policyID int64) (bool, error) {
	_, executions, err := replication.OperationCtl.ListExecutions(&models.ExecutionQuery{
		PolicyID: policyID,
	})
	if err != nil {
		return false, err
	}
	for _, execution := range executions {
		if execution.Status != models.ExecutionStatusInProgress {
			continue
		}
		_, tasks, err := replication.OperationCtl.ListTasks(&models.TaskQuery{
			ExecutionID: execution.ID,
		})
		if err != nil {
			return false, err
		}
		for _, task := range tasks {
			if isTaskRunning(task) {
				return true, nil
			}
		}
	}
	return false, nil
}

// return true if the status of the task is running or pending
func isTaskRunning(task *models.Task) bool {
	if task == nil {
		return false
	}
	switch task.Status {
	case models.TaskStatusSucceed,
		models.TaskStatusStopped,
		models.TaskStatusFailed:
		return false
	}
	return true
}

// ignore the credential for the registries
func populateRegistries(registryMgr registry.Manager, policy *model.Policy) error {
	if err := event.PopulateRegistries(registryMgr, policy); err != nil {
		return err
	}
	if policy.SrcRegistry != nil {
		hideAccessSecret(policy.SrcRegistry.Credential)
	}
	if policy.DestRegistry != nil {
		hideAccessSecret(policy.DestRegistry.Credential)
	}
	return nil
}
<|endoftext|>


<commit_before>package brokers

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"

	"github.com/RichardKnop/machinery/Godeps/_workspace/src/github.com/streadway/amqp"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/signatures"
	"github.com/RichardKnop/machinery/v1/utils"
)

// AMQPBroker represents an AMQP broker
type AMQPBroker struct {
	config    *config.Config
	retryFunc func()
	stopChan  chan int
}

// NewAMQPBroker creates new AMQPBroker instance
func NewAMQPBroker(cnf *config.Config) Broker {
	return Broker(&AMQPBroker{
		config: cnf,
	})
}

// StartConsuming enters a loop and waits for incoming messages
func (amqpBroker *AMQPBroker) StartConsuming(consumerTag string, taskProcessor TaskProcessor) (bool, error) {
	if amqpBroker.retryFunc == nil {
		amqpBroker.retryFunc = utils.RetryClosure()
	}

	conn, channel, queue, _, err := amqpBroker.open()
	if err != nil {
		amqpBroker.retryFunc()
		return true, err // retry true
	}

	amqpBroker.retryFunc = utils.RetryClosure()

	defer amqpBroker.close(channel, conn)

	amqpBroker.stopChan = make(chan int)

	if err := channel.Qos(
		3,     // prefetch count
		0,     // prefetch size
		false, // global
	); err != nil {
		return false, fmt.Errorf("Channel Qos: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return false, fmt.Errorf("Queue Consume: %s", err)
	}

	log.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := amqpBroker.consume(deliveries, taskProcessor); err != nil {
		return true, err // retry true
	}

	return false, nil
}

// StopConsuming quits the loop
func (amqpBroker *AMQPBroker) StopConsuming() {
	// Notifying the stop channel stops consuming of messages
	amqpBroker.stopChan <- 1
}

// Publish places a new message on the default queue
func (amqpBroker *AMQPBroker) Publish(signature *signatures.TaskSignature) error {
	conn, channel, _, confirmsChan, err := amqpBroker.open()
	if err != nil {
		return err
	}

	defer amqpBroker.close(channel, conn)

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON Encode Message: %v", err)
	}

	signature.AdjustRoutingKey(
		amqpBroker.config.ExchangeType,
		amqpBroker.config.BindingKey,
		amqpBroker.config.DefaultQueue,
	)
	if err := channel.Publish(
		amqpBroker.config.Exchange, // exchange
		signature.RoutingKey,       // routing key
		false,                      // mandatory
		false,                      // immediate
		amqp.Publishing{
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// Consume a single message
func (amqpBroker *AMQPBroker) consumeOne(d amqp.Delivery, taskProcessor TaskProcessor, errorsChan chan error) {
	if len(d.Body) == 0 {
		d.Nack(false, false)                                   // multiple, requeue
		errorsChan <- errors.New("Received an empty message.") // RabbitMQ down?
		return
	}

	log.Printf("Received new message: %s", d.Body)

	signature := signatures.TaskSignature{}
	if err := json.Unmarshal(d.Body, &signature); err != nil {
		d.Nack(false, false) // multiple, requeue
		errorsChan <- err
		return
	}

	d.Ack(false) // multiple

	if err := taskProcessor.Process(&signature); err != nil {
		errorsChan <- err
	}
}

// Consumes messages...
func (amqpBroker *AMQPBroker) consume(deliveries <-chan amqp.Delivery, taskProcessor TaskProcessor) error {
	errorsChan := make(chan error)
	for {
		select {
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				amqpBroker.consumeOne(d, taskProcessor, errorsChan)
			}()
		case <-amqpBroker.stopChan:
			return nil
		}
	}
}

// Connects to the message queue, opens a channel, declares a queue
func (amqpBroker *AMQPBroker) open() (*amqp.Connection, *amqp.Channel, amqp.Queue, <-chan amqp.Confirmation, error) {
	var conn *amqp.Connection
	var channel *amqp.Channel
	var queue amqp.Queue
	var err error

	conn, err = amqp.Dial(amqpBroker.config.Broker)
	if err != nil {
		return conn, channel, queue, nil, fmt.Errorf("Dial: %s", err)
	}

	channel, err = conn.Channel()
	if err != nil {
		return conn, channel, queue, nil, fmt.Errorf("Channel: %s", err)
	}

	if err := channel.ExchangeDeclare(
		amqpBroker.config.Exchange,     // name of the exchange
		amqpBroker.config.ExchangeType, // type
		true,  // durable
		false, // delete when complete
		false, // internal
		false, // noWait
		nil,   // arguments
	); err != nil {
		return conn, channel, queue, nil, fmt.Errorf("Exchange: %s", err)
	}

	queue, err = channel.QueueDeclare(
		amqpBroker.config.DefaultQueue, // name
		true,  // durable
		false, // delete when unused
		false, // exclusive
		false, // no-wait
		nil,   // arguments
	)
	if err != nil {
		return conn, channel, queue, nil, fmt.Errorf("Queue Declare: %s", err)
	}

	if err := channel.QueueBind(
		queue.Name,                   // name of the queue
		amqpBroker.config.BindingKey, // binding key
		amqpBroker.config.Exchange,   // source exchange
		false, // noWait
		nil,   // arguments
	); err != nil {
		return conn, channel, queue, nil, fmt.Errorf("Queue Bind: %s", err)
	}

	confirmsChan := make(chan amqp.Confirmation, 1)

	// Enable publish confirmations
	if err := channel.Confirm(false); err != nil {
		close(confirmsChan)
		return conn, channel, queue, nil, fmt.Errorf("Channel could not be put into confirm mode: %s", err)
	}

	return conn, channel, queue, channel.NotifyPublish(confirmsChan), nil
}

// Closes the connection
func (amqpBroker *AMQPBroker) close(channel *amqp.Channel, conn *amqp.Connection) error {
	if err := channel.Close(); err != nil {
		return fmt.Errorf("Channel Close: %s", err)
	}

	if err := conn.Close(); err != nil {
		return fmt.Errorf("Connection Close: %s", err)
	}

	return nil
}
<commit_msg>AMPQ. Thread-safe connection reusing<commit_after>package brokers

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"sync"

	"github.com/RichardKnop/machinery/Godeps/_workspace/src/github.com/streadway/amqp"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/signatures"
	"github.com/RichardKnop/machinery/v1/utils"
)

var once sync.Once
var conn *amqp.Connection
var channel *amqp.Channel
var queue amqp.Queue

// AMQPBroker represents an AMQP broker
type AMQPBroker struct {
	config    *config.Config
	retryFunc func()
	stopChan  chan int
}

// NewAMQPBroker creates new AMQPBroker instance
func NewAMQPBroker(cnf *config.Config) Broker {
	return Broker(&AMQPBroker{
		config: cnf,
	})
}

// StartConsuming enters a loop and waits for incoming messages
func (amqpBroker *AMQPBroker) StartConsuming(consumerTag string, taskProcessor TaskProcessor) (bool, error) {
	if amqpBroker.retryFunc == nil {
		amqpBroker.retryFunc = utils.RetryClosure()
	}

	_, channel, queue, _, err := amqpBroker.open()
	if err != nil {
		amqpBroker.retryFunc()
		return true, err // retry true
	}

	amqpBroker.retryFunc = utils.RetryClosure()

	amqpBroker.stopChan = make(chan int)

	if err := channel.Qos(
		3,     // prefetch count
		0,     // prefetch size
		false, // global
	); err != nil {
		return false, fmt.Errorf("Channel Qos: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return false, fmt.Errorf("Queue Consume: %s", err)
	}

	log.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := amqpBroker.consume(deliveries, taskProcessor); err != nil {
		return true, err // retry true
	}

	return false, nil
}

// StopConsuming quits the loop
func (amqpBroker *AMQPBroker) StopConsuming() {
	// Notifying the stop channel stops consuming of messages
	amqpBroker.stopChan <- 1
}

// Publish places a new message on the default queue
func (amqpBroker *AMQPBroker) Publish(signature *signatures.TaskSignature) error {
	_, channel, _, confirmsChan, err := amqpBroker.open()
	if err != nil {
		return err
	}

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON Encode Message: %v", err)
	}

	signature.AdjustRoutingKey(
		amqpBroker.config.ExchangeType,
		amqpBroker.config.BindingKey,
		amqpBroker.config.DefaultQueue,
	)
	if err := channel.Publish(
		amqpBroker.config.Exchange, // exchange
		signature.RoutingKey,       // routing key
		false,                      // mandatory
		false,                      // immediate
		amqp.Publishing{
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// Consume a single message
func (amqpBroker *AMQPBroker) consumeOne(d amqp.Delivery, taskProcessor TaskProcessor, errorsChan chan error) {
	if len(d.Body) == 0 {
		d.Nack(false, false)                                   // multiple, requeue
		errorsChan <- errors.New("Received an empty message.") // RabbitMQ down?
		return
	}

	log.Printf("Received new message: %s", d.Body)

	signature := signatures.TaskSignature{}
	if err := json.Unmarshal(d.Body, &signature); err != nil {
		d.Nack(false, false) // multiple, requeue
		errorsChan <- err
		return
	}

	d.Ack(false) // multiple

	if err := taskProcessor.Process(&signature); err != nil {
		errorsChan <- err
	}
}

// Consumes messages...
func (amqpBroker *AMQPBroker) consume(deliveries <-chan amqp.Delivery, taskProcessor TaskProcessor) error {
	errorsChan := make(chan error)
	for {
		select {
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				amqpBroker.consumeOne(d, taskProcessor, errorsChan)
			}()
		case <-amqpBroker.stopChan:
			return nil
		}
	}
}

// Connects to the message queue, opens a channel, declares a queue
func (amqpBroker *AMQPBroker) connect() {

	var err error
	fmt.Println("connecting...")
	conn, err = amqp.Dial(amqpBroker.config.Broker)
	if err != nil {
		fmt.Printf("Dial: %s\n", err)
	}

	channel, err = conn.Channel()
	if err != nil {
		fmt.Printf("Channel: %s\n", err)
	}

	if err := channel.ExchangeDeclare(
		amqpBroker.config.Exchange,     // name of the exchange
		amqpBroker.config.ExchangeType, // type
		true,  // durable
		false, // delete when complete
		false, // internal
		false, // noWait
		nil,   // arguments
	); err != nil {
		fmt.Printf("Exchange: %s\n", err)
	}

	queue, err = channel.QueueDeclare(
		amqpBroker.config.DefaultQueue, // name
		true,  // durable
		false, // delete when unused
		false, // exclusive
		false, // no-wait
		nil,   // arguments
	)
	if err != nil {
		fmt.Printf("Queue Declare: %s\n", err)
	}

	if err := channel.QueueBind(
		queue.Name,                   // name of the queue
		amqpBroker.config.BindingKey, // binding key
		amqpBroker.config.Exchange,   // source exchange
		false, // noWait
		nil,   // arguments
	); err != nil {
		fmt.Printf("Queue Bind: %s\n", err)
	}

}

// Connects to the message queue, opens a channel, declares a queue
func (amqpBroker *AMQPBroker) open() (*amqp.Connection, *amqp.Channel, amqp.Queue, <-chan amqp.Confirmation, error) {
	if conn == nil {
		connected := make(chan bool)
		go func() {
			once.Do(amqpBroker.connect)
			connected <- true
		}()
		<-connected
	}

	if conn == nil {
		return conn, channel, queue, nil, fmt.Errorf("Can't connect to the server")
	}

	confirmsChan := make(chan amqp.Confirmation, 1)
	// Enable publish confirmations
	if err := channel.Confirm(false); err != nil {
		close(confirmsChan)
		fmt.Printf("Channel could not be put into confirm mode: %s\n", err)
	}

	return conn, channel, queue, channel.NotifyPublish(confirmsChan), nil
}
<|endoftext|>


<commit_before>package db_test

import (
	"cred-alert/mysqlrunner"
	"fmt"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

func TestDB(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "DB Suite")
}

var dbRunner mysqlrunner.Runner

var _ = BeforeSuite(func() {
	dbRunner = mysqlrunner.Runner{
		DBName: fmt.Sprintf("testdb_%d", GinkgoParallelNode()),
	}
	dbRunner.Setup()
})

var _ = AfterSuite(func() {
	dbRunner.Teardown()
})
<commit_msg>Truncate test db in AfterEach<commit_after>package db_test

import (
	"cred-alert/mysqlrunner"
	"fmt"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

func TestDB(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "DB Suite")
}

var dbRunner mysqlrunner.Runner

var _ = BeforeSuite(func() {
	dbRunner = mysqlrunner.Runner{
		DBName: fmt.Sprintf("testdb_%d", GinkgoParallelNode()),
	}
	dbRunner.Setup()
})

var _ = AfterSuite(func() {
	dbRunner.Teardown()
})

var _ = AfterEach(func() {
	dbRunner.Truncate()
})
<|endoftext|>


<commit_before>// Package proto contains protocol buffers that are exchanged between the client
// and server. Note, some of these types have additional helper methods.
package proto

import (
	"bytes"
	"crypto"
	"fmt"

	"github.com/google/go-tpm/tpm2"
	"github.com/google/go-tpm/tpmutil"
)

// CheckIfSubsetOf verifies if the pcrs PCRs are a valid "subset" of the provided
// "superset" of PCRs. The PCR values must match (if present), and all PCRs must
// be present in the superset. This function will return an error containing the
// first missing or mismatched PCR number.
func (x *Pcrs) CheckIfSubsetOf(superset *Pcrs) error {
	if x.GetHash() != superset.GetHash() {
		return fmt.Errorf("PCR hash algo not matching: %v, %v", x.GetHash(), superset.GetHash())
	}
	for pcrNum, pcrVal := range x.GetPcrs() {
		if expectedVal, ok := superset.GetPcrs()[pcrNum]; ok {
			if !bytes.Equal(expectedVal, pcrVal) {
				return fmt.Errorf("PCR %d mismatch: expected %v, got %v", pcrNum, expectedVal, pcrVal)
			}
		} else {
			return fmt.Errorf("PCR %d mismatch: value missing from the superset PCRs", pcrNum)
		}
	}
	return nil
}

// PCRSelection returns the corresponding tpm2.PCRSelection for the PCR data.
func (x *Pcrs) PCRSelection() tpm2.PCRSelection {
	sel := tpm2.PCRSelection{Hash: tpm2.Algorithm(x.GetHash())}

	for pcrNum := range x.GetPcrs() {
		sel.PCRs = append(sel.PCRs, int(pcrNum))
	}
	return sel
}

// HasSamePCRSelection checks if the Pcrs has the same PCRSelection as the
// provided given tpm2.PCRSelection (including the hash algorithm).
func (x *Pcrs) HasSamePCRSelection(pcrSel tpm2.PCRSelection) bool {
	if tpm2.Algorithm(x.Hash) != pcrSel.Hash {
		return false
	}
	if len(x.GetPcrs()) != len(pcrSel.PCRs) {
		return false
	}
	for _, p := range pcrSel.PCRs {
		if _, ok := x.Pcrs[uint32(p)]; !ok {
			return false
		}
	}
	return true
}

// ComputePCRSessionAuth calculates the authorization value for the given PCRs.
func (x *Pcrs) ComputePCRSessionAuth(hashAlg crypto.Hash) []byte {
	// Start with all zeros, we only use a single policy command on our session.
	oldDigest := make([]byte, hashAlg.Size())
	ccPolicyPCR, _ := tpmutil.Pack(tpm2.CmdPolicyPCR)

	// Extend the policy digest, see TPM2_PolicyPCR in Part 3 of the spec.
	hash := hashAlg.New()
	hash.Write(oldDigest)
	hash.Write(ccPolicyPCR)
	hash.Write(encodePCRSelection(x.PCRSelection()))
	hash.Write(x.ComputePCRDigest(hashAlg))
	newDigest := hash.Sum(nil)
	return newDigest[:]
}

// ComputePCRDigest computes the digest of the Pcrs. Note that the digest hash
// algorithm may differ from the PCRs' hash (which denotes the PCR bank).
func (x *Pcrs) ComputePCRDigest(hashAlg crypto.Hash) []byte {
	hash := hashAlg.New()
	for i := 0; i < 24; i++ {
		if pcrValue, exists := x.Pcrs[uint32(i)]; exists {
			hash.Write(pcrValue)
		}
	}
	return hash.Sum(nil)
}

// Encode a tpm2.PCRSelection as if it were a TPML_PCR_SELECTION
func encodePCRSelection(sel tpm2.PCRSelection) []byte {
	// Encode count, pcrSelections.hash and pcrSelections.sizeofSelect fields
	buf, _ := tpmutil.Pack(uint32(1), sel.Hash, byte(3))
	// Encode pcrSelect bitmask
	pcrBits := make([]byte, 3)
	for _, pcr := range sel.PCRs {
		byteNum := pcr / 8
		bytePos := 1 << uint(pcr%8)
		pcrBits[byteNum] |= byte(bytePos)
	}

	return append(buf, pcrBits...)
}
<commit_msg>proto: Add Pcrs.PrettyFormat<commit_after>// Package proto contains protocol buffers that are exchanged between the client
// and server. Note, some of these types have additional helper methods.
package proto

import (
	"bytes"
	"crypto"
	"fmt"
	"io"

	"github.com/google/go-tpm/tpm2"
	"github.com/google/go-tpm/tpmutil"
)

const minPCRIndex = uint32(0)

func (x *Pcrs) maxPCRIndex() uint32 {
	max := minPCRIndex
	for idx := range x.Pcrs {
		if idx > max {
			max = idx
		}
	}
	return max
}

// PrettyFormat writes a multiline representation of the PCR values to w.
func (x *Pcrs) PrettyFormat(w io.Writer) error {
	if _, err := fmt.Fprintf(w, "%v:\n", x.Hash); err != nil {
		return err
	}
	for idx := minPCRIndex; idx <= x.maxPCRIndex(); idx++ {
		if val, ok := x.Pcrs[idx]; ok {
			if _, err := fmt.Fprintf(w, "  %2d: 0x%X\n", idx, val); err != nil {
				return err
			}
		}
	}
	return nil
}

// CheckIfSubsetOf verifies if the pcrs PCRs are a valid "subset" of the provided
// "superset" of PCRs. The PCR values must match (if present), and all PCRs must
// be present in the superset. This function will return an error containing the
// first missing or mismatched PCR number.
func (x *Pcrs) CheckIfSubsetOf(superset *Pcrs) error {
	if x.GetHash() != superset.GetHash() {
		return fmt.Errorf("PCR hash algo not matching: %v, %v", x.GetHash(), superset.GetHash())
	}
	for pcrNum, pcrVal := range x.GetPcrs() {
		if expectedVal, ok := superset.GetPcrs()[pcrNum]; ok {
			if !bytes.Equal(expectedVal, pcrVal) {
				return fmt.Errorf("PCR %d mismatch: expected %v, got %v", pcrNum, expectedVal, pcrVal)
			}
		} else {
			return fmt.Errorf("PCR %d mismatch: value missing from the superset PCRs", pcrNum)
		}
	}
	return nil
}

// PCRSelection returns the corresponding tpm2.PCRSelection for the PCR data.
func (x *Pcrs) PCRSelection() tpm2.PCRSelection {
	sel := tpm2.PCRSelection{Hash: tpm2.Algorithm(x.GetHash())}

	for pcrNum := range x.GetPcrs() {
		sel.PCRs = append(sel.PCRs, int(pcrNum))
	}
	return sel
}

// HasSamePCRSelection checks if the Pcrs has the same PCRSelection as the
// provided given tpm2.PCRSelection (including the hash algorithm).
func (x *Pcrs) HasSamePCRSelection(pcrSel tpm2.PCRSelection) bool {
	if tpm2.Algorithm(x.Hash) != pcrSel.Hash {
		return false
	}
	if len(x.GetPcrs()) != len(pcrSel.PCRs) {
		return false
	}
	for _, p := range pcrSel.PCRs {
		if _, ok := x.Pcrs[uint32(p)]; !ok {
			return false
		}
	}
	return true
}

// ComputePCRSessionAuth calculates the authorization value for the given PCRs.
func (x *Pcrs) ComputePCRSessionAuth(hashAlg crypto.Hash) []byte {
	// Start with all zeros, we only use a single policy command on our session.
	oldDigest := make([]byte, hashAlg.Size())
	ccPolicyPCR, _ := tpmutil.Pack(tpm2.CmdPolicyPCR)

	// Extend the policy digest, see TPM2_PolicyPCR in Part 3 of the spec.
	hash := hashAlg.New()
	hash.Write(oldDigest)
	hash.Write(ccPolicyPCR)
	hash.Write(encodePCRSelection(x.PCRSelection()))
	hash.Write(x.ComputePCRDigest(hashAlg))
	newDigest := hash.Sum(nil)
	return newDigest[:]
}

// ComputePCRDigest computes the digest of the Pcrs. Note that the digest hash
// algorithm may differ from the PCRs' hash (which denotes the PCR bank).
func (x *Pcrs) ComputePCRDigest(hashAlg crypto.Hash) []byte {
	hash := hashAlg.New()
	for i := 0; i < 24; i++ {
		if pcrValue, exists := x.Pcrs[uint32(i)]; exists {
			hash.Write(pcrValue)
		}
	}
	return hash.Sum(nil)
}

// Encode a tpm2.PCRSelection as if it were a TPML_PCR_SELECTION
func encodePCRSelection(sel tpm2.PCRSelection) []byte {
	// Encode count, pcrSelections.hash and pcrSelections.sizeofSelect fields
	buf, _ := tpmutil.Pack(uint32(1), sel.Hash, byte(3))
	// Encode pcrSelect bitmask
	pcrBits := make([]byte, 3)
	for _, pcr := range sel.PCRs {
		byteNum := pcr / 8
		bytePos := 1 << uint(pcr%8)
		pcrBits[byteNum] |= byte(bytePos)
	}

	return append(buf, pcrBits...)
}
<|endoftext|>


<commit_before>package database

import (
	"context"
	"fmt"
	"math/rand"
	"time"

	"cloud.google.com/go/spanner"
	database "cloud.google.com/go/spanner/admin/database/apiv1"
	"github.com/rs/zerolog/log"
	"github.com/spf13/viper"
	adminpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1"
	"google.golang.org/grpc/codes"
)

type Spanner struct {
	Database
	client *spanner.Client
	admin  *database.DatabaseAdminClient
}

// SetupSpanner initializes the spanner clients.
func NewSpanner(project, instance, db string) (*Spanner, error) {
	ctx := context.Background()
	sp := &Spanner{}

	adminClient, err := database.NewDatabaseAdminClient(ctx)
	if err != nil {
		return nil, err
	}

	sp.admin = adminClient

	// Create the databases if they don't exist.
	err = sp.createSpannerDatabase(ctx, project, instance, db)

	if err != nil {
		return nil, err
	}

	dbstr := fmt.Sprintf("projects/%s/instances/%s/databases/%s",
		project,
		instance,
		db)

	client, err := spanner.NewClient(context.Background(), dbstr)
	if err != nil {
		return nil, err
	}

	sp.client = client
	return sp, err
}

func (s *Spanner) Read(msg *Message) error {

	stmt := spanner.NewStatement(`
			SELECT SUM(a.Count) as Count
			FROM sheep as a
			WHERE a.Keyspace=@Keyspace
			AND a.Key=@Key
			AND a.Name=@Name
		`)
	stmt.Params["Keyspace"] = msg.Keyspace
	stmt.Params["Key"] = msg.Key
	stmt.Params["Name"] = msg.Name

	// TODO: Expose this stale time.
	iter := s.client.
		Single().
		WithTimestampBound(spanner.MaxStaleness(5*time.Second)).
		Query(context.Background(), stmt)
	defer iter.Stop()

	row, err := iter.Next()

	if err != nil {
		return err
	}

	var value spanner.NullInt64
	err = row.ColumnByName("Count", &value)

	if err != nil {
		return err
	}

	if value.Valid {
		msg.Value = value.Int64
		return nil
	}

	return &spanner.Error{
		Code: codes.NotFound,
		Desc: "counter not found",
	}
}

func (s *Spanner) Save(message *Message) error {
	ctx := context.WithValue(context.Background(), contextKey("message"), message)
	_, err := s.client.ReadWriteTransaction(ctx, s.doSave)
	return err
}

// Here's where the magic happens. Save out message!
func (s *Spanner) doSave(ctx context.Context, rw *spanner.ReadWriteTransaction) error {
	msg := ctx.Value(contextKey("message")).(*Message)
	shards := viper.GetInt("spanner.shards")
	shard := rand.Intn(shards)

	// First, let's check and see if our message has been written.
	row, err := rw.ReadRow(context.Background(), "sheep_transaction", spanner.Key{msg.Keyspace, msg.Key, msg.Name, msg.UUID}, []string{"Applied"})
	if err != nil {
		if spanner.ErrCode(err) != codes.NotFound {
			return err
		}
	} else {
		var ap bool
		err = row.ColumnByName("Applied", &ap)
		if err != nil {
			return err
		}
		if ap {
			return nil
		}
	}

	// Let's get our current count
	var move int64
	row, err = rw.ReadRow(context.Background(), "sheep", spanner.Key{msg.Keyspace, msg.Key, msg.Name, shard}, []string{"Count"})
	if err != nil {
		if spanner.ErrCode(err) != codes.NotFound {
			return err
		}
	} else {
		row.ColumnByName("Count", &move)
	}

	// Now we'll do our operation.
	switch msg.Operation {
	case "incr":
		move++
	case "decr":
		move--
	case "set":
		move = msg.Value
	default:
		return &spanner.Error{
			Code: codes.InvalidArgument,
			Desc: "Invalid operation sent from message '" + msg.Operation + "', aborting transaction!",
		}
	}

	m := []*spanner.Mutation{}

	log.Debug().Int("shard", shard).Msg("shard selected for op")
	if msg.Operation == "set" {
		for i := 0; i < shards; i++ {
			m = append(m, spanner.InsertOrUpdate(
				"sheep",
				[]string{"Keyspace", "Key", "Name", "Shard", "Count"},
				[]interface{}{msg.Keyspace, msg.Key, msg.Name, i, move},
			))
		}
	} else {
		m = append(m, spanner.InsertOrUpdate(
			"sheep",
			[]string{"Keyspace", "Key", "Name", "Shard", "Count"},
			[]interface{}{msg.Keyspace, msg.Key, msg.Name, shard, move},
		))
	}

	m = append(m, spanner.InsertOrUpdate(
		"sheep_transaction",
		[]string{"Keyspace", "Key", "Name", "Shard", "UUID", "Time"},
		[]interface{}{msg.Keyspace, msg.Key, msg.Name, shard, msg.UUID, time.Now()}))

	// ...and write!
	return rw.BufferWrite(m)

}

func (s *Spanner) createSpannerDatabase(ctx context.Context, project, instance, db string) error {
	// Create our database if it doesn't exist.
	_, err := s.admin.GetDatabase(ctx, &adminpb.GetDatabaseRequest{
		Name: "projects/" + project + "/instances/" + instance + "/databases/" + db})
	if err != nil {
		// Database doesn't exist, or error.
		op, err := s.admin.CreateDatabase(ctx, &adminpb.CreateDatabaseRequest{
			Parent:          "projects/" + project + "/instances/" + instance,
			CreateStatement: "CREATE DATABASE `" + db + "`",
			ExtraStatements: []string{
				`CREATE TABLE sheep (
							Keyspace 	STRING(MAX) NOT NULL,
							Key 			STRING(MAX) NOT NULL,
							Name			STRING(MAX) NOT NULL,
							Shard     INT64       NOT NULL,
							Count 		INT64       NOT NULL
					) PRIMARY KEY (Keyspace, Key, Name, Shard)`,
				`CREATE TABLE sheep_transaction (
							Keyspace 	STRING(MAX) NOT NULL,
							Key 			STRING(MAX) NOT NULL,
							Name			STRING(MAX) NOT NULL,
							Shard     INT64       NOT NULL,
							UUID 			STRING(128) NOT NULL,
							Time      TIMESTAMP   NOT NULL
					) PRIMARY KEY (Keyspace, Key, Name, Shard, UUID),
						INTERLEAVE IN PARENT sheep ON DELETE CASCADE`,
			},
		})

		if err != nil {
			return err
		}

		_, err = op.Wait(ctx)

		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>de-interleave transaction table<commit_after>package database

import (
	"context"
	"fmt"
	"math/rand"
	"time"

	"cloud.google.com/go/spanner"
	database "cloud.google.com/go/spanner/admin/database/apiv1"
	"github.com/rs/zerolog/log"
	"github.com/spf13/viper"
	adminpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1"
	"google.golang.org/grpc/codes"
)

type Spanner struct {
	Database
	client *spanner.Client
	admin  *database.DatabaseAdminClient
}

// SetupSpanner initializes the spanner clients.
func NewSpanner(project, instance, db string) (*Spanner, error) {
	ctx := context.Background()
	sp := &Spanner{}

	adminClient, err := database.NewDatabaseAdminClient(ctx)
	if err != nil {
		return nil, err
	}

	sp.admin = adminClient

	// Create the databases if they don't exist.
	err = sp.createSpannerDatabase(ctx, project, instance, db)

	if err != nil {
		return nil, err
	}

	dbstr := fmt.Sprintf("projects/%s/instances/%s/databases/%s",
		project,
		instance,
		db)

	client, err := spanner.NewClient(context.Background(), dbstr)
	if err != nil {
		return nil, err
	}

	sp.client = client
	return sp, err
}

func (s *Spanner) Read(msg *Message) error {

	stmt := spanner.NewStatement(`
			SELECT SUM(a.Count) as Count
			FROM sheep as a
			WHERE a.Keyspace=@Keyspace
			AND a.Key=@Key
			AND a.Name=@Name
		`)
	stmt.Params["Keyspace"] = msg.Keyspace
	stmt.Params["Key"] = msg.Key
	stmt.Params["Name"] = msg.Name

	// TODO: Expose this stale time.
	iter := s.client.
		Single().
		WithTimestampBound(spanner.MaxStaleness(5*time.Second)).
		Query(context.Background(), stmt)
	defer iter.Stop()

	row, err := iter.Next()

	if err != nil {
		return err
	}

	var value spanner.NullInt64
	err = row.ColumnByName("Count", &value)

	if err != nil {
		return err
	}

	if value.Valid {
		msg.Value = value.Int64
		return nil
	}

	return &spanner.Error{
		Code: codes.NotFound,
		Desc: "counter not found",
	}
}

func (s *Spanner) Save(message *Message) error {
	ctx := context.WithValue(context.Background(), contextKey("message"), message)
	_, err := s.client.ReadWriteTransaction(ctx, s.doSave)
	return err
}

// Here's where the magic happens. Save out message!
func (s *Spanner) doSave(ctx context.Context, rw *spanner.ReadWriteTransaction) error {
	msg := ctx.Value(contextKey("message")).(*Message)
	shards := viper.GetInt("spanner.shards")
	shard := rand.Intn(shards)

	// First, let's check and see if our message has been written.
	row, err := rw.ReadRow(context.Background(), "sheep_transaction", spanner.Key{msg.Keyspace, msg.Key, msg.Name, msg.UUID}, []string{"UUID"})
	if err != nil {
		if spanner.ErrCode(err) != codes.NotFound {
			return err
		}
	} else if err == nil {
		// We need to return if err is nil, this means
		// the UUID was found.
		return nil
	}

	// Let's get our current count
	var move int64
	row, err = rw.ReadRow(context.Background(), "sheep", spanner.Key{msg.Keyspace, msg.Key, msg.Name, shard}, []string{"Count"})
	if err != nil {
		if spanner.ErrCode(err) != codes.NotFound {
			return err
		}
	} else {
		row.ColumnByName("Count", &move)
	}

	// Now we'll do our operation.
	switch msg.Operation {
	case "incr":
		move++
	case "decr":
		move--
	case "set":
		move = msg.Value
	default:
		return &spanner.Error{
			Code: codes.InvalidArgument,
			Desc: "Invalid operation sent from message '" + msg.Operation + "', aborting transaction!",
		}
	}

	m := []*spanner.Mutation{}

	log.Debug().Int("shard", shard).Msg("shard selected for op")
	if msg.Operation == "set" {
		for i := 0; i < shards; i++ {
			m = append(m, spanner.InsertOrUpdate(
				"sheep",
				[]string{"Keyspace", "Key", "Name", "Shard", "Count"},
				[]interface{}{msg.Keyspace, msg.Key, msg.Name, i, move},
			))
		}
	} else {
		m = append(m, spanner.InsertOrUpdate(
			"sheep",
			[]string{"Keyspace", "Key", "Name", "Shard", "Count"},
			[]interface{}{msg.Keyspace, msg.Key, msg.Name, shard, move},
		))
	}

	m = append(m, spanner.InsertOrUpdate(
		"sheep_transaction",
		[]string{"Keyspace", "Key", "Name", "UUID", "Time"},
		[]interface{}{msg.Keyspace, msg.Key, msg.Name, msg.UUID, time.Now()}))

	// ...and write!
	return rw.BufferWrite(m)

}

func (s *Spanner) createSpannerDatabase(ctx context.Context, project, instance, db string) error {
	// Create our database if it doesn't exist.
	_, err := s.admin.GetDatabase(ctx, &adminpb.GetDatabaseRequest{
		Name: "projects/" + project + "/instances/" + instance + "/databases/" + db})
	if err != nil {
		// Database doesn't exist, or error.
		op, err := s.admin.CreateDatabase(ctx, &adminpb.CreateDatabaseRequest{
			Parent:          "projects/" + project + "/instances/" + instance,
			CreateStatement: "CREATE DATABASE `" + db + "`",
			ExtraStatements: []string{
				`CREATE TABLE sheep (
							Keyspace 	STRING(MAX) NOT NULL,
							Key 			STRING(MAX) NOT NULL,
							Name			STRING(MAX) NOT NULL,
							Shard     INT64       NOT NULL,
							Count 		INT64       NOT NULL
					) PRIMARY KEY (Keyspace, Key, Name, Shard)`,
				`CREATE TABLE sheep_transaction (
							Keyspace 	STRING(MAX) NOT NULL,
							Key 			STRING(MAX) NOT NULL,
							Name			STRING(MAX) NOT NULL,
							UUID 			STRING(128) NOT NULL,
							Time      TIMESTAMP   NOT NULL
					) PRIMARY KEY (Keyspace, Key, Name, UUID)`,
			},
		})

		if err != nil {
			return err
		}

		_, err = op.Wait(ctx)

		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC 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 partial

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/google/go-containerregistry/v1"
	"github.com/google/go-containerregistry/v1/v1util"
)

// WithRawConfigFile defines the subset of v1.Image used by these helper methods
type WithRawConfigFile interface {
	// RawConfigFile returns the serialized bytes of this image's config file.
	RawConfigFile() ([]byte, error)
}

// ConfigFile is a helper for implementing v1.Image
func ConfigFile(i WithRawConfigFile) (*v1.ConfigFile, error) {
	b, err := i.RawConfigFile()
	if err != nil {
		return nil, err
	}
	return v1.ParseConfigFile(v1util.NopReadCloser(bytes.NewBuffer(b)))
}

// ConfigName is a helper for implementing v1.Image
func ConfigName(i WithRawConfigFile) (v1.Hash, error) {
	b, err := i.RawConfigFile()
	if err != nil {
		return v1.Hash{}, err
	}
	h, _, err := v1.SHA256(v1util.NopReadCloser(bytes.NewBuffer(b)))
	return h, err
}

// WithConfigFile defines the subset of v1.Image used by these helper methods
type WithConfigFile interface {
	// ConfigFile returns this image's config file.
	ConfigFile() (*v1.ConfigFile, error)
}

// DiffIDs is a helper for implementing v1.Image
func DiffIDs(i WithConfigFile) ([]v1.Hash, error) {
	cfg, err := i.ConfigFile()
	if err != nil {
		return nil, err
	}
	dids := make([]v1.Hash, len(cfg.RootFS.DiffIDs))
	for i, did := range cfg.RootFS.DiffIDs {
		dids[len(dids)-i-1] = did
	}
	return dids, nil
}

// RawConfigFile is a helper for implementing v1.Image
func RawConfigFile(i WithConfigFile) ([]byte, error) {
	cfg, err := i.ConfigFile()
	if err != nil {
		return nil, err
	}
	return json.Marshal(cfg)
}

// WithUncompressedLayer defines the subset of v1.Image used by these helper methods
type WithUncompressedLayer interface {
	// UncompressedLayer is like UncompressedBlob, but takes the "diff id".
	UncompressedLayer(v1.Hash) (io.ReadCloser, error)
}

// Layer is the same as Blob, but takes the "diff id".
func Layer(wul WithUncompressedLayer, h v1.Hash) (io.ReadCloser, error) {
	rc, err := wul.UncompressedLayer(h)
	if err != nil {
		return nil, err
	}
	return v1util.GzipReadCloser(rc)
}

// WithRawManifest defines the subset of v1.Image used by these helper methods
type WithRawManifest interface {
	// RawManifest returns the serialized bytes of this image's config file.
	RawManifest() ([]byte, error)
}

// Digest is a helper for implementing v1.Image
func Digest(i WithRawManifest) (v1.Hash, error) {
	mb, err := i.RawManifest()
	if err != nil {
		return v1.Hash{}, nil
	}
	digest, _, err := v1.SHA256(ioutil.NopCloser(bytes.NewReader(mb)))
	return digest, err
}

// Manifest is a helper for implementing v1.Image
func Manifest(i WithRawManifest) (*v1.Manifest, error) {
	b, err := i.RawManifest()
	if err != nil {
		return nil, err
	}
	return v1.ParseManifest(v1util.NopReadCloser(bytes.NewBuffer(b)))
}

// WithManifest defines the subset of v1.Image used by these helper methods
type WithManifest interface {
	// Manifest returns this image's Manifest object.
	Manifest() (*v1.Manifest, error)
}

// RawManifest is a helper for implementing v1.Image
func RawManifest(i WithManifest) ([]byte, error) {
	m, err := i.Manifest()
	if err != nil {
		return nil, err
	}
	return json.Marshal(m)
}

// FSLayers is a helper for implementing v1.Image
func FSLayers(i WithManifest) ([]v1.Hash, error) {
	m, err := i.Manifest()
	if err != nil {
		return nil, err
	}
	fsl := make([]v1.Hash, len(m.Layers))
	for i, l := range m.Layers {
		fsl[len(fsl)-i-1] = l.Digest
	}
	return fsl, nil
}

// BlobSet is a helper for implementing v1.Image
func BlobSet(i WithManifest) (map[v1.Hash]struct{}, error) {
	m, err := i.Manifest()
	if err != nil {
		return nil, err
	}
	bs := make(map[v1.Hash]struct{})
	for _, l := range m.Layers {
		bs[l.Digest] = struct{}{}
	}
	bs[m.Config.Digest] = struct{}{}
	return bs, nil
}

// BlobSize is a helper for implementing v1.Image
func BlobSize(i WithManifest, h v1.Hash) (int64, error) {
	m, err := i.Manifest()
	if err != nil {
		return -1, err
	}
	for _, l := range m.Layers {
		if l.Digest == h {
			return l.Size, nil
		}
	}
	return -1, fmt.Errorf("blob %v not found", h)
}

// WithManifestAndConfigFile defines the subset of v1.Image used by these helper methods
type WithManifestAndConfigFile interface {
	WithConfigFile

	// Manifest returns this image's Manifest object.
	Manifest() (*v1.Manifest, error)
}

// BlobToDiffID is a helper for mapping between compressed
// and uncompressed blob hashes.
func BlobToDiffID(i WithManifestAndConfigFile, h v1.Hash) (v1.Hash, error) {
	blobs, err := FSLayers(i)
	if err != nil {
		return v1.Hash{}, err
	}
	diffIDs, err := DiffIDs(i)
	if err != nil {
		return v1.Hash{}, err
	}
	if len(blobs) != len(diffIDs) {
		return v1.Hash{}, fmt.Errorf("mismatched fs layers (%d) and diff ids (%d)", len(blobs), len(diffIDs))
	}
	for i, blob := range blobs {
		if blob == h {
			return diffIDs[i], nil
		}
	}
	return v1.Hash{}, fmt.Errorf("unknown blob %v", h)
}

// DiffIDToBlob is a helper for mapping between uncompressed
// and compressed blob hashes.
func DiffIDToBlob(wm WithManifestAndConfigFile, h v1.Hash) (v1.Hash, error) {
	blobs, err := FSLayers(wm)
	if err != nil {
		return v1.Hash{}, err
	}
	diffIDs, err := DiffIDs(wm)
	if err != nil {
		return v1.Hash{}, err
	}
	if len(blobs) != len(diffIDs) {
		return v1.Hash{}, fmt.Errorf("mismatched fs layers (%d) and diff ids (%d)", len(blobs), len(diffIDs))
	}
	for i, diffID := range diffIDs {
		if diffID == h {
			return blobs[i], nil
		}
	}
	return v1.Hash{}, fmt.Errorf("unknown diffID %v", h)

}

// WithBlob defines the subset of v1.Image used by these helper methods
type WithBlob interface {
	// Blob returns a ReadCloser for streaming the blob's content.
	Blob(v1.Hash) (io.ReadCloser, error)
}

// UncompressedBlob returns a ReadCloser for streaming the blob's content uncompressed.
func UncompressedBlob(b WithBlob, h v1.Hash) (io.ReadCloser, error) {
	rc, err := b.Blob(h)
	if err != nil {
		return nil, err
	}
	return v1util.GunzipReadCloser(rc)
}
<commit_msg>Fix v1/partial Digest swallowing err (#49)<commit_after>// Copyright 2018 Google LLC 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 partial

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/google/go-containerregistry/v1"
	"github.com/google/go-containerregistry/v1/v1util"
)

// WithRawConfigFile defines the subset of v1.Image used by these helper methods
type WithRawConfigFile interface {
	// RawConfigFile returns the serialized bytes of this image's config file.
	RawConfigFile() ([]byte, error)
}

// ConfigFile is a helper for implementing v1.Image
func ConfigFile(i WithRawConfigFile) (*v1.ConfigFile, error) {
	b, err := i.RawConfigFile()
	if err != nil {
		return nil, err
	}
	return v1.ParseConfigFile(v1util.NopReadCloser(bytes.NewBuffer(b)))
}

// ConfigName is a helper for implementing v1.Image
func ConfigName(i WithRawConfigFile) (v1.Hash, error) {
	b, err := i.RawConfigFile()
	if err != nil {
		return v1.Hash{}, err
	}
	h, _, err := v1.SHA256(v1util.NopReadCloser(bytes.NewBuffer(b)))
	return h, err
}

// WithConfigFile defines the subset of v1.Image used by these helper methods
type WithConfigFile interface {
	// ConfigFile returns this image's config file.
	ConfigFile() (*v1.ConfigFile, error)
}

// DiffIDs is a helper for implementing v1.Image
func DiffIDs(i WithConfigFile) ([]v1.Hash, error) {
	cfg, err := i.ConfigFile()
	if err != nil {
		return nil, err
	}
	dids := make([]v1.Hash, len(cfg.RootFS.DiffIDs))
	for i, did := range cfg.RootFS.DiffIDs {
		dids[len(dids)-i-1] = did
	}
	return dids, nil
}

// RawConfigFile is a helper for implementing v1.Image
func RawConfigFile(i WithConfigFile) ([]byte, error) {
	cfg, err := i.ConfigFile()
	if err != nil {
		return nil, err
	}
	return json.Marshal(cfg)
}

// WithUncompressedLayer defines the subset of v1.Image used by these helper methods
type WithUncompressedLayer interface {
	// UncompressedLayer is like UncompressedBlob, but takes the "diff id".
	UncompressedLayer(v1.Hash) (io.ReadCloser, error)
}

// Layer is the same as Blob, but takes the "diff id".
func Layer(wul WithUncompressedLayer, h v1.Hash) (io.ReadCloser, error) {
	rc, err := wul.UncompressedLayer(h)
	if err != nil {
		return nil, err
	}
	return v1util.GzipReadCloser(rc)
}

// WithRawManifest defines the subset of v1.Image used by these helper methods
type WithRawManifest interface {
	// RawManifest returns the serialized bytes of this image's config file.
	RawManifest() ([]byte, error)
}

// Digest is a helper for implementing v1.Image
func Digest(i WithRawManifest) (v1.Hash, error) {
	mb, err := i.RawManifest()
	if err != nil {
		return v1.Hash{}, err
	}
	digest, _, err := v1.SHA256(ioutil.NopCloser(bytes.NewReader(mb)))
	return digest, err
}

// Manifest is a helper for implementing v1.Image
func Manifest(i WithRawManifest) (*v1.Manifest, error) {
	b, err := i.RawManifest()
	if err != nil {
		return nil, err
	}
	return v1.ParseManifest(v1util.NopReadCloser(bytes.NewBuffer(b)))
}

// WithManifest defines the subset of v1.Image used by these helper methods
type WithManifest interface {
	// Manifest returns this image's Manifest object.
	Manifest() (*v1.Manifest, error)
}

// RawManifest is a helper for implementing v1.Image
func RawManifest(i WithManifest) ([]byte, error) {
	m, err := i.Manifest()
	if err != nil {
		return nil, err
	}
	return json.Marshal(m)
}

// FSLayers is a helper for implementing v1.Image
func FSLayers(i WithManifest) ([]v1.Hash, error) {
	m, err := i.Manifest()
	if err != nil {
		return nil, err
	}
	fsl := make([]v1.Hash, len(m.Layers))
	for i, l := range m.Layers {
		fsl[len(fsl)-i-1] = l.Digest
	}
	return fsl, nil
}

// BlobSet is a helper for implementing v1.Image
func BlobSet(i WithManifest) (map[v1.Hash]struct{}, error) {
	m, err := i.Manifest()
	if err != nil {
		return nil, err
	}
	bs := make(map[v1.Hash]struct{})
	for _, l := range m.Layers {
		bs[l.Digest] = struct{}{}
	}
	bs[m.Config.Digest] = struct{}{}
	return bs, nil
}

// BlobSize is a helper for implementing v1.Image
func BlobSize(i WithManifest, h v1.Hash) (int64, error) {
	m, err := i.Manifest()
	if err != nil {
		return -1, err
	}
	for _, l := range m.Layers {
		if l.Digest == h {
			return l.Size, nil
		}
	}
	return -1, fmt.Errorf("blob %v not found", h)
}

// WithManifestAndConfigFile defines the subset of v1.Image used by these helper methods
type WithManifestAndConfigFile interface {
	WithConfigFile

	// Manifest returns this image's Manifest object.
	Manifest() (*v1.Manifest, error)
}

// BlobToDiffID is a helper for mapping between compressed
// and uncompressed blob hashes.
func BlobToDiffID(i WithManifestAndConfigFile, h v1.Hash) (v1.Hash, error) {
	blobs, err := FSLayers(i)
	if err != nil {
		return v1.Hash{}, err
	}
	diffIDs, err := DiffIDs(i)
	if err != nil {
		return v1.Hash{}, err
	}
	if len(blobs) != len(diffIDs) {
		return v1.Hash{}, fmt.Errorf("mismatched fs layers (%d) and diff ids (%d)", len(blobs), len(diffIDs))
	}
	for i, blob := range blobs {
		if blob == h {
			return diffIDs[i], nil
		}
	}
	return v1.Hash{}, fmt.Errorf("unknown blob %v", h)
}

// DiffIDToBlob is a helper for mapping between uncompressed
// and compressed blob hashes.
func DiffIDToBlob(wm WithManifestAndConfigFile, h v1.Hash) (v1.Hash, error) {
	blobs, err := FSLayers(wm)
	if err != nil {
		return v1.Hash{}, err
	}
	diffIDs, err := DiffIDs(wm)
	if err != nil {
		return v1.Hash{}, err
	}
	if len(blobs) != len(diffIDs) {
		return v1.Hash{}, fmt.Errorf("mismatched fs layers (%d) and diff ids (%d)", len(blobs), len(diffIDs))
	}
	for i, diffID := range diffIDs {
		if diffID == h {
			return blobs[i], nil
		}
	}
	return v1.Hash{}, fmt.Errorf("unknown diffID %v", h)

}

// WithBlob defines the subset of v1.Image used by these helper methods
type WithBlob interface {
	// Blob returns a ReadCloser for streaming the blob's content.
	Blob(v1.Hash) (io.ReadCloser, error)
}

// UncompressedBlob returns a ReadCloser for streaming the blob's content uncompressed.
func UncompressedBlob(b WithBlob, h v1.Hash) (io.ReadCloser, error) {
	rc, err := b.Blob(h)
	if err != nil {
		return nil, err
	}
	return v1util.GunzipReadCloser(rc)
}
<|endoftext|>


<commit_before>package elasticthought

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"path"
	"strings"

	"github.com/couchbaselabs/cbfs/client"
	"github.com/couchbaselabs/logg"
)

// Worker job that splits a dataset into training/test set
type DatasetSplitter struct {
	Configuration Configuration
	Dataset       Dataset
}

// Run this job
func (d DatasetSplitter) Run() {

	switch d.Dataset.isSplittable() {
	case true:
		d.SplitDatafile()
	default:
		d.DownloadDatafiles()
	}

}

func (d DatasetSplitter) SplitDatafile() {

	// Find the datafile object associated with dataset
	db := d.Configuration.DbConnection()
	datafile, err := d.Dataset.GetSplittableDatafile(db)
	if err != nil {
		errMsg := fmt.Errorf("Error looking up datafile: %+v.  Error: %v", d.Dataset, err)
		d.recordProcessingError(errMsg)
		return
	}

	// Open the url -- content type should be application/x-gzip
	tr, err := openTarGzStream(datafile.Url)
	if err != nil {
		errMsg := fmt.Errorf("Error opening tar.gz streams: %v", err)
		d.recordProcessingError(errMsg)
		return
	}

	// Create pipes
	prTrain, pwTrain := io.Pipe()
	prTest, pwTest := io.Pipe()

	// Wrap in gzip writers
	pwGzTest := gzip.NewWriter(pwTest)
	pwGzTrain := gzip.NewWriter(pwTrain)

	// Create tar writers on the write end of the pipes
	tarWriterTesting := tar.NewWriter(pwGzTest)
	tarWriterTraining := tar.NewWriter(pwGzTrain)

	// Create a cbfs client
	cbfs, err := cbfsclient.New(d.Configuration.CbfsUrl)
	options := cbfsclient.PutOptions{
		ContentType: "application/x-gzip",
	}
	if err != nil {
		errMsg := fmt.Errorf("Error creating cbfs client: %v", err)
		d.recordProcessingError(errMsg)
		return
	}

	// Figure out where to store these on cbfs
	destTraining := d.Dataset.TrainingArtifactPath()
	destTesting := d.Dataset.TestingArtifactPath()

	// Spawn a goroutine that will read from tar.gz reader coming from url data
	// and write to the training and test tar writers (which are on write ends of pipe)
	transformDoneChan := make(chan error, 1)
	go func() {

		// Must close _underlying_ piped writers, or the piped readers will
		// never get an EOF.  Closing just the tar writers that wrap the underlying
		// piped writers is not enough.
		defer pwTest.Close()
		defer pwTrain.Close()
		defer pwGzTest.Close()
		defer pwGzTrain.Close()

		logg.LogTo("DATASET_SPLITTER", "Calling transform")
		err = d.transform(tr, tarWriterTraining, tarWriterTesting)
		if err != nil {
			errMsg := fmt.Errorf("Error transforming tar stream: %v", err)
			logg.LogError(errMsg)
			transformDoneChan <- errMsg
			return
		}

		transformDoneChan <- nil

	}()

	// Spawn goroutines to read off the read ends of the pipe and store in cbfs
	cbfsTrainDoneChan := make(chan error, 1)
	cbfsTestDoneChan := make(chan error, 1)
	go func() {
		if err := cbfs.Put("", destTesting, prTest, options); err != nil {
			errMsg := fmt.Errorf("Error writing %v to cbfs: %v", destTesting, err)
			logg.LogError(errMsg)
			cbfsTestDoneChan <- errMsg
			return

		}
		logg.LogTo("DATASET_SPLITTER", "Wrote %v to cbfs", destTesting)
		cbfsTestDoneChan <- nil
	}()
	go func() {
		if err := cbfs.Put("", destTraining, prTrain, options); err != nil {
			errMsg := fmt.Errorf("Error writing %v to cbfs: %v", destTraining, err)
			logg.LogError(errMsg)
			cbfsTrainDoneChan <- errMsg
			return
		}
		logg.LogTo("DATASET_SPLITTER", "Wrote %v to cbfs", destTraining)
		cbfsTrainDoneChan <- nil
	}()

	// Wait for the results from all the goroutines
	cbfsTrainResult := <-cbfsTrainDoneChan
	cbfsTestResult := <-cbfsTestDoneChan
	transformResult := <-transformDoneChan

	// If any results had an error, log it and return
	results := []error{transformResult, cbfsTestResult, cbfsTrainResult}
	for _, result := range results {
		if result != nil {
			logg.LogTo("DATASET_SPLITTER", "Setting dataset to failed: %v", result)
			d.Dataset.Failed(db, fmt.Errorf("%v", result))
			return
		}
	}

	// Update the state of the dataset to be finished
	d.Dataset.FinishedSuccessfully(db)

}

func (d DatasetSplitter) DownloadDatafiles() {

	// Create a cbfs client
	cbfs, err := cbfsclient.New(d.Configuration.CbfsUrl)
	options := cbfsclient.PutOptions{
		ContentType: "application/x-gzip",
	}
	if err != nil {
		errMsg := fmt.Errorf("Error creating cbfs client: %v", err)
		d.recordProcessingError(errMsg)
		return
	}

	db := d.Configuration.DbConnection()

	source2destEntries := []struct {
		Url      string
		DestPath string
	}{
		{
			Url:      d.Dataset.GetTrainingDatafileUrl(db),
			DestPath: d.Dataset.TrainingArtifactPath(),
		},
		{
			Url:      d.Dataset.GetTestingDatafileUrl(db),
			DestPath: d.Dataset.TestingArtifactPath(),
		},
	}

	for _, source2destEntry := range source2destEntries {

		// open tar.gz stream to source
		// Open the url -- content type should be application/x-gzip
		tr, err := openTarGzStream(source2destEntry.Url)
		if err != nil {
			errMsg := fmt.Errorf("Error opening tar.gz stream to: %v. Err %v", source2destEntry.Url, err)
			d.recordProcessingError(errMsg)
			return
		}

		if err := cbfs.Put("", source2destEntry.DestPath, tr, options); err != nil {
			errMsg := fmt.Errorf("Error writing %v to cbfs: %v", source2destEntry.DestPath, err)
			d.recordProcessingError(errMsg)
			return
		}

	}

	// Update the state of the dataset to be finished
	d.Dataset.FinishedSuccessfully(db)

}

// Codereview: de-dupe
func (d DatasetSplitter) recordProcessingError(err error) {
	logg.LogError(err)
	db := d.Configuration.DbConnection()
	if err := d.Dataset.Failed(db, err); err != nil {
		errMsg := fmt.Errorf("Error setting dataset as failed: %v", err)
		logg.LogError(errMsg)
	}
}

// Read from source tar stream and write training and test to given tar writers
func (d DatasetSplitter) transform(source *tar.Reader, train, test *tar.Writer) error {

	splitter := d.splitter(train, test)

	for {
		hdr, err := source.Next()

		if err == io.EOF {
			// end of tar archive
			break
		}

		if err != nil {
			return err
		}

		tw := splitter(hdr.Name)

		if err := tw.WriteHeader(hdr); err != nil {
			return err
		}
		_, err = io.Copy(tw, source)
		if err != nil {
			return err
		}

	}

	// close writers
	if err := train.Close(); err != nil {
		errMsg := fmt.Errorf("Error closing tar writer: %v", err)
		logg.LogError(errMsg)
		return err
	}
	if err := test.Close(); err != nil {
		errMsg := fmt.Errorf("Error closing tar reader: %v", err)
		logg.LogError(errMsg)
		return err
	}

	return nil
}

func (d DatasetSplitter) splitter(train, test *tar.Writer) func(string) *tar.Writer {

	trainingRatio := int(d.Dataset.TrainingDataset.SplitPercentage * 100)
	testRatio := int(d.Dataset.TestDataset.SplitPercentage * 100)

	ratio := [2]int{trainingRatio, testRatio}

	dircounts := make(map[string][2]int)
	return func(file string) *tar.Writer {
		dir := path.Dir(file)
		counts := dircounts[dir]
		count0ratio1 := counts[0] * ratio[1]
		count1ratio0 := counts[1] * ratio[0]
		if count0ratio1 <= count1ratio0 {
			counts[0]++
			dircounts[dir] = counts
			return train
		} else {
			counts[1]++
			dircounts[dir] = counts
			return test
		}
	}
}

// Validate that the source tar stream conforms to expected specs
func (d DatasetSplitter) validate(source *tar.Reader) (bool, error) {

	// validation rules:
	// 1. has at least 2 files
	// 2. the depth of each file is 2 (folder/filename.xxx)

	numFiles := 0
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return false, err
		}
		numFiles += 1

		pathComponents := strings.Split(hdr.Name, "/")
		if len(pathComponents) != 2 {
			return false, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

	}

	if numFiles < 2 {
		return false, fmt.Errorf("Archive must contain at least 2 files")
	}

	return true, nil
}
<commit_msg>close stream<commit_after>package elasticthought

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"path"
	"strings"

	"github.com/couchbaselabs/cbfs/client"
	"github.com/couchbaselabs/logg"
)

// Worker job that splits a dataset into training/test set
type DatasetSplitter struct {
	Configuration Configuration
	Dataset       Dataset
}

// Run this job
func (d DatasetSplitter) Run() {

	switch d.Dataset.isSplittable() {
	case true:
		d.SplitDatafile()
	default:
		d.DownloadDatafiles()
	}

}

func (d DatasetSplitter) SplitDatafile() {

	// Find the datafile object associated with dataset
	db := d.Configuration.DbConnection()
	datafile, err := d.Dataset.GetSplittableDatafile(db)
	if err != nil {
		errMsg := fmt.Errorf("Error looking up datafile: %+v.  Error: %v", d.Dataset, err)
		d.recordProcessingError(errMsg)
		return
	}

	// Open the url -- content type should be application/x-gzip
	tr, err := openTarGzStream(datafile.Url)
	if err != nil {
		errMsg := fmt.Errorf("Error opening tar.gz streams: %v", err)
		d.recordProcessingError(errMsg)
		return
	}

	// Create pipes
	prTrain, pwTrain := io.Pipe()
	prTest, pwTest := io.Pipe()

	// Wrap in gzip writers
	pwGzTest := gzip.NewWriter(pwTest)
	pwGzTrain := gzip.NewWriter(pwTrain)

	// Create tar writers on the write end of the pipes
	tarWriterTesting := tar.NewWriter(pwGzTest)
	tarWriterTraining := tar.NewWriter(pwGzTrain)

	// Create a cbfs client
	cbfs, err := cbfsclient.New(d.Configuration.CbfsUrl)
	options := cbfsclient.PutOptions{
		ContentType: "application/x-gzip",
	}
	if err != nil {
		errMsg := fmt.Errorf("Error creating cbfs client: %v", err)
		d.recordProcessingError(errMsg)
		return
	}

	// Figure out where to store these on cbfs
	destTraining := d.Dataset.TrainingArtifactPath()
	destTesting := d.Dataset.TestingArtifactPath()

	// Spawn a goroutine that will read from tar.gz reader coming from url data
	// and write to the training and test tar writers (which are on write ends of pipe)
	transformDoneChan := make(chan error, 1)
	go func() {

		// Must close _underlying_ piped writers, or the piped readers will
		// never get an EOF.  Closing just the tar writers that wrap the underlying
		// piped writers is not enough.
		defer pwTest.Close()
		defer pwTrain.Close()
		defer pwGzTest.Close()
		defer pwGzTrain.Close()

		logg.LogTo("DATASET_SPLITTER", "Calling transform")
		err = d.transform(tr, tarWriterTraining, tarWriterTesting)
		if err != nil {
			errMsg := fmt.Errorf("Error transforming tar stream: %v", err)
			logg.LogError(errMsg)
			transformDoneChan <- errMsg
			return
		}

		transformDoneChan <- nil

	}()

	// Spawn goroutines to read off the read ends of the pipe and store in cbfs
	cbfsTrainDoneChan := make(chan error, 1)
	cbfsTestDoneChan := make(chan error, 1)
	go func() {
		if err := cbfs.Put("", destTesting, prTest, options); err != nil {
			errMsg := fmt.Errorf("Error writing %v to cbfs: %v", destTesting, err)
			logg.LogError(errMsg)
			cbfsTestDoneChan <- errMsg
			return

		}
		logg.LogTo("DATASET_SPLITTER", "Wrote %v to cbfs", destTesting)
		cbfsTestDoneChan <- nil
	}()
	go func() {
		if err := cbfs.Put("", destTraining, prTrain, options); err != nil {
			errMsg := fmt.Errorf("Error writing %v to cbfs: %v", destTraining, err)
			logg.LogError(errMsg)
			cbfsTrainDoneChan <- errMsg
			return
		}
		logg.LogTo("DATASET_SPLITTER", "Wrote %v to cbfs", destTraining)
		cbfsTrainDoneChan <- nil
	}()

	// Wait for the results from all the goroutines
	cbfsTrainResult := <-cbfsTrainDoneChan
	cbfsTestResult := <-cbfsTestDoneChan
	transformResult := <-transformDoneChan

	// If any results had an error, log it and return
	results := []error{transformResult, cbfsTestResult, cbfsTrainResult}
	for _, result := range results {
		if result != nil {
			logg.LogTo("DATASET_SPLITTER", "Setting dataset to failed: %v", result)
			d.Dataset.Failed(db, fmt.Errorf("%v", result))
			return
		}
	}

	// Update the state of the dataset to be finished
	d.Dataset.FinishedSuccessfully(db)

}

func (d DatasetSplitter) DownloadDatafiles() {

	// Create a cbfs client
	cbfs, err := cbfsclient.New(d.Configuration.CbfsUrl)
	options := cbfsclient.PutOptions{
		ContentType: "application/x-gzip",
	}
	if err != nil {
		errMsg := fmt.Errorf("Error creating cbfs client: %v", err)
		d.recordProcessingError(errMsg)
		return
	}

	db := d.Configuration.DbConnection()

	source2destEntries := []struct {
		Url      string
		DestPath string
	}{
		{
			Url:      d.Dataset.GetTrainingDatafileUrl(db),
			DestPath: d.Dataset.TrainingArtifactPath(),
		},
		{
			Url:      d.Dataset.GetTestingDatafileUrl(db),
			DestPath: d.Dataset.TestingArtifactPath(),
		},
	}

	for _, source2destEntry := range source2destEntries {

		// open tar.gz stream to source
		// Open the url -- content type should be application/x-gzip

		func() {

			resp, err := http.Get(source2destEntry.Url)
			defer resp.Body.Close()

			if err != nil {
				errMsg := fmt.Errorf("Error opening stream to: %v. Err %v", source2destEntry.Url, err)
				d.recordProcessingError(errMsg)
				return
			}
			gzipReader, err := gzip.NewReader(resp.Body)
			if err != nil {
				errMsg := fmt.Errorf("Error opening gz stream to: %v. Err %v", source2destEntry.Url, err)
				d.recordProcessingError(errMsg)
				return
			}
			tr := tar.NewReader(gzipReader)

			if err := cbfs.Put("", source2destEntry.DestPath, tr, options); err != nil {
				errMsg := fmt.Errorf("Error writing %v to cbfs: %v", source2destEntry.DestPath, err)
				d.recordProcessingError(errMsg)
				return
			}

			logg.LogTo("DATASET_SPLITTER", "Wrote %v to cbfs", source2destEntry.DestPath)

		}()

	}

	// Update the state of the dataset to be finished
	d.Dataset.FinishedSuccessfully(db)

}

// Codereview: de-dupe
func (d DatasetSplitter) recordProcessingError(err error) {
	logg.LogError(err)
	db := d.Configuration.DbConnection()
	if err := d.Dataset.Failed(db, err); err != nil {
		errMsg := fmt.Errorf("Error setting dataset as failed: %v", err)
		logg.LogError(errMsg)
	}
}

// Read from source tar stream and write training and test to given tar writers
func (d DatasetSplitter) transform(source *tar.Reader, train, test *tar.Writer) error {

	splitter := d.splitter(train, test)

	for {
		hdr, err := source.Next()

		if err == io.EOF {
			// end of tar archive
			break
		}

		if err != nil {
			return err
		}

		tw := splitter(hdr.Name)

		if err := tw.WriteHeader(hdr); err != nil {
			return err
		}
		_, err = io.Copy(tw, source)
		if err != nil {
			return err
		}

	}

	// close writers
	if err := train.Close(); err != nil {
		errMsg := fmt.Errorf("Error closing tar writer: %v", err)
		logg.LogError(errMsg)
		return err
	}
	if err := test.Close(); err != nil {
		errMsg := fmt.Errorf("Error closing tar reader: %v", err)
		logg.LogError(errMsg)
		return err
	}

	return nil
}

func (d DatasetSplitter) splitter(train, test *tar.Writer) func(string) *tar.Writer {

	trainingRatio := int(d.Dataset.TrainingDataset.SplitPercentage * 100)
	testRatio := int(d.Dataset.TestDataset.SplitPercentage * 100)

	ratio := [2]int{trainingRatio, testRatio}

	dircounts := make(map[string][2]int)
	return func(file string) *tar.Writer {
		dir := path.Dir(file)
		counts := dircounts[dir]
		count0ratio1 := counts[0] * ratio[1]
		count1ratio0 := counts[1] * ratio[0]
		if count0ratio1 <= count1ratio0 {
			counts[0]++
			dircounts[dir] = counts
			return train
		} else {
			counts[1]++
			dircounts[dir] = counts
			return test
		}
	}
}

// Validate that the source tar stream conforms to expected specs
func (d DatasetSplitter) validate(source *tar.Reader) (bool, error) {

	// validation rules:
	// 1. has at least 2 files
	// 2. the depth of each file is 2 (folder/filename.xxx)

	numFiles := 0
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return false, err
		}
		numFiles += 1

		pathComponents := strings.Split(hdr.Name, "/")
		if len(pathComponents) != 2 {
			return false, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

	}

	if numFiles < 2 {
		return false, fmt.Errorf("Archive must contain at least 2 files")
	}

	return true, nil
}
<|endoftext|>


<commit_before>package db

import (
	"fmt"
	"testing"

	"github.com/lukaszbudnik/migrator/config"
	"github.com/stretchr/testify/assert"
)

func TestDBCreateDialectMSSQLDriver(t *testing.T) {
	config := &config.Config{}
	config.Driver = "sqlserver"
	dialect := newDialect(config)
	assert.IsType(t, &msSQLDialect{}, dialect)
}

func TestMSSQLLastInsertIdSupported(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"
	dialect := newDialect(config)

	lastInsertIDSupported := dialect.LastInsertIDSupported()

	assert.False(t, lastInsertIDSupported)
}

func TestMSSQLGetMigrationInsertSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"

	dialect := newDialect(config)

	insertMigrationSQL := dialect.GetMigrationInsertSQL()

	assert.Equal(t, "insert into migrator.migrator_migrations (name, source_dir, filename, type, db_schema, contents, checksum, version_id) values (@p1, @p2, @p3, @p4, @p5, @p6, @p7, @p8)", insertMigrationSQL)
}

func TestMSSQLGetTenantInsertSQLDefault(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"
	dialect := newDialect(config)
	connector := baseConnector{newTestContext(), config, dialect, nil}
	defer connector.Dispose()

	tenantInsertSQL := connector.getTenantInsertSQL()

	assert.Equal(t, "insert into migrator.migrator_tenants (name) values (@p1)", tenantInsertSQL)
}

func TestMSSQLDialectGetCreateTenantsTableSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"

	dialect := newDialect(config)

	createTenantsTableSQL := dialect.GetCreateTenantsTableSQL()

	expected := `
IF NOT EXISTS (select * from information_schema.tables where table_schema = 'migrator' and table_name = 'migrator_tenants')
BEGIN
  create table [migrator].migrator_tenants (
    id int identity (1,1) primary key,
    name varchar(200) not null,
    created datetime default CURRENT_TIMESTAMP
  );
END
`

	assert.Equal(t, expected, createTenantsTableSQL)
}

func TestMSSQLDialectGetCreateMigrationsTableSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"

	dialect := newDialect(config)

	createMigrationsTableSQL := dialect.GetCreateMigrationsTableSQL()

	expected := `
IF NOT EXISTS (select * from information_schema.tables where table_schema = 'migrator' and table_name = 'migrator_migrations')
BEGIN
  create table [migrator].migrator_migrations (
    id int identity (1,1) primary key,
    name varchar(200) not null,
    source_dir varchar(200) not null,
    filename varchar(200) not null,
    type int not null,
    db_schema varchar(200) not null,
    created datetime default CURRENT_TIMESTAMP,
		contents text,
		checksum varchar(64)
  );
END
`

	assert.Equal(t, expected, createMigrationsTableSQL)
}

func TestMSSQLDialectGetCreateSchemaSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"

	dialect := newDialect(config)

	createSchemaSQL := dialect.GetCreateSchemaSQL("def")

	expected := `
IF NOT EXISTS (select * from information_schema.schemata where schema_name = 'def')
BEGIN
  EXEC sp_executesql N'create schema def';
END
`

	assert.Equal(t, expected, createSchemaSQL)
}

func TestMSSQLGetVersionInsertSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"
	dialect := newDialect(config)

	versionInsertSQL := dialect.GetVersionInsertSQL()

	assert.Equal(t, "insert into migrator.migrator_versions (name) output inserted.id values (@p1)", versionInsertSQL)
}

func TestMSSQLGetCreateVersionsTableSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"
	dialect := newDialect(config)

	actual := dialect.GetCreateVersionsTableSQL()

	expected :=
		`
if not exists (select * from information_schema.tables where table_schema = 'migrator' and table_name = 'migrator_versions')
begin
  declare @cn nvarchar(200);
  create table [migrator].migrator_versions (
    id int identity (1,1) primary key,
    name varchar(200) not null,
    created datetime default CURRENT_TIMESTAMP
  );
  -- workaround for MSSQL not finding a newly created column
  -- when creating initial version default value is set to 1
  alter table [migrator].migrator_migrations add version_id int not null default 1;
  if exists (select * from [migrator].migrator_migrations)
  begin
    insert into [migrator].migrator_versions (name) values ('Initial version');
  end
  -- change version_id to not null
  alter table [migrator].migrator_migrations
    alter column version_id int not null;
  alter table [migrator].migrator_migrations
    add constraint migrator_versions_version_id_fk foreign key (version_id) references [migrator].migrator_versions (id) on delete cascade;
  create index migrator_migrations_version_id_idx on [migrator].migrator_migrations (version_id);
  -- remove workaround default value
  select @cn = name from sys.default_constraints where parent_object_id = object_id('[migrator].migrator_migrations') and name like '%ver%';
  EXEC ('alter table [migrator].migrator_migrations drop constraint ' + @cn);
end
`

	fmt.Println(actual)

	assert.Equal(t, expected, actual[0])
}
<commit_msg>removed test fmt.Println()<commit_after>package db

import (
	"testing"

	"github.com/lukaszbudnik/migrator/config"
	"github.com/stretchr/testify/assert"
)

func TestDBCreateDialectMSSQLDriver(t *testing.T) {
	config := &config.Config{}
	config.Driver = "sqlserver"
	dialect := newDialect(config)
	assert.IsType(t, &msSQLDialect{}, dialect)
}

func TestMSSQLLastInsertIdSupported(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"
	dialect := newDialect(config)

	lastInsertIDSupported := dialect.LastInsertIDSupported()

	assert.False(t, lastInsertIDSupported)
}

func TestMSSQLGetMigrationInsertSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"

	dialect := newDialect(config)

	insertMigrationSQL := dialect.GetMigrationInsertSQL()

	assert.Equal(t, "insert into migrator.migrator_migrations (name, source_dir, filename, type, db_schema, contents, checksum, version_id) values (@p1, @p2, @p3, @p4, @p5, @p6, @p7, @p8)", insertMigrationSQL)
}

func TestMSSQLGetTenantInsertSQLDefault(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"
	dialect := newDialect(config)
	connector := baseConnector{newTestContext(), config, dialect, nil}
	defer connector.Dispose()

	tenantInsertSQL := connector.getTenantInsertSQL()

	assert.Equal(t, "insert into migrator.migrator_tenants (name) values (@p1)", tenantInsertSQL)
}

func TestMSSQLDialectGetCreateTenantsTableSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"

	dialect := newDialect(config)

	createTenantsTableSQL := dialect.GetCreateTenantsTableSQL()

	expected := `
IF NOT EXISTS (select * from information_schema.tables where table_schema = 'migrator' and table_name = 'migrator_tenants')
BEGIN
  create table [migrator].migrator_tenants (
    id int identity (1,1) primary key,
    name varchar(200) not null,
    created datetime default CURRENT_TIMESTAMP
  );
END
`

	assert.Equal(t, expected, createTenantsTableSQL)
}

func TestMSSQLDialectGetCreateMigrationsTableSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"

	dialect := newDialect(config)

	createMigrationsTableSQL := dialect.GetCreateMigrationsTableSQL()

	expected := `
IF NOT EXISTS (select * from information_schema.tables where table_schema = 'migrator' and table_name = 'migrator_migrations')
BEGIN
  create table [migrator].migrator_migrations (
    id int identity (1,1) primary key,
    name varchar(200) not null,
    source_dir varchar(200) not null,
    filename varchar(200) not null,
    type int not null,
    db_schema varchar(200) not null,
    created datetime default CURRENT_TIMESTAMP,
		contents text,
		checksum varchar(64)
  );
END
`

	assert.Equal(t, expected, createMigrationsTableSQL)
}

func TestMSSQLDialectGetCreateSchemaSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"

	dialect := newDialect(config)

	createSchemaSQL := dialect.GetCreateSchemaSQL("def")

	expected := `
IF NOT EXISTS (select * from information_schema.schemata where schema_name = 'def')
BEGIN
  EXEC sp_executesql N'create schema def';
END
`

	assert.Equal(t, expected, createSchemaSQL)
}

func TestMSSQLGetVersionInsertSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"
	dialect := newDialect(config)

	versionInsertSQL := dialect.GetVersionInsertSQL()

	assert.Equal(t, "insert into migrator.migrator_versions (name) output inserted.id values (@p1)", versionInsertSQL)
}

func TestMSSQLGetCreateVersionsTableSQL(t *testing.T) {
	config, err := config.FromFile("../test/migrator.yaml")
	assert.Nil(t, err)

	config.Driver = "sqlserver"
	dialect := newDialect(config)

	actual := dialect.GetCreateVersionsTableSQL()

	expected :=
		`
if not exists (select * from information_schema.tables where table_schema = 'migrator' and table_name = 'migrator_versions')
begin
  declare @cn nvarchar(200);
  create table [migrator].migrator_versions (
    id int identity (1,1) primary key,
    name varchar(200) not null,
    created datetime default CURRENT_TIMESTAMP
  );
  -- workaround for MSSQL not finding a newly created column
  -- when creating initial version default value is set to 1
  alter table [migrator].migrator_migrations add version_id int not null default 1;
  if exists (select * from [migrator].migrator_migrations)
  begin
    insert into [migrator].migrator_versions (name) values ('Initial version');
  end
  -- change version_id to not null
  alter table [migrator].migrator_migrations
    alter column version_id int not null;
  alter table [migrator].migrator_migrations
    add constraint migrator_versions_version_id_fk foreign key (version_id) references [migrator].migrator_versions (id) on delete cascade;
  create index migrator_migrations_version_id_idx on [migrator].migrator_migrations (version_id);
  -- remove workaround default value
  select @cn = name from sys.default_constraints where parent_object_id = object_id('[migrator].migrator_migrations') and name like '%ver%';
  EXEC ('alter table [migrator].migrator_migrations drop constraint ' + @cn);
end
`

	assert.Equal(t, expected, actual[0])
}
<|endoftext|>