<commit_before>// Copyright 2016 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package main

import (
	"database/sql"
	"errors"
	"log"
	"math/rand"
	"time"

	"github.com/cockroachdb/cockroach/util"
	_ "github.com/lib/pq"
)

const (
	// TODO(spencer): update the CREATE DATABASE statement in the schema
	//   to pull out the database specified in the DB URL and use it instead
	//   of "photos" below.
	photosSchema = `
CREATE DATABASE IF NOT EXISTS photos;

CREATE TABLE IF NOT EXISTS users (
  id           INT,
  photoCount   INT,
  commentCount INT,
  name         STRING,
  address      STRING,

  PRIMARY KEY (id)
);

CREATE TABLE IF NOT EXISTS photos (
  id           BYTES DEFAULT EXPERIMENTAL_UUID_V4(),
  userID       INT,
  commentCount INT,
  caption      STRING,
  latitude     FLOAT,
  longitude    FLOAT,
  timestamp    TIMESTAMP,

  PRIMARY KEY (id)
);

CREATE UNIQUE INDEX IF NOT EXISTS photosByUserID ON photos (userID, timestamp);

CREATE TABLE IF NOT EXISTS comments (
  photoID   BYTES,
  commentID BYTES DEFAULT EXPERIMENTAL_UUID_V4(),
  userID    INT,
  message   STRING,
  timestamp TIMESTAMP,

  PRIMARY KEY (photoID, commentID)
);

CREATE UNIQUE INDEX IF NOT EXISTS commentsByPhotoID ON comments (photoID, timestamp);
`
)

// openDB opens the database connection according to the context.
func openDB(ctx Context) (*sql.DB, error) {
	return sql.Open("postgres", ctx.DBUrl)
}

// initSchema creates the database schema if it doesn't exist.
func initSchema(db *sql.DB) error {
	_, err := db.Exec(photosSchema)
	return err
}

// dropDatabase drops the database.
func dropDatabase(db *sql.DB) error {
	_, err := db.Exec("DROP DATABASE IF EXISTS photos;")
	return err
}

// execute runs a closure within a transaction. On failure,
// the transaction is aborted and rolled back; on success,
// the transaction is committed. If the error string matches
// CockroachDB's restart directive, the same transaction is
// re-used to retry executing the closure.
//
// NOTE: the supplied exec closure should not have external
//       side effects beyond changes to the database.
func executeTxn(db *sql.DB, exec func(*sql.Tx) error) error {
	tx, err := ctx.DB.Begin()
	if err != nil {
		return err
	}
	for {
		if err := exec(tx); err != nil {
			// TODO(spencer): check error return here and if retry txn is
			// warranted, continue.
			/*
				if err.Error().Matches("retry cockroachdb txn") {
					if _, err = tx.Exec(`RETRY TRANSACTION;`); err == nil {
						continue
					}
				}
			*/
			_ = tx.Rollback()
			return err
		}
		return tx.Commit()
	}
}

const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

func randString(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = letterBytes[rand.Intn(len(letterBytes))]
	}
	return string(b)
}

var noUserErr = errors.New("no user found")

// userExists looks up a user by ID.
func userExists(tx *sql.Tx, userID int) (bool, error) {
	var id int
	const selectSQL = `
SELECT id FROM users WHERE id = $1;
`
	err := tx.QueryRow(selectSQL, userID).Scan(&id)
	switch err {
	case sql.ErrNoRows:
		return false, nil
	case nil:
		return true, nil
	default:
		return false, err
	}
}

// findClosestUserByID selects the first user which exists with
// id >= userID. Returns the found user ID or an error.
func findClosestUserByID(tx *sql.Tx, userID int) (int, error) {
	var id int
	const selectSQL = `
SELECT id FROM users WHERE id >= $1 ORDER BY id LIMIT 1;
`
	err := tx.QueryRow(selectSQL, userID).Scan(&id)
	if err == sql.ErrNoRows {
		return 0, noUserErr
	}
	return id, err
}

// createUser creates a new user with random name and address strings.
func createUser(tx *sql.Tx, userID int) error {
	exists, err := userExists(tx, userID)
	if err != nil || exists {
		return err
	}
	const insertSQL = `
INSERT INTO users VALUES ($1, 0, 0, $2, $3);
`
	const minNameLen = 1
	const maxNameLen = 30
	const minAddrLen = 20
	const maxAddrLen = 100
	name := randString(minNameLen + rand.Intn(maxNameLen-minNameLen))
	addr := randString(minAddrLen + rand.Intn(maxAddrLen-minAddrLen))
	_, err = tx.Exec(insertSQL, userID, name, addr)
	return err
}

// createPhoto looks up or creates a new user to match userID (it's
// the only method in this interface which doesn't match an existing
// user except for createUser). It then creates a new photo for the
// new or pre-existing user.
func createPhoto(tx *sql.Tx, userID int) error {
	if err := createUser(tx, userID); err != nil {
		return err
	}

	const insertSQL = `
INSERT INTO photos VALUES (DEFAULT, $1, 0, $2, $3, $4, NOW());
`
	const minCaptionLen = 10
	const maxCaptionLen = 200
	caption := randString(minCaptionLen + rand.Intn(maxCaptionLen-minCaptionLen))
	latitude := rand.Float32() * 90
	longitude := rand.Float32() * 180
	if _, err := tx.Exec(insertSQL, userID, caption, latitude, longitude); err != nil {
		return err
	}

	const updateSQL = `
UPDATE users SET photoCount = photoCount + 1 WHERE id = $1;
`
	if _, err := tx.Exec(updateSQL, userID); err != nil {
		return err
	}
	return nil
}

// createComment chooses a random photo from a user with the closest
// matching user ID and generates a random author ID to author the
// comment. Counts are updated on the photo and author user.
func createComment(tx *sql.Tx, userID int) error {
	photoID, err := chooseRandomPhoto(tx, userID)
	if err != nil {
		return err
	}
	authorID := rand.Intn(userID) + 1

	const insertSQL = `
INSERT INTO comments VALUES ($1, DEFAULT, $2, $3, NOW());
`
	const minMessageLen = 32
	const maxMessageLen = 1024
	message := randString(minMessageLen + rand.Intn(maxMessageLen-minMessageLen))
	if _, err := tx.Exec(insertSQL, photoID, authorID, message); err != nil {
		log.Printf("insert into photos failed: %s", err)
		return err
	}

	const updatePhotoSQL = `
UPDATE photos SET commentCount = commentCount + 1 WHERE id = $1;
`
	if _, err := tx.Exec(updatePhotoSQL, photoID); err != nil {
		return err
	}

	const updateUserSQL = `
UPDATE users SET commentCount = commentCount + 1 WHERE id = $1;
`
	if _, err := tx.Exec(updateUserSQL, authorID); err != nil {
		return err
	}
	return nil
}

// listPhotos queries up to 100 photos, sorted by timestamp in
// descending order, for the first user with ID >= userID. If photoIDs
// is not nil, stores the queried photo IDs in photoIDs.
func listPhotos(tx *sql.Tx, userID int, photoIDs *[][]byte) error {
	var err error
	userID, err = findClosestUserByID(tx, userID)
	if err != nil {
		return err
	}
	const selectSQL = `
SELECT id, caption, commentCount, latitude, longitude, timestamp FROM photos WHERE userID = $1 ORDER BY timestamp DESC LIMIT 100`
	rows, err := tx.Query(selectSQL, userID)
	switch {
	case err == sql.ErrNoRows:
		return nil
	case err != nil:
		return err
	}
	defer rows.Close()
	// Count and process the result set so we make sure work is done to
	// stream the results.
	var count int
	for rows.Next() {
		if err := rows.Err(); err != nil {
			return err
		}
		var id []byte
		var caption string
		var cCount int
		var lat, lon float64
		var ts time.Time
		if err := rows.Scan(&id, &caption, &cCount, &lat, &lon, &ts); err != nil {
			return util.Errorf("failed to scan result set for user %d: %s", userID, err)
		}
		count++
		if photoIDs != nil {
			*photoIDs = append(*photoIDs, id)
		}
	}
	//log.Printf("selected %d photos for user %d", count, userID)
	return nil
}

// chooseRandomPhoto selects a random photo for the specified
// user or an existing user with the closest user ID. Returns
// the photo ID or an error.
func chooseRandomPhoto(tx *sql.Tx, userID int) ([]byte, error) {
	photoIDs := [][]byte{}
	if err := listPhotos(tx, userID, &photoIDs); err != nil {
		return nil, err
	}
	if len(photoIDs) == 0 {
		return nil, nil
	}
	photoID := photoIDs[rand.Intn(len(photoIDs))]
	return photoID, nil
}

// listComments chooses a random photo and lists up to 100 of its
// comments. Returns the photoID or an error. If the commentIDs slice
// is not nil, it's set to the queried comments' IDs.
func listComments(tx *sql.Tx, userID int, commentIDs *[][]byte) ([]byte, error) {
	photoID, err := chooseRandomPhoto(tx, userID)
	if err != nil {
		return nil, err
	}
	const selectSQL = `
SELECT commentID, userID, message, timestamp FROM comments WHERE photoID = $1 ORDER BY timestamp DESC LIMIT 100`
	rows, err := tx.Query(selectSQL, photoID)
	switch {
	case err == sql.ErrNoRows:
		return photoID, nil
	case err != nil:
		return nil, err
	}
	defer rows.Close()
	// Count and process the result set so we make sure work is done to
	// stream the results.
	var count int
	for rows.Next() {
		if err := rows.Err(); err != nil {
			return nil, err
		}
		var commentID []byte
		var message string
		var userID int
		var ts time.Time
		if err := rows.Scan(&commentID, &userID, &message, &ts); err != nil {
			return nil, util.Errorf("failed to scan result set for photo %q: %s", photoID, err)
		}
		count++
		if commentIDs != nil {
			*commentIDs = append(*commentIDs, commentID)
		}
	}
	//log.Printf("selected %d comments for photo %q", count, photoID)
	return photoID, nil
}

// chooseRandomComment selects a random comment for the specified
// user or an existing user with the closest user ID. Returns
// the photo and comment IDs or an error.
func chooseRandomComment(tx *sql.Tx, userID int) ([]byte, []byte, error) {
	commentIDs := [][]byte{}
	photoID, err := listComments(tx, userID, &commentIDs)
	if err != nil {
		return nil, nil, err
	}
	if len(commentIDs) == 0 {
		return photoID, nil, nil
	}
	commentID := commentIDs[rand.Intn(len(commentIDs))]
	return photoID, commentID, nil
}

func updatePhoto(tx *sql.Tx, userID int) error {
	photoID, err := chooseRandomPhoto(tx, userID)
	if err != nil {
		return err
	}

	const updatePhotoSQL = `
UPDATE photos SET caption = $1 WHERE id = $2;
`
	const minCaptionLen = 10
	const maxCaptionLen = 200
	caption := randString(minCaptionLen + rand.Intn(maxCaptionLen-minCaptionLen))
	if _, err := tx.Exec(updatePhotoSQL, caption, photoID); err != nil {
		return err
	}
	return nil
}

func updateComment(tx *sql.Tx, userID int) error {
	photoID, commentID, err := chooseRandomComment(tx, userID)
	if err != nil {
		return err
	}

	const updateCommentSQL = `
UPDATE comments SET message = $1 WHERE photoID = $2 AND commentID = $3;
`
	const minMessageLen = 10
	const maxMessageLen = 200
	message := randString(minMessageLen + rand.Intn(maxMessageLen-minMessageLen))
	if _, err := tx.Exec(updateCommentSQL, message, photoID, commentID); err != nil {
		return err
	}
	return nil
}

func deletePhoto(tx *sql.Tx, userID int) error {
	photoID, err := chooseRandomPhoto(tx, userID)
	if err != nil {
		return err
	}
	const deletePhotoSQL = `
DELETE FROM photos WHERE id = $1;
`
	if _, err := tx.Exec(deletePhotoSQL, photoID); err != nil {
		return err
	}

	const updateSQL = `
UPDATE users SET photoCount = photoCount - 1 WHERE id = $1;
`
	if _, err := tx.Exec(updateSQL, userID); err != nil {
		return err
	}
	return nil
}

func deleteComment(tx *sql.Tx, userID int) error {
	photoID, commentID, err := chooseRandomComment(tx, userID)
	if err != nil {
		return err
	}
	const deleteCommentSQL = `
DELETE FROM comments WHERE photoID = $1 AND commentID = $2;
`
	if _, err := tx.Exec(deleteCommentSQL, photoID, commentID); err != nil {
		return err
	}

	const updatePhotoSQL = `
UPDATE photos SET commentCount = commentCount - 1 WHERE id = $1;
`
	if _, err := tx.Exec(updatePhotoSQL, photoID); err != nil {
		return err
	}

	const updateUserSQL = `
UPDATE users SET commentCount = commentCount - 1 WHERE id = $1;
`
	if _, err := tx.Exec(updateUserSQL, userID); err != nil {
		return err
	}
	return nil
}
<commit_msg>Use more efficient query<commit_after>// Copyright 2016 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package main

import (
	"database/sql"
	"errors"
	"log"
	"math/rand"
	"time"

	"github.com/cockroachdb/cockroach/util"
	_ "github.com/lib/pq"
)

const (
	// TODO(spencer): update the CREATE DATABASE statement in the schema
	//   to pull out the database specified in the DB URL and use it instead
	//   of "photos" below.
	photosSchema = `
CREATE DATABASE IF NOT EXISTS photos;

CREATE TABLE IF NOT EXISTS users (
  id           INT,
  photoCount   INT,
  commentCount INT,
  name         STRING,
  address      STRING,

  PRIMARY KEY (id)
);

CREATE TABLE IF NOT EXISTS photos (
  id           BYTES DEFAULT EXPERIMENTAL_UUID_V4(),
  userID       INT,
  commentCount INT,
  caption      STRING,
  latitude     FLOAT,
  longitude    FLOAT,
  timestamp    TIMESTAMP,

  PRIMARY KEY (id)
);

CREATE UNIQUE INDEX IF NOT EXISTS photosByUserID ON photos (userID, timestamp);

CREATE TABLE IF NOT EXISTS comments (
  photoID   BYTES,
  commentID BYTES DEFAULT EXPERIMENTAL_UUID_V4(),
  userID    INT,
  message   STRING,
  timestamp TIMESTAMP,

  PRIMARY KEY (photoID, commentID)
);

CREATE UNIQUE INDEX IF NOT EXISTS commentsByPhotoID ON comments (photoID, timestamp);
`
)

// openDB opens the database connection according to the context.
func openDB(ctx Context) (*sql.DB, error) {
	return sql.Open("postgres", ctx.DBUrl)
}

// initSchema creates the database schema if it doesn't exist.
func initSchema(db *sql.DB) error {
	_, err := db.Exec(photosSchema)
	return err
}

// dropDatabase drops the database.
func dropDatabase(db *sql.DB) error {
	_, err := db.Exec("DROP DATABASE IF EXISTS photos;")
	return err
}

// execute runs a closure within a transaction. On failure,
// the transaction is aborted and rolled back; on success,
// the transaction is committed. If the error string matches
// CockroachDB's restart directive, the same transaction is
// re-used to retry executing the closure.
//
// NOTE: the supplied exec closure should not have external
//       side effects beyond changes to the database.
func executeTxn(db *sql.DB, exec func(*sql.Tx) error) error {
	tx, err := ctx.DB.Begin()
	if err != nil {
		return err
	}
	for {
		if err := exec(tx); err != nil {
			// TODO(spencer): check error return here and if retry txn is
			// warranted, continue.
			/*
				if err.Error().Matches("retry cockroachdb txn") {
					if _, err = tx.Exec(`RETRY TRANSACTION;`); err == nil {
						continue
					}
				}
			*/
			_ = tx.Rollback()
			return err
		}
		return tx.Commit()
	}
}

const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

func randString(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = letterBytes[rand.Intn(len(letterBytes))]
	}
	return string(b)
}

var noUserErr = errors.New("no user found")

// userExists looks up a user by ID.
func userExists(tx *sql.Tx, userID int) (bool, error) {
	var id int
	const selectSQL = `
SELECT id FROM users WHERE id = $1;
`
	err := tx.QueryRow(selectSQL, userID).Scan(&id)
	switch err {
	case sql.ErrNoRows:
		return false, nil
	case nil:
		return true, nil
	default:
		return false, err
	}
}

// findClosestUserByID selects the first user which exists with
// id >= userID. Returns the found user ID or an error.
func findClosestUserByID(tx *sql.Tx, userID int) (int, error) {
	var id int
	const selectSQL = `
SELECT id FROM users WHERE id >= $1 ORDER BY id LIMIT 1;
`
	err := tx.QueryRow(selectSQL, userID).Scan(&id)
	if err == sql.ErrNoRows {
		return 0, noUserErr
	}
	return id, err
}

// createUser creates a new user with random name and address strings.
func createUser(tx *sql.Tx, userID int) error {
	exists, err := userExists(tx, userID)
	if err != nil || exists {
		return err
	}
	const insertSQL = `
INSERT INTO users VALUES ($1, 0, 0, $2, $3);
`
	const minNameLen = 1
	const maxNameLen = 30
	const minAddrLen = 20
	const maxAddrLen = 100
	name := randString(minNameLen + rand.Intn(maxNameLen-minNameLen))
	addr := randString(minAddrLen + rand.Intn(maxAddrLen-minAddrLen))
	_, err = tx.Exec(insertSQL, userID, name, addr)
	return err
}

// createPhoto looks up or creates a new user to match userID (it's
// the only method in this interface which doesn't match an existing
// user except for createUser). It then creates a new photo for the
// new or pre-existing user.
func createPhoto(tx *sql.Tx, userID int) error {
	if err := createUser(tx, userID); err != nil {
		return err
	}

	const insertSQL = `
INSERT INTO photos VALUES (DEFAULT, $1, 0, $2, $3, $4, NOW());
`
	const minCaptionLen = 10
	const maxCaptionLen = 200
	caption := randString(minCaptionLen + rand.Intn(maxCaptionLen-minCaptionLen))
	latitude := rand.Float32() * 90
	longitude := rand.Float32() * 180
	if _, err := tx.Exec(insertSQL, userID, caption, latitude, longitude); err != nil {
		return err
	}

	const updateSQL = `
UPDATE users SET photoCount = photoCount + 1 WHERE id = $1;
`
	if _, err := tx.Exec(updateSQL, userID); err != nil {
		return err
	}
	return nil
}

// createComment chooses a random photo from a user with the closest
// matching user ID and generates a random author ID to author the
// comment. Counts are updated on the photo and author user.
func createComment(tx *sql.Tx, userID int) error {
	photoID, err := chooseRandomPhoto(tx, userID)
	if err != nil {
		return err
	}
	authorID := rand.Intn(userID) + 1

	const insertSQL = `
INSERT INTO comments VALUES ($1, DEFAULT, $2, $3, NOW());
`
	const minMessageLen = 32
	const maxMessageLen = 1024
	message := randString(minMessageLen + rand.Intn(maxMessageLen-minMessageLen))
	if _, err := tx.Exec(insertSQL, photoID, authorID, message); err != nil {
		log.Printf("insert into photos failed: %s", err)
		return err
	}

	const updatePhotoSQL = `
UPDATE photos SET commentCount = commentCount + 1 WHERE id = $1;
`
	if _, err := tx.Exec(updatePhotoSQL, photoID); err != nil {
		return err
	}

	const updateUserSQL = `
UPDATE users SET commentCount = commentCount + 1 WHERE id = $1;
`
	if _, err := tx.Exec(updateUserSQL, authorID); err != nil {
		return err
	}
	return nil
}

// listPhotos queries up to 100 photos, sorted by timestamp in
// descending order, for the first user with ID >= userID. If photoIDs
// is not nil, stores the queried photo IDs in photoIDs.
func listPhotos(tx *sql.Tx, userID int, photoIDs *[][]byte) error {
	var err error
	userID, err = findClosestUserByID(tx, userID)
	if err != nil {
		return err
	}
	const selectSQL = `
SELECT id, caption, commentCount, latitude, longitude, timestamp FROM photos WHERE userID = $1 ORDER BY timestamp DESC LIMIT 100`
	rows, err := tx.Query(selectSQL, userID)
	switch {
	case err == sql.ErrNoRows:
		return nil
	case err != nil:
		return err
	}
	defer rows.Close()
	// Count and process the result set so we make sure work is done to
	// stream the results.
	var count int
	for rows.Next() {
		if err := rows.Err(); err != nil {
			return err
		}
		var id []byte
		var caption string
		var cCount int
		var lat, lon float64
		var ts time.Time
		if err := rows.Scan(&id, &caption, &cCount, &lat, &lon, &ts); err != nil {
			return util.Errorf("failed to scan result set for user %d: %s", userID, err)
		}
		count++
		if photoIDs != nil {
			*photoIDs = append(*photoIDs, id)
		}
	}
	//log.Printf("selected %d photos for user %d", count, userID)
	return nil
}

// chooseRandomPhoto selects a random photo for the specified
// user or an existing user with the closest user ID. Returns
// the photo ID or an error.
func chooseRandomPhoto(tx *sql.Tx, userID int) ([]byte, error) {
	photoIDs := [][]byte{}
	if err := listPhotos(tx, userID, &photoIDs); err != nil {
		return nil, err
	}
	if len(photoIDs) == 0 {
		return nil, nil
	}
	photoID := photoIDs[rand.Intn(len(photoIDs))]
	return photoID, nil
}

// listComments chooses a random photo and lists up to 100 of its
// comments. Returns the photoID or an error. If the commentIDs slice
// is not nil, it's set to the queried comments' IDs.
func listComments(tx *sql.Tx, userID int, commentIDs *[][]byte) ([]byte, error) {
	photoID, err := chooseRandomPhoto(tx, userID)
	if err != nil {
		return nil, err
	}
	const selectSQL = `SELECT commentID, userID, message, timestamp FROM comments ` +
		`WHERE photoID = $1 AND commentID IN ` +
		`(SELECT commentID FROM comments WHERE photoID = $1 ORDER BY timestamp DESC LIMIT 100)` +
		`ORDER BY timestamp DESC`
	rows, err := tx.Query(selectSQL, photoID)
	switch {
	case err == sql.ErrNoRows:
		return photoID, nil
	case err != nil:
		return nil, err
	}
	defer rows.Close()
	// Count and process the result set so we make sure work is done to
	// stream the results.
	var count int
	for rows.Next() {
		if err := rows.Err(); err != nil {
			return nil, err
		}
		var commentID []byte
		var message string
		var userID int
		var ts time.Time
		if err := rows.Scan(&commentID, &userID, &message, &ts); err != nil {
			return nil, util.Errorf("failed to scan result set for photo %q: %s", photoID, err)
		}
		count++
		if commentIDs != nil {
			*commentIDs = append(*commentIDs, commentID)
		}
	}
	//log.Printf("selected %d comments for photo %q", count, photoID)
	return photoID, nil
}

// chooseRandomComment selects a random comment for the specified
// user or an existing user with the closest user ID. Returns
// the photo and comment IDs or an error.
func chooseRandomComment(tx *sql.Tx, userID int) ([]byte, []byte, error) {
	commentIDs := [][]byte{}
	photoID, err := listComments(tx, userID, &commentIDs)
	if err != nil {
		return nil, nil, err
	}
	if len(commentIDs) == 0 {
		return photoID, nil, nil
	}
	commentID := commentIDs[rand.Intn(len(commentIDs))]
	return photoID, commentID, nil
}

func updatePhoto(tx *sql.Tx, userID int) error {
	photoID, err := chooseRandomPhoto(tx, userID)
	if err != nil {
		return err
	}

	const updatePhotoSQL = `
UPDATE photos SET caption = $1 WHERE id = $2;
`
	const minCaptionLen = 10
	const maxCaptionLen = 200
	caption := randString(minCaptionLen + rand.Intn(maxCaptionLen-minCaptionLen))
	if _, err := tx.Exec(updatePhotoSQL, caption, photoID); err != nil {
		return err
	}
	return nil
}

func updateComment(tx *sql.Tx, userID int) error {
	photoID, commentID, err := chooseRandomComment(tx, userID)
	if err != nil {
		return err
	}

	const updateCommentSQL = `
UPDATE comments SET message = $1 WHERE photoID = $2 AND commentID = $3;
`
	const minMessageLen = 10
	const maxMessageLen = 200
	message := randString(minMessageLen + rand.Intn(maxMessageLen-minMessageLen))
	if _, err := tx.Exec(updateCommentSQL, message, photoID, commentID); err != nil {
		return err
	}
	return nil
}

func deletePhoto(tx *sql.Tx, userID int) error {
	photoID, err := chooseRandomPhoto(tx, userID)
	if err != nil {
		return err
	}
	const deletePhotoSQL = `
DELETE FROM photos WHERE id = $1;
`
	if _, err := tx.Exec(deletePhotoSQL, photoID); err != nil {
		return err
	}

	const updateSQL = `
UPDATE users SET photoCount = photoCount - 1 WHERE id = $1;
`
	if _, err := tx.Exec(updateSQL, userID); err != nil {
		return err
	}
	return nil
}

func deleteComment(tx *sql.Tx, userID int) error {
	photoID, commentID, err := chooseRandomComment(tx, userID)
	if err != nil {
		return err
	}
	const deleteCommentSQL = `
DELETE FROM comments WHERE photoID = $1 AND commentID = $2;
`
	if _, err := tx.Exec(deleteCommentSQL, photoID, commentID); err != nil {
		return err
	}

	const updatePhotoSQL = `
UPDATE photos SET commentCount = commentCount - 1 WHERE id = $1;
`
	if _, err := tx.Exec(updatePhotoSQL, photoID); err != nil {
		return err
	}

	const updateUserSQL = `
UPDATE users SET commentCount = commentCount - 1 WHERE id = $1;
`
	if _, err := tx.Exec(updateUserSQL, userID); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2017 The Neugram Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package format

import (
	"bytes"

	"neugram.io/ng/stmt"
)

func (p *printer) stmt(s stmt.Stmt) {
	switch s := s.(type) {
	case *stmt.Simple:
		p.expr(s.Expr)
	case *stmt.Return:
		p.buf.WriteString("return")
		if len(s.Exprs) > 0 {
			p.buf.WriteByte(' ')
		}
		for i, e := range s.Exprs {
			if i > 0 {
				p.buf.WriteString(", ")
			}
			p.expr(e)
		}
	default:
		p.printf("format: unknown stmt %T: ", s)
		WriteDebug(p.buf, s)
	}
}

func WriteStmt(buf *bytes.Buffer, s stmt.Stmt) {
	p := &printer{
		buf: buf,
	}
	p.stmt(s)
}

func Stmt(e stmt.Stmt) string {
	buf := new(bytes.Buffer)
	WriteStmt(buf, e)
	return buf.String()
}
<commit_msg>format: add support for Assign, Send, Select and Block<commit_after>// Copyright 2017 The Neugram Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package format

import (
	"bytes"

	"neugram.io/ng/stmt"
)

func (p *printer) stmt(s stmt.Stmt) {
	switch s := s.(type) {
	case *stmt.Simple:
		p.expr(s.Expr)
	case *stmt.Return:
		p.buf.WriteString("return")
		if len(s.Exprs) > 0 {
			p.buf.WriteByte(' ')
		}
		for i, e := range s.Exprs {
			if i > 0 {
				p.buf.WriteString(", ")
			}
			p.expr(e)
		}
	case *stmt.Assign:
		for i, e := range s.Left {
			if i > 0 {
				p.buf.WriteString(", ")
			}
			p.expr(e)
		}
		p.buf.WriteString(" ")
		if s.Decl {
			p.buf.WriteString(":")
		}
		p.buf.WriteString("= ")
		for i, e := range s.Right {
			if i > 0 {
				p.buf.WriteString(", ")
			}
			p.expr(e)
		}
	case *stmt.Send:
		p.expr(s.Chan)
		p.buf.WriteString("<-")
		p.expr(s.Value)
	case *stmt.Select:
		p.buf.WriteString("select {")
		for _, c := range s.Cases {
			switch c.Default {
			case true:
				p.buf.WriteString("default:")
			default:
				p.buf.WriteString("case ")
				p.stmt(c.Stmt)
				p.buf.WriteString(":")
			}
			p.stmt(c.Body)
		}
		p.buf.WriteString("}")
	case *stmt.Block:
		p.buf.WriteString("{")
		for _, s := range s.Stmts {
			p.stmt(s)
		}
		p.buf.WriteString("}")
	default:
		p.printf("format: unknown stmt %T: ", s)
		WriteDebug(p.buf, s)
	}
}

func WriteStmt(buf *bytes.Buffer, s stmt.Stmt) {
	p := &printer{
		buf: buf,
	}
	p.stmt(s)
}

func Stmt(e stmt.Stmt) string {
	buf := new(bytes.Buffer)
	WriteStmt(buf, e)
	return buf.String()
}
<|endoftext|>


<commit_before>// Copyright 2018 The etcd 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 tester

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

	"github.com/coreos/etcd/tools/functional-tester/rpcpb"

	"go.uber.org/zap"
)

// Failure defines failure injection interface.
// To add a fail case:
//  1. implement "Failure" interface
//  2. define fail case name in "rpcpb.FailureCase"
type Failure interface {
	// Inject injeccts the failure into the testing cluster at the given
	// round. When calling the function, the cluster should be in health.
	Inject(clus *Cluster) error
	// Recover recovers the injected failure caused by the injection of the
	// given round and wait for the recovery of the testing cluster.
	Recover(clus *Cluster) error
	// Desc returns a description of the failure
	Desc() string
	// FailureCase returns "rpcpb.FailureCase" enum type.
	FailureCase() rpcpb.FailureCase
}

type injectMemberFunc func(*Cluster, int) error
type recoverMemberFunc func(*Cluster, int) error

type failureByFunc struct {
	desc          string
	failureCase   rpcpb.FailureCase
	injectMember  injectMemberFunc
	recoverMember recoverMemberFunc
}

func (f *failureByFunc) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	return f.failureCase.String()
}

func (f *failureByFunc) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

type failureFollower struct {
	failureByFunc
	last int
	lead int
}

func (f *failureFollower) updateIndex(clus *Cluster) error {
	idx, err := clus.GetLeader()
	if err != nil {
		return err
	}
	f.lead = idx

	n := len(clus.Members)
	if f.last == -1 { // first run
		f.last = clus.rd % n
		if f.last == f.lead {
			f.last = (f.last + 1) % n
		}
	} else {
		f.last = (f.last + 1) % n
		if f.last == f.lead {
			f.last = (f.last + 1) % n
		}
	}
	return nil
}

func (f *failureFollower) Inject(clus *Cluster) error {
	if err := f.updateIndex(clus); err != nil {
		return err
	}
	return f.injectMember(clus, f.last)
}

func (f *failureFollower) Recover(clus *Cluster) error {
	return f.recoverMember(clus, f.last)
}

func (f *failureFollower) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	return f.failureCase.String()
}

func (f *failureFollower) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

type failureLeader struct {
	failureByFunc
	last int
	lead int
}

func (f *failureLeader) updateIndex(clus *Cluster) error {
	idx, err := clus.GetLeader()
	if err != nil {
		return err
	}
	f.lead = idx
	f.last = idx
	return nil
}

func (f *failureLeader) Inject(clus *Cluster) error {
	if err := f.updateIndex(clus); err != nil {
		return err
	}
	return f.injectMember(clus, f.last)
}

func (f *failureLeader) Recover(clus *Cluster) error {
	return f.recoverMember(clus, f.last)
}

func (f *failureLeader) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

type failureQuorum failureByFunc

func (f *failureQuorum) Inject(clus *Cluster) error {
	for i := range killMap(len(clus.Members), clus.rd) {
		if err := f.injectMember(clus, i); err != nil {
			return err
		}
	}
	return nil
}

func (f *failureQuorum) Recover(clus *Cluster) error {
	for i := range killMap(len(clus.Members), clus.rd) {
		if err := f.recoverMember(clus, i); err != nil {
			return err
		}
	}
	return nil
}

func (f *failureQuorum) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	return f.failureCase.String()
}

func (f *failureQuorum) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

func killMap(size int, seed int) map[int]bool {
	m := make(map[int]bool)
	r := rand.New(rand.NewSource(int64(seed)))
	majority := size/2 + 1
	for {
		m[r.Intn(size)] = true
		if len(m) >= majority {
			return m
		}
	}
}

type failureAll failureByFunc

func (f *failureAll) Inject(clus *Cluster) error {
	for i := range clus.Members {
		if err := f.injectMember(clus, i); err != nil {
			return err
		}
	}
	return nil
}

func (f *failureAll) Recover(clus *Cluster) error {
	for i := range clus.Members {
		if err := f.recoverMember(clus, i); err != nil {
			return err
		}
	}
	return nil
}

func (f *failureAll) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	return f.failureCase.String()
}

func (f *failureAll) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

// failureUntilSnapshot injects a failure and waits for a snapshot event
type failureUntilSnapshot struct {
	desc        string
	failureCase rpcpb.FailureCase
	Failure
}

var slowCases = map[rpcpb.FailureCase]bool{
	rpcpb.FailureCase_DELAY_PEER_PORT_TX_RX_ONE_FOLLOWER_UNTIL_TRIGGER_SNAPSHOT: true,
	rpcpb.FailureCase_DELAY_PEER_PORT_TX_RX_LEADER_UNTIL_TRIGGER_SNAPSHOT:       true,
}

func (f *failureUntilSnapshot) Inject(clus *Cluster) error {
	if err := f.Failure.Inject(clus); err != nil {
		return err
	}
	if len(clus.Members) < 3 {
		return nil
	}

	snapshotCount := clus.Members[0].Etcd.SnapshotCount

	now := time.Now()
	clus.lg.Info(
		"trigger snapshot START",
		zap.Int("round", clus.rd),
		zap.Int("case", clus.cs),
		zap.String("desc", f.Desc()),
		zap.Int64("etcd-snapshot-count", snapshotCount),
	)

	// maxRev may fail since failure just injected, retry if failed.
	startRev, err := clus.maxRev()
	for i := 0; i < 10 && startRev == 0; i++ {
		startRev, err = clus.maxRev()
	}
	if startRev == 0 {
		return err
	}
	lastRev := startRev

	// healthy cluster could accept 1000 req/sec at least.
	// 3x time to trigger snapshot.
	retries := int(snapshotCount) / 1000 * 3
	if v, ok := slowCases[f.FailureCase()]; v && ok {
		// slow network takes more retries
		retries *= 4
	}

	for i := 0; i < retries; i++ {
		lastRev, _ = clus.maxRev()
		// If the number of proposals committed is bigger than snapshot count,
		// a new snapshot should have been created.
		diff := lastRev - startRev
		if diff > snapshotCount {
			clus.lg.Info(
				"trigger snapshot PASS",
				zap.Int("round", clus.rd),
				zap.Int("case", clus.cs),
				zap.Int("retries", i),
				zap.String("desc", f.Desc()),
				zap.Int64("committed-entries", diff),
				zap.Int64("etcd-snapshot-count", snapshotCount),
				zap.Int64("last-revision", lastRev),
				zap.Duration("took", time.Since(now)),
			)
			return nil
		}

		clus.lg.Info(
			"trigger snapshot PROGRESS",
			zap.Int("retries", i),
			zap.Int64("committed-entries", diff),
			zap.Int64("etcd-snapshot-count", snapshotCount),
			zap.Int64("last-revision", lastRev),
			zap.Duration("took", time.Since(now)),
		)
		time.Sleep(time.Second)
	}

	return fmt.Errorf("cluster too slow: only %d commits in %d retries", lastRev-startRev, retries)
}

func (f *failureUntilSnapshot) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	if f.failureCase.String() != "" {
		return f.failureCase.String()
	}
	return f.Failure.Desc()
}

func (f *failureUntilSnapshot) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}
<commit_msg>functional-tester/tester: reduce retries in slow network<commit_after>// Copyright 2018 The etcd 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 tester

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

	"github.com/coreos/etcd/tools/functional-tester/rpcpb"

	"go.uber.org/zap"
)

// Failure defines failure injection interface.
// To add a fail case:
//  1. implement "Failure" interface
//  2. define fail case name in "rpcpb.FailureCase"
type Failure interface {
	// Inject injeccts the failure into the testing cluster at the given
	// round. When calling the function, the cluster should be in health.
	Inject(clus *Cluster) error
	// Recover recovers the injected failure caused by the injection of the
	// given round and wait for the recovery of the testing cluster.
	Recover(clus *Cluster) error
	// Desc returns a description of the failure
	Desc() string
	// FailureCase returns "rpcpb.FailureCase" enum type.
	FailureCase() rpcpb.FailureCase
}

type injectMemberFunc func(*Cluster, int) error
type recoverMemberFunc func(*Cluster, int) error

type failureByFunc struct {
	desc          string
	failureCase   rpcpb.FailureCase
	injectMember  injectMemberFunc
	recoverMember recoverMemberFunc
}

func (f *failureByFunc) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	return f.failureCase.String()
}

func (f *failureByFunc) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

type failureFollower struct {
	failureByFunc
	last int
	lead int
}

func (f *failureFollower) updateIndex(clus *Cluster) error {
	idx, err := clus.GetLeader()
	if err != nil {
		return err
	}
	f.lead = idx

	n := len(clus.Members)
	if f.last == -1 { // first run
		f.last = clus.rd % n
		if f.last == f.lead {
			f.last = (f.last + 1) % n
		}
	} else {
		f.last = (f.last + 1) % n
		if f.last == f.lead {
			f.last = (f.last + 1) % n
		}
	}
	return nil
}

func (f *failureFollower) Inject(clus *Cluster) error {
	if err := f.updateIndex(clus); err != nil {
		return err
	}
	return f.injectMember(clus, f.last)
}

func (f *failureFollower) Recover(clus *Cluster) error {
	return f.recoverMember(clus, f.last)
}

func (f *failureFollower) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	return f.failureCase.String()
}

func (f *failureFollower) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

type failureLeader struct {
	failureByFunc
	last int
	lead int
}

func (f *failureLeader) updateIndex(clus *Cluster) error {
	idx, err := clus.GetLeader()
	if err != nil {
		return err
	}
	f.lead = idx
	f.last = idx
	return nil
}

func (f *failureLeader) Inject(clus *Cluster) error {
	if err := f.updateIndex(clus); err != nil {
		return err
	}
	return f.injectMember(clus, f.last)
}

func (f *failureLeader) Recover(clus *Cluster) error {
	return f.recoverMember(clus, f.last)
}

func (f *failureLeader) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

type failureQuorum failureByFunc

func (f *failureQuorum) Inject(clus *Cluster) error {
	for i := range killMap(len(clus.Members), clus.rd) {
		if err := f.injectMember(clus, i); err != nil {
			return err
		}
	}
	return nil
}

func (f *failureQuorum) Recover(clus *Cluster) error {
	for i := range killMap(len(clus.Members), clus.rd) {
		if err := f.recoverMember(clus, i); err != nil {
			return err
		}
	}
	return nil
}

func (f *failureQuorum) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	return f.failureCase.String()
}

func (f *failureQuorum) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

func killMap(size int, seed int) map[int]bool {
	m := make(map[int]bool)
	r := rand.New(rand.NewSource(int64(seed)))
	majority := size/2 + 1
	for {
		m[r.Intn(size)] = true
		if len(m) >= majority {
			return m
		}
	}
}

type failureAll failureByFunc

func (f *failureAll) Inject(clus *Cluster) error {
	for i := range clus.Members {
		if err := f.injectMember(clus, i); err != nil {
			return err
		}
	}
	return nil
}

func (f *failureAll) Recover(clus *Cluster) error {
	for i := range clus.Members {
		if err := f.recoverMember(clus, i); err != nil {
			return err
		}
	}
	return nil
}

func (f *failureAll) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	return f.failureCase.String()
}

func (f *failureAll) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}

// failureUntilSnapshot injects a failure and waits for a snapshot event
type failureUntilSnapshot struct {
	desc        string
	failureCase rpcpb.FailureCase
	Failure
}

var slowCases = map[rpcpb.FailureCase]bool{
	rpcpb.FailureCase_DELAY_PEER_PORT_TX_RX_ONE_FOLLOWER_UNTIL_TRIGGER_SNAPSHOT: true,
	rpcpb.FailureCase_DELAY_PEER_PORT_TX_RX_LEADER_UNTIL_TRIGGER_SNAPSHOT:       true,
}

func (f *failureUntilSnapshot) Inject(clus *Cluster) error {
	if err := f.Failure.Inject(clus); err != nil {
		return err
	}
	if len(clus.Members) < 3 {
		return nil
	}

	snapshotCount := clus.Members[0].Etcd.SnapshotCount

	now := time.Now()
	clus.lg.Info(
		"trigger snapshot START",
		zap.Int("round", clus.rd),
		zap.Int("case", clus.cs),
		zap.String("desc", f.Desc()),
		zap.Int64("etcd-snapshot-count", snapshotCount),
	)

	// maxRev may fail since failure just injected, retry if failed.
	startRev, err := clus.maxRev()
	for i := 0; i < 10 && startRev == 0; i++ {
		startRev, err = clus.maxRev()
	}
	if startRev == 0 {
		return err
	}
	lastRev := startRev

	// healthy cluster could accept 1000 req/sec at least.
	// 3x time to trigger snapshot.
	retries := int(snapshotCount) / 1000 * 3
	if v, ok := slowCases[f.FailureCase()]; v && ok {
		// slow network takes more retries
		retries *= 2
	}

	for i := 0; i < retries; i++ {
		lastRev, _ = clus.maxRev()
		// If the number of proposals committed is bigger than snapshot count,
		// a new snapshot should have been created.
		diff := lastRev - startRev
		if diff > snapshotCount {
			clus.lg.Info(
				"trigger snapshot PASS",
				zap.Int("round", clus.rd),
				zap.Int("case", clus.cs),
				zap.Int("retries", i),
				zap.String("desc", f.Desc()),
				zap.Int64("committed-entries", diff),
				zap.Int64("etcd-snapshot-count", snapshotCount),
				zap.Int64("last-revision", lastRev),
				zap.Duration("took", time.Since(now)),
			)
			return nil
		}

		clus.lg.Info(
			"trigger snapshot PROGRESS",
			zap.Int("retries", i),
			zap.Int64("committed-entries", diff),
			zap.Int64("etcd-snapshot-count", snapshotCount),
			zap.Int64("last-revision", lastRev),
			zap.Duration("took", time.Since(now)),
		)
		time.Sleep(time.Second)
	}

	return fmt.Errorf("cluster too slow: only %d commits in %d retries", lastRev-startRev, retries)
}

func (f *failureUntilSnapshot) Desc() string {
	if f.desc != "" {
		return f.desc
	}
	if f.failureCase.String() != "" {
		return f.failureCase.String()
	}
	return f.Failure.Desc()
}

func (f *failureUntilSnapshot) FailureCase() rpcpb.FailureCase {
	return f.failureCase
}
<|endoftext|>


<commit_before>package middlewares

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/cozy/cozy-stack/pkg/config/config"
	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/labstack/echo/v4"
)

// BasicAuth use HTTP basic authentication to authenticate a user. The secret
// of the user should be stored in a file with the specified name, stored in
// one of the the config.Paths directories.
//
// The format of the secret is the same as our hashed passwords in database: a
// scrypt hash with a salt contained in the value.
func BasicAuth(secretFileName string) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			if c.QueryParam("Trace") == "true" {
				t := time.Now()
				defer func() {
					elapsed := time.Since(t)
					logger.
						WithDomain("admin").
						WithField("nspace", "trace").
						Printf("Check basic auth: %v", elapsed)
				}()
			}

			_, passphrase, ok := c.Request().BasicAuth()
			if !ok {
				return echo.NewHTTPError(http.StatusUnauthorized, "missing basic auth")
			}

			shadowFile, err := config.FindConfigFile(secretFileName)
			if err != nil {
				return echo.NewHTTPError(http.StatusInternalServerError, err)
			}

			f, err := os.Open(shadowFile)
			if err != nil {
				return echo.NewHTTPError(http.StatusInternalServerError, err)
			}
			defer f.Close()

			b, err := ioutil.ReadAll(f)
			if err != nil {
				return echo.NewHTTPError(http.StatusInternalServerError, err)
			}
			b = bytes.TrimSpace(b)

			needUpdate, err := crypto.CompareHashAndPassphrase(b, []byte(passphrase))
			if err != nil {
				return echo.NewHTTPError(http.StatusForbidden, "bad passphrase")
			}
			if needUpdate {
				logger.
					WithDomain("admin").
					Warnf("Passphrase hash from %q needs update and should be regenerated", secretFileName)
			}

			return next(c)
		}
	}
}
<commit_msg>Fix trace in BasicAuth middleware (#2646)<commit_after>package middlewares

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/cozy/cozy-stack/pkg/config/config"
	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/labstack/echo/v4"
)

// BasicAuth use HTTP basic authentication to authenticate a user. The secret
// of the user should be stored in a file with the specified name, stored in
// one of the the config.Paths directories.
//
// The format of the secret is the same as our hashed passwords in database: a
// scrypt hash with a salt contained in the value.
func BasicAuth(secretFileName string) echo.MiddlewareFunc {
	check := func(next echo.HandlerFunc, c echo.Context) error {
		if c.QueryParam("Trace") == "true" {
			t := time.Now()
			defer func() {
				elapsed := time.Since(t)
				logger.
					WithDomain("admin").
					WithField("nspace", "trace").
					Printf("Check basic auth: %v", elapsed)
			}()
		}

		_, passphrase, ok := c.Request().BasicAuth()
		if !ok {
			return echo.NewHTTPError(http.StatusUnauthorized, "missing basic auth")
		}

		shadowFile, err := config.FindConfigFile(secretFileName)
		if err != nil {
			return echo.NewHTTPError(http.StatusInternalServerError, err)
		}

		f, err := os.Open(shadowFile)
		if err != nil {
			return echo.NewHTTPError(http.StatusInternalServerError, err)
		}
		defer f.Close()

		b, err := ioutil.ReadAll(f)
		if err != nil {
			return echo.NewHTTPError(http.StatusInternalServerError, err)
		}
		b = bytes.TrimSpace(b)

		needUpdate, err := crypto.CompareHashAndPassphrase(b, []byte(passphrase))
		if err != nil {
			return echo.NewHTTPError(http.StatusForbidden, "bad passphrase")
		}
		if needUpdate {
			logger.
				WithDomain("admin").
				Warnf("Passphrase hash from %q needs update and should be regenerated", secretFileName)
		}

		return nil
	}

	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			if err := check(next, c); err != nil {
				return err
			}

			return next(c)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"

	"github.com/boltdb/bolt"
	ui "github.com/gosuri/uiprogress"
)

// A result is the product of reading and summing a file using MD5.
type result struct {
	path string
	sum  []byte
	err  error
}

// walkFiles starts a goroutine to walk the directory tree at root and send the
// path of each regular file on the string channel.  It sends the result of the
// walk on the error channel.  If done is closed, walkFiles abandons its work.
func walkFiles(done <-chan struct{}, root string) (<-chan string, <-chan error) {
	paths := make(chan string)
	errc := make(chan error, 1)
	go func() { // HL
		// Close the paths channel after Walk returns.
		defer close(paths) // HL
		// No select needed for this send, since errc is buffered.
		errc <- filepath.Walk(root, func(path string, info os.FileInfo, err error) error { // HL
			if err != nil {
				log.Println("eoor")
				return err
			}
			if !info.Mode().IsRegular() {
				return nil
			}

			select {
			case paths <- path: // HL
			case <-done: // HL
				return errors.New("walk canceled")
			}
			return nil
		})
	}()
	return paths, errc
}

// digester reads path names from paths and sends digests of the corresponding
// files on c until either paths or done is closed.
func digester(done <-chan struct{}, paths <-chan string, c chan<- result) {
	for path := range paths { // HLpaths
		result := result{path: path}
		f, err := os.Open(path)
		if err != nil {
			result.err = err
		} else {
			h := md5.New()
			if _, err := io.Copy(h, f); err != nil {
				result.err = err
			}
			result.sum = h.Sum(nil)
		}
		select {
		case c <- result:
		case <-done:
			return

		}

	}
}

// MD5All reads all the files in the file tree rooted at root and returns a map
// from file path to the MD5 sum of the file's contents.  If the directory walk
// fails or any read operation fails, MD5All returns an error.  In that case,
// MD5All does not wait for inflight read operations to complete.
func MD5All(root string, dbName string, bucketName string, stats *Statistics, bar *ui.Bar) error {

	db, err := bolt.Open(dbname, 0600, nil)
	if err != nil {
		return err
	}
	defer db.Close()

	// MD5All closes the done channel when it returns; it may do so before
	// receiving all the values from c and errc.
	done := make(chan struct{})
	defer close(done)

	paths, errc := walkFiles(done, root)

	// Start a fixed number of goroutines to read and digest files.
	c := make(chan result) // HLc
	var wg sync.WaitGroup

	wg.Add(digesterNr)
	log.Printf("launching %d digester", digesterNr)
	for i := 0; i < digesterNr; i++ {
		go func() {
			digester(done, paths, c) // HLc
			wg.Done()
		}()
	}
	go func() {
		wg.Wait()
		close(c) // HLc
	}()
	// End of pipeline. OMIT
	err = db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		e := b.Bucket([]byte("errors"))
		if b == nil {
			return fmt.Errorf("gettings bucket")
		}
		for r := range c { // HLrange
			bar.Incr()
			stats.Files++
			if r.err != nil {
				stats.Errors++
				if err := e.Put([]byte(r.path), []byte(r.err.Error())); err != nil {
					log.Println(err)
				}
			} else {
				if err := b.Put([]byte(r.path), r.sum[:]); err != nil {
					log.Println(err)
				}

			}

		}
		return nil
	})
	if err != nil {
		return err
	}
	if err := <-errc; err != nil {
		log.Println(err)
		return err
	}
	return nil
}
<commit_msg>shadow declaration<commit_after>package main

import (
	"crypto/md5"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"

	"github.com/boltdb/bolt"
	ui "github.com/gosuri/uiprogress"
)

// A result is the product of reading and summing a file using MD5.
type result struct {
	path string
	sum  []byte
	err  error
}

// walkFiles starts a goroutine to walk the directory tree at root and send the
// path of each regular file on the string channel.  It sends the result of the
// walk on the error channel.  If done is closed, walkFiles abandons its work.
func walkFiles(done <-chan struct{}, root string) (<-chan string, <-chan error) {
	paths := make(chan string)
	errc := make(chan error, 1)
	go func() { // HL
		// Close the paths channel after Walk returns.
		defer close(paths) // HL
		// No select needed for this send, since errc is buffered.
		errc <- filepath.Walk(root, func(path string, info os.FileInfo, err error) error { // HL
			if err != nil {
				log.Println("eoor")
				return err
			}
			if !info.Mode().IsRegular() {
				return nil
			}

			select {
			case paths <- path: // HL
			case <-done: // HL
				return errors.New("walk canceled")
			}
			return nil
		})
	}()
	return paths, errc
}

// digester reads path names from paths and sends digests of the corresponding
// files on c until either paths or done is closed.
func digester(done <-chan struct{}, paths <-chan string, c chan<- result) {
	for path := range paths { // HLpaths
		res := result{path: path}
		f, err := os.Open(path)
		if err != nil {
			res.err = err
		} else {
			h := md5.New()
			if _, err := io.Copy(h, f); err != nil {
				res.err = err
			}
			res.sum = h.Sum(nil)
		}
		select {
		case c <- res:
		case <-done:
			return

		}

	}
}

// MD5All reads all the files in the file tree rooted at root and returns a map
// from file path to the MD5 sum of the file's contents.  If the directory walk
// fails or any read operation fails, MD5All returns an error.  In that case,
// MD5All does not wait for inflight read operations to complete.
func MD5All(root string, dbName string, bucketName string, stats *Statistics, bar *ui.Bar) error {

	db, err := bolt.Open(dbname, 0600, nil)
	if err != nil {
		return err
	}
	defer db.Close()

	// MD5All closes the done channel when it returns; it may do so before
	// receiving all the values from c and errc.
	done := make(chan struct{})
	defer close(done)

	paths, errc := walkFiles(done, root)

	// Start a fixed number of goroutines to read and digest files.
	c := make(chan result) // HLc
	var wg sync.WaitGroup

	wg.Add(digesterNr)
	log.Printf("launching %d digester", digesterNr)
	for i := 0; i < digesterNr; i++ {
		go func() {
			digester(done, paths, c) // HLc
			wg.Done()
		}()
	}
	go func() {
		wg.Wait()
		close(c) // HLc
	}()
	// End of pipeline. OMIT
	err = db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		e := b.Bucket([]byte("errors"))
		if b == nil {
			return fmt.Errorf("gettings bucket")
		}
		for r := range c { // HLrange
			bar.Incr()
			stats.Files++
			if r.err != nil {
				stats.Errors++
				if err = e.Put([]byte(r.path), []byte(r.err.Error())); err != nil {
					log.Println(err)
				}
			} else {
				if err = b.Put([]byte(r.path), r.sum[:]); err != nil {
					log.Println(err)
				}

			}

		}
		return nil
	})
	if err != nil {
		return err
	}
	if err := <-errc; err != nil {
		log.Println(err)
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Runs the servicerunner binary and checks that it outputs a JSON line to
// stdout with the expected variables.
package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"testing"
)

func check(t *testing.T, err error) {
	if err != nil {
		t.Fatal(err)
	}
}

func TestMain(t *testing.T) {
	tmpdir, err := ioutil.TempDir("", "servicerunner_test")
	check(t, err)
	defer os.RemoveAll(tmpdir)
	os.Setenv("TMPDIR", tmpdir)

	bin := path.Join(tmpdir, "servicerunner")
	fmt.Println("Building", bin)
	check(t, exec.Command("veyron", "go", "build", "-o", bin, "veyron.io/veyron/veyron/tools/servicerunner").Run())

	cmd := exec.Command(bin)
	stdout, err := cmd.StdoutPipe()
	check(t, err)
	check(t, cmd.Start())

	line, err := bufio.NewReader(stdout).ReadBytes('\n')
	check(t, err)
	vars := map[string]string{}
	check(t, json.Unmarshal(line, &vars))
	fmt.Println(vars)
	for _, name := range []string{"VEYRON_CREDENTIALS", "MT_NAME", "PROXY_ADDR", "WSPR_ADDR"} {
		if _, ok := vars[name]; !ok {
			t.Error("Missing", name)
		}
	}

	check(t, cmd.Process.Kill())
}
<commit_msg>s/veyron/v23/ for the service-runner test.<commit_after>// Runs the servicerunner binary and checks that it outputs a JSON line to
// stdout with the expected variables.
package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"testing"
)

func check(t *testing.T, err error) {
	if err != nil {
		t.Fatal(err)
	}
}

func TestMain(t *testing.T) {
	tmpdir, err := ioutil.TempDir("", "servicerunner_test")
	check(t, err)
	defer os.RemoveAll(tmpdir)
	os.Setenv("TMPDIR", tmpdir)

	bin := path.Join(tmpdir, "servicerunner")
	fmt.Println("Building", bin)
	check(t, exec.Command("v23", "go", "build", "-o", bin, "veyron.io/veyron/veyron/tools/servicerunner").Run())

	cmd := exec.Command(bin)
	stdout, err := cmd.StdoutPipe()
	check(t, err)
	check(t, cmd.Start())

	line, err := bufio.NewReader(stdout).ReadBytes('\n')
	check(t, err)
	vars := map[string]string{}
	check(t, json.Unmarshal(line, &vars))
	fmt.Println(vars)
	for _, name := range []string{"VEYRON_CREDENTIALS", "MT_NAME", "PROXY_ADDR", "WSPR_ADDR"} {
		if _, ok := vars[name]; !ok {
			t.Error("Missing", name)
		}
	}

	check(t, cmd.Process.Kill())
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/workspaces"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"

	"github.com/terraform-providers/terraform-provider-aws/aws/internal/keyvaluetags"
)

func dataSourceAwsWorkspacesWorkspace() *schema.Resource {
	return &schema.Resource{
		Read: dataSourceAwsWorkspacesWorkspaceRead,

		Schema: map[string]*schema.Schema{
			"bundle_id": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"directory_id": {
				Type:         schema.TypeString,
				Optional:     true,
				RequiredWith: []string{"user_name"},
			},
			"ip_address": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"computer_name": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"state": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"root_volume_encryption_enabled": {
				Type:     schema.TypeBool,
				Computed: true,
			},
			"user_name": {
				Type:         schema.TypeString,
				Optional:     true,
				Computed:     true,
				RequiredWith: []string{"directory_id"},
			},
			"user_volume_encryption_enabled": {
				Type:     schema.TypeBool,
				Computed: true,
			},
			"volume_encryption_key": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"workspace_id": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"directory_id", "user_name"},
			},
			"workspace_properties": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"compute_type_name": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"root_volume_size_gib": {
							Type:     schema.TypeInt,
							Computed: true,
						},
						"running_mode": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"running_mode_auto_stop_timeout_in_minutes": {
							Type:     schema.TypeInt,
							Computed: true,
						},
						"user_volume_size_gib": {
							Type:     schema.TypeInt,
							Computed: true,
						},
					},
				},
			},
			"tags": tagsSchema(),
		},
	}
}

func dataSourceAwsWorkspacesWorkspaceRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn
	ignoreTagsConfig := meta.(*AWSClient).IgnoreTagsConfig

	var workspace *workspaces.Workspace

	if workspaceID, ok := d.GetOk("workspace_id"); ok {
		resp, err := conn.DescribeWorkspaces(&workspaces.DescribeWorkspacesInput{
			WorkspaceIds: aws.StringSlice([]string{workspaceID.(string)}),
		})
		if err != nil {
			return err
		}

		if len(resp.Workspaces) != 1 {
			return fmt.Errorf("expected 1 result for WorkSpace  %q, found %d", workspaceID, len(resp.Workspaces))
		}

		workspace = resp.Workspaces[0]

		if workspace == nil {
			return fmt.Errorf("no WorkSpace with ID %q found", workspaceID)
		}
	}

	if directoryID, ok := d.GetOk("directory_id"); ok {
		userName := d.Get("user_name").(string)
		resp, err := conn.DescribeWorkspaces(&workspaces.DescribeWorkspacesInput{
			DirectoryId: aws.String(directoryID.(string)),
			UserName:    aws.String(userName),
		})
		if err != nil {
			return err
		}

		if len(resp.Workspaces) != 1 {
			return fmt.Errorf("expected 1 result for %q WorkSpace in %q directory, found %d", userName, directoryID, len(resp.Workspaces))
		}

		workspace = resp.Workspaces[0]

		if workspace == nil {
			return fmt.Errorf("no %q WorkSpace in %q directory found", userName, directoryID)
		}
	}

	d.SetId(aws.StringValue(workspace.WorkspaceId))
	d.Set("bundle_id", aws.StringValue(workspace.BundleId))
	d.Set("directory_id", aws.StringValue(workspace.DirectoryId))
	d.Set("ip_address", aws.StringValue(workspace.IpAddress))
	d.Set("computer_name", aws.StringValue(workspace.ComputerName))
	d.Set("state", aws.StringValue(workspace.State))
	d.Set("root_volume_encryption_enabled", aws.BoolValue(workspace.RootVolumeEncryptionEnabled))
	d.Set("user_name", aws.StringValue(workspace.UserName))
	d.Set("user_volume_encryption_enabled", aws.BoolValue(workspace.UserVolumeEncryptionEnabled))
	d.Set("volume_encryption_key", aws.StringValue(workspace.VolumeEncryptionKey))
	if err := d.Set("workspace_properties", flattenWorkspaceProperties(workspace.WorkspaceProperties)); err != nil {
		return fmt.Errorf("error setting workspace properties: %s", err)
	}

	tags, err := keyvaluetags.WorkspacesListTags(conn, d.Id())
	if err != nil {
		return fmt.Errorf("error listing tags: %s", err)
	}

	if err := d.Set("tags", tags.IgnoreAws().IgnoreConfig(ignoreTagsConfig).Map()); err != nil {
		return fmt.Errorf("error setting tags: %s", err)
	}

	return nil
}
<commit_msg>@review<commit_after>package aws

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/workspaces"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"

	"github.com/terraform-providers/terraform-provider-aws/aws/internal/keyvaluetags"
)

func dataSourceAwsWorkspacesWorkspace() *schema.Resource {
	return &schema.Resource{
		Read: dataSourceAwsWorkspacesWorkspaceRead,

		Schema: map[string]*schema.Schema{
			"bundle_id": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"directory_id": {
				Type:          schema.TypeString,
				Computed:      true,
				Optional:      true,
				RequiredWith:  []string{"user_name"},
				ConflictsWith: []string{"workspace_id"},
			},
			"ip_address": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"computer_name": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"state": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"root_volume_encryption_enabled": {
				Type:     schema.TypeBool,
				Computed: true,
			},
			"user_name": {
				Type:          schema.TypeString,
				Computed:      true,
				Optional:      true,
				RequiredWith:  []string{"directory_id"},
				ConflictsWith: []string{"workspace_id"},
			},
			"user_volume_encryption_enabled": {
				Type:     schema.TypeBool,
				Computed: true,
			},
			"volume_encryption_key": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"workspace_id": {
				Type:          schema.TypeString,
				Computed:      true,
				Optional:      true,
				ConflictsWith: []string{"directory_id", "user_name"},
			},
			"workspace_properties": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"compute_type_name": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"root_volume_size_gib": {
							Type:     schema.TypeInt,
							Computed: true,
						},
						"running_mode": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"running_mode_auto_stop_timeout_in_minutes": {
							Type:     schema.TypeInt,
							Computed: true,
						},
						"user_volume_size_gib": {
							Type:     schema.TypeInt,
							Computed: true,
						},
					},
				},
			},
			"tags": tagsSchemaComputed(),
		},
	}
}

func dataSourceAwsWorkspacesWorkspaceRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn
	ignoreTagsConfig := meta.(*AWSClient).IgnoreTagsConfig

	var workspace *workspaces.Workspace

	if workspaceID, ok := d.GetOk("workspace_id"); ok {
		resp, err := conn.DescribeWorkspaces(&workspaces.DescribeWorkspacesInput{
			WorkspaceIds: aws.StringSlice([]string{workspaceID.(string)}),
		})
		if err != nil {
			return err
		}

		if len(resp.Workspaces) != 1 {
			return fmt.Errorf("expected 1 result for WorkSpace  %q, found %d", workspaceID, len(resp.Workspaces))
		}

		workspace = resp.Workspaces[0]

		if workspace == nil {
			return fmt.Errorf("no WorkSpace with ID %q found", workspaceID)
		}
	}

	if directoryID, ok := d.GetOk("directory_id"); ok {
		userName := d.Get("user_name").(string)
		resp, err := conn.DescribeWorkspaces(&workspaces.DescribeWorkspacesInput{
			DirectoryId: aws.String(directoryID.(string)),
			UserName:    aws.String(userName),
		})
		if err != nil {
			return err
		}

		if len(resp.Workspaces) != 1 {
			return fmt.Errorf("expected 1 result for %q WorkSpace in %q directory, found %d", userName, directoryID, len(resp.Workspaces))
		}

		workspace = resp.Workspaces[0]

		if workspace == nil {
			return fmt.Errorf("no %q WorkSpace in %q directory found", userName, directoryID)
		}
	}

	d.SetId(aws.StringValue(workspace.WorkspaceId))
	d.Set("bundle_id", aws.StringValue(workspace.BundleId))
	d.Set("directory_id", aws.StringValue(workspace.DirectoryId))
	d.Set("ip_address", aws.StringValue(workspace.IpAddress))
	d.Set("computer_name", aws.StringValue(workspace.ComputerName))
	d.Set("state", aws.StringValue(workspace.State))
	d.Set("root_volume_encryption_enabled", aws.BoolValue(workspace.RootVolumeEncryptionEnabled))
	d.Set("user_name", aws.StringValue(workspace.UserName))
	d.Set("user_volume_encryption_enabled", aws.BoolValue(workspace.UserVolumeEncryptionEnabled))
	d.Set("volume_encryption_key", aws.StringValue(workspace.VolumeEncryptionKey))
	if err := d.Set("workspace_properties", flattenWorkspaceProperties(workspace.WorkspaceProperties)); err != nil {
		return fmt.Errorf("error setting workspace properties: %s", err)
	}

	tags, err := keyvaluetags.WorkspacesListTags(conn, d.Id())
	if err != nil {
		return fmt.Errorf("error listing tags: %s", err)
	}

	if err := d.Set("tags", tags.IgnoreAws().IgnoreConfig(ignoreTagsConfig).Map()); err != nil {
		return fmt.Errorf("error setting tags: %s", err)
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package uniter

import (
	"errors"
)

// This module implements a subset of the interface provided by
// state.RelationUnit, as needed by the uniter API.
// Most of this is pretty much a verbatim copy of the code in
// state/relationunit.go, except for a few API-specific changes.

// TODO: Only the calls that do not call the API are implemented,
// without tests for now. Others are added as placeholders,
// the actual implementation will come in a follow-up.

// RelationUnit holds information about a single unit in a relation,
// and allows clients to conveniently access unit-specific
// functionality.
type RelationUnit struct {
	st       *State
	relation *Relation
	unit     *Unit
	endpoint Endpoint
	scope    string
}

// Relation returns the relation associated with the unit.
func (ru *RelationUnit) Relation() *Relation {
	return ru.relation
}

// Endpoint returns the relation endpoint that defines the unit's
// participation in the relation.
func (ru *RelationUnit) Endpoint() Endpoint {
	return ru.endpoint
}

// PrivateAddress returns the private address of the unit and whether
// it is valid.
//
// NOTE: This differs from state.RelationUnit.PrivateAddress() by
// returning an error as well, because it needs to make an API call.
func (ru *RelationUnit) PrivateAddress() (string, bool, error) {
	return ru.unit.PrivateAddress()
}

// TODO: Translate ErrCannotEnterScope and ErrCannotEnterScopeYet
// errors from EnterScope() into their api/params equivalents with the
// respective codes ( CodeCannotEnterScope and
// CodeCannotEnterScopeYet) and message contents, matching the part of
// the error message after the colon.

// EnterScope ensures that the unit has entered its scope in the relation.
// When the unit has already entered its relation scope, EnterScope will report
// success but make no changes to state.
//
// Otherwise, assuming both the relation and the unit are alive, it will enter
// scope.
//
// If the unit is a principal and the relation has container scope, EnterScope
// will also create the required subordinate unit, if it does not already exist;
// this is because there's no point having a principal in scope if there is no
// corresponding subordinate to join it.
//
// Once a unit has entered a scope, it stays in scope without further
// intervention; the relation will not be able to become Dead until all units
// have departed its scopes.
//
// NOTE: Unlike state.RelatioUnit.EnterScope(), this method does not take
// settings, because uniter only uses this to supply the unit's private
// address, but this is not done at the server-side by the API.
func (ru *RelationUnit) EnterScope() error {
	// TODO: Call Uniter.EnterScope(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// LeaveScope signals that the unit has left its scope in the relation.
// After the unit has left its relation scope, it is no longer a member
// of the relation; if the relation is dying when its last member unit
// leaves, it is removed immediately. It is not an error to leave a scope
// that the unit is not, or never was, a member of.
func (ru *RelationUnit) LeaveScope() error {
	// TODO: Call Uniter.LeaveScope(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// Settings returns a Settings which allows access to the unit's settings
// within the relation.
func (ru *RelationUnit) Settings() (*Settings, error) {
	// TODO: Call Uniter.ReadSettings(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// ReadSettings returns a map holding the settings of the unit with the
// supplied name within this relation. An error will be returned if the
// relation no longer exists, or if the unit's service is not part of the
// relation, or the settings are invalid; but mere non-existence of the
// unit is not grounds for an error, because the unit settings are
// guaranteed to persist for the lifetime of the relation, regardless
// of the lifetime of the unit.
func (ru *RelationUnit) ReadSettings(uname string) (m map[string]interface{}, err error) {
	// TODO: Convert uname to a remoteUnitTag, then call
	// Uniter.ReadRemoteSettings(), passing ru.relation.Tag(),
	// ru.unit.Tag() (as local unit), and remoteUnitTag as arguments.
	panic("not implemented")
}

// Watch returns a watcher that notifies of changes to conterpart units in
// the relation.
// TODO: This is commented out because RelationUnitsWatcher is not yet
// implemented at the server-side.
//func (ru *RelationUnit) Watch() *RelationUnitsWatcher {
//}
<commit_msg>Changes after review<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package uniter

// This module implements a subset of the interface provided by
// state.RelationUnit, as needed by the uniter API.
// Most of this is pretty much a verbatim copy of the code in
// state/relationunit.go, except for a few API-specific changes.

// TODO: Only the calls that do not call the API are implemented,
// without tests for now. Others are added as placeholders,
// the actual implementation will come in a follow-up.

// RelationUnit holds information about a single unit in a relation,
// and allows clients to conveniently access unit-specific
// functionality.
type RelationUnit struct {
	st       *State
	relation *Relation
	unit     *Unit
	endpoint Endpoint
	scope    string
}

// Relation returns the relation associated with the unit.
func (ru *RelationUnit) Relation() *Relation {
	return ru.relation
}

// Endpoint returns the relation endpoint that defines the unit's
// participation in the relation.
func (ru *RelationUnit) Endpoint() Endpoint {
	return ru.endpoint
}

// PrivateAddress returns the private address of the unit and whether
// it is valid.
//
// NOTE: This differs from state.RelationUnit.PrivateAddress() by
// returning an error as well, because it needs to make an API call.
func (ru *RelationUnit) PrivateAddress() (string, bool, error) {
	return ru.unit.PrivateAddress()
}

// TODO: Translate ErrCannotEnterScope and ErrCannotEnterScopeYet
// errors from EnterScope() into their api/params equivalents with the
// respective codes ( CodeCannotEnterScope and
// CodeCannotEnterScopeYet) and message contents, matching the part of
// the error message after the colon.

// EnterScope ensures that the unit has entered its scope in the relation.
// When the unit has already entered its relation scope, EnterScope will report
// success but make no changes to state.
//
// Otherwise, assuming both the relation and the unit are alive, it will enter
// scope.
//
// If the unit is a principal and the relation has container scope, EnterScope
// will also create the required subordinate unit, if it does not already exist;
// this is because there's no point having a principal in scope if there is no
// corresponding subordinate to join it.
//
// Once a unit has entered a scope, it stays in scope without further
// intervention; the relation will not be able to become Dead until all units
// have departed its scopes.
//
// NOTE: Unlike state.RelatioUnit.EnterScope(), this method does not take
// settings, because uniter only uses this to supply the unit's private
// address, but this is not done at the server-side by the API.
func (ru *RelationUnit) EnterScope() error {
	// TODO: Call Uniter.EnterScope(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// LeaveScope signals that the unit has left its scope in the relation.
// After the unit has left its relation scope, it is no longer a member
// of the relation; if the relation is dying when its last member unit
// leaves, it is removed immediately. It is not an error to leave a scope
// that the unit is not, or never was, a member of.
func (ru *RelationUnit) LeaveScope() error {
	// TODO: Call Uniter.LeaveScope(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// Settings returns a Settings which allows access to the unit's settings
// within the relation.
func (ru *RelationUnit) Settings() (*Settings, error) {
	// TODO: Call Uniter.ReadSettings(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// ReadSettings returns a map holding the settings of the unit with the
// supplied name within this relation. An error will be returned if the
// relation no longer exists, or if the unit's service is not part of the
// relation, or the settings are invalid; but mere non-existence of the
// unit is not grounds for an error, because the unit settings are
// guaranteed to persist for the lifetime of the relation, regardless
// of the lifetime of the unit.
func (ru *RelationUnit) ReadSettings(uname string) (m map[string]interface{}, err error) {
	// TODO: Convert uname to a remoteUnitTag, then call
	// Uniter.ReadRemoteSettings(), passing ru.relation.Tag(),
	// ru.unit.Tag() (as local unit), and remoteUnitTag as arguments.
	panic("not implemented")
}

// Watch returns a watcher that notifies of changes to conterpart units in
// the relation.
// TODO: This is commented out because RelationUnitsWatcher is not yet
// implemented at the server-side.
//func (ru *RelationUnit) Watch() *RelationUnitsWatcher {
//}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"reflect"

	"github.com/juju/utils/set"
	gc "gopkg.in/check.v1"
	"gopkg.in/juju/charm.v6-unstable"
)

type MigrationSuite struct{}

var _ = gc.Suite(&MigrationSuite{})

func (s *MigrationSuite) TestKnownCollections(c *gc.C) {
	completedCollections := set.NewStrings(
		annotationsC,
		modelsC,
		modelUsersC,
		modelUserLastConnectionC,
		settingsC,
		statusesC,

		// machine
		instanceDataC,
		machinesC,
		openedPortsC,

		// service / unit
		servicesC,
		unitsC,
		meterStatusC, // red / green status for metrics of units

		// settings reference counts are only used for services
		settingsrefsC,

		// relation
		relationsC,
		relationScopesC,
	)

	ignoredCollections := set.NewStrings(
		// We don't export the controller model at this stage.
		controllersC,
		// Users aren't migrated.
		usersC,
		userLastLoginC,
		// userenvnameC is just to provide a unique key constraint.
		usermodelnameC,
		// Metrics aren't migrated.
		metricsC,
		// leaseC is deprecated in favour of leasesC.
		leaseC,
		// Backup and restore information is not migrated.
		restoreInfoC,
		// upgradeInfoC is used to coordinate upgrades and schema migrations,
		// and aren't needed for model migrations.
		upgradeInfoC,
		// Not exported, but the tools will possibly need to be either bundled
		// with the representation or sent separately.
		toolsmetadataC,
		// Transaction stuff.
		"txns",
		"txns.log",

		// We don't import any of the migration collections.
		modelMigrationsC,
		modelMigrationStatusC,
		modelMigrationsActiveC,

		// The container ref document is primarily there to keep track
		// of a particular machine's containers. The migration format
		// uses object containment for this purpose.
		containerRefsC,
		// The min units collection is only used to trigger a watcher
		// in order to have the service add or remove units if the minimum
		// number of units is changed. The Service doc has all we need
		// for migratino.
		minUnitsC,
		// This is a transitory collection of units that need to be assigned
		// to machines.
		assignUnitC,
	)

	// THIS SET WILL BE REMOVED WHEN MIGRATIONS ARE COMPLETE
	todoCollections := set.NewStrings(
		// model
		blocksC,
		cleanupsC,
		cloudimagemetadataC,
		sequenceC,

		// machine
		rebootC,

		// service / unit
		charmsC,
		leasesC,
		"payloads",
		"resources",
		endpointBindingsC,

		// storage
		blockDevicesC,
		filesystemsC,
		filesystemAttachmentsC,
		storageInstancesC,
		storageAttachmentsC,
		storageConstraintsC,
		volumesC,
		volumeAttachmentsC,

		// network
		ipaddressesC,
		networksC,
		networkInterfacesC,
		requestedNetworksC,
		subnetsC,
		spacesC,

		// actions
		actionsC,
		actionNotificationsC,
		actionresultsC,

		// done as part of machines/services/units
		constraintsC,
		statusesHistoryC,

		// uncategorised
		metricsManagerC, // should really be copied across
	)

	envCollections := set.NewStrings()
	for name := range allCollections() {
		envCollections.Add(name)
	}

	known := completedCollections.Union(ignoredCollections)

	remainder := envCollections.Difference(known)
	remainder = remainder.Difference(todoCollections)

	// If this test fails, it means that a new collection has been added
	// but migrations for it has not been done. This is a Bad Thing™.
	c.Assert(remainder, gc.HasLen, 0)
}

func (s *MigrationSuite) TestModelDocFields(c *gc.C) {
	fields := set.NewStrings(
		// UUID and Mame are constructed from the model config.
		"UUID",
		"Name",
		// Life will always be alive, or we won't be migrating.
		"Life",
		"Owner",
		"LatestAvailableTools",
		// ServerUUID is recreated when the new model is created in the
		// new controller (yay name changes).
		"ServerUUID",
		// Both of the times for dying and death are empty as the model
		// is alive.
		"TimeOfDying",
		"TimeOfDeath",
	)
	s.AssertExportedFields(c, modelDoc{}, fields)
}

func (s *MigrationSuite) TestEnvUserDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ID is the same as UserName (but lowercased)
		"ID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Tracked fields:
		"UserName",
		"DisplayName",
		"CreatedBy",
		"DateCreated",
		"ReadOnly",
	)
	s.AssertExportedFields(c, modelUserDoc{}, fields)
}

func (s *MigrationSuite) TestEnvUserLastConnectionDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ID is the same as UserName (but lowercased)
		"ID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// UserName is captured in the migration.User.
		"UserName",
		"LastConnection",
	)
	s.AssertExportedFields(c, modelUserLastConnectionDoc{}, fields)
}

func (s *MigrationSuite) TestMachineDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + machine id
		"DocID",
		// ID is the machine id
		"Id",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Life is always alive, confirmed by export precheck.
		"Life",

		"Addresses",
		"ContainerType",
		"Jobs",
		"MachineAddresses",
		"Nonce",
		"PasswordHash",
		"Placement",
		"PreferredPrivateAddress",
		"PreferredPublicAddress",
		"Series",
		"SupportedContainers",
		"SupportedContainersKnown",
		"Tools",

		// Ignored at this stage, could be an issue if mongo 3.0 isn't
		// available.
		"StopMongoUntilVersion",
	)
	todo := set.NewStrings(
		"Principals",
		"Volumes",
		"NoVote",
		"Clean",
		"Filesystems",
		"HasVote",
	)
	s.AssertExportedFields(c, machineDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestInstanceDataFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + machine id
		"DocID",
		"MachineId",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",

		"InstanceId",
		"Status",
		"Arch",
		"Mem",
		"RootDisk",
		"CpuCores",
		"CpuPower",
		"Tags",
		"AvailZone",
	)
	s.AssertExportedFields(c, instanceData{}, fields)
}

func (s *MigrationSuite) TestServiceDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + name
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Always alive, not explicitly exported.
		"Life",
		// OwnerTag is deprecated and should be deleted.
		"OwnerTag",
		// TxnRevno is mgo internals and should not be migrated.
		"TxnRevno",

		"Name",
		"Series",
		"Subordinate",
		"CharmURL",
		"ForceCharm",
		"Exposed",
		"MinUnits",
		"MetricCredentials",
		// UnitCount is handled by the number of units for the exported service.
		"UnitCount",
	)
	todo := set.NewStrings(
		"RelationCount",
	)
	s.AssertExportedFields(c, serviceDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestSettingsRefsDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",

		"RefCount",
	)
	s.AssertExportedFields(c, settingsRefsDoc{}, fields)
}

func (s *MigrationSuite) TestUnitDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		"Name",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Service is implicit in the migration structure through containment.
		"Service",
		// Series and CharmURL also come from the service.
		"Series",
		"CharmURL",
		"Principal",
		"Subordinates",
		"MachineId",
		// Resolved is not migrated as we check that all is good before we start.
		"Resolved",
		"Tools",
		// Life isn't migrated as we only migrate live things.
		"Life",
		// TxnRevno isn't migrated.
		"TxnRevno",
		"PasswordHash",
		// Obsolete and not migrated.
		"Ports",
		"PublicAddress",
		"PrivateAddress",
	)
	todo := set.NewStrings(
		"StorageAttachmentCount",
	)

	s.AssertExportedFields(c, unitDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestPortsDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// MachineId is implicit in the migration structure through containment.
		"MachineID",
		"NetworkName",
		"Ports",
		// TxnRevno isn't migrated.
		"TxnRevno",
	)
	s.AssertExportedFields(c, portsDoc{}, fields)
}

func (s *MigrationSuite) TestMeterStatusDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		"Code",
		"Info",
	)
	s.AssertExportedFields(c, meterStatusDoc{}, fields)
}

func (s *MigrationSuite) TestRelationDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		"Key",
		"Id",
		"Endpoints",
		// Life isn't exported, only alive.
		"Life",
		// UnitCount isn't explicitly exported, but defined by the stored
		// unit settings data for the relation endpoint.
		"UnitCount",
	)
	s.AssertExportedFields(c, relationDoc{}, fields)
	// We also need to check the Endpoint and nested charm.Relation field.
	endpointFields := set.NewStrings("ServiceName", "Relation")
	s.AssertExportedFields(c, Endpoint{}, endpointFields)
	charmRelationFields := set.NewStrings(
		"Name",
		"Role",
		"Interface",
		"Optional",
		"Limit",
		"Scope",
	)
	s.AssertExportedFields(c, charm.Relation{}, charmRelationFields)
}

func (s *MigrationSuite) TestRelationScopeDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		"Key",
		// Departing isn't exported as we only deal with live, stable systems.
		"Departing",
	)
	s.AssertExportedFields(c, relationScopeDoc{}, fields)
}

func (s *MigrationSuite) TestAnnatatorDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		"GlobalKey",
		"Tag",
		"Annotations",
	)
	s.AssertExportedFields(c, annotatorDoc{}, fields)
}

func (s *MigrationSuite) AssertExportedFields(c *gc.C, doc interface{}, fields set.Strings) {
	expected := getExportedFields(doc)
	unknown := expected.Difference(fields)
	// If this test fails, it means that extra fields have been added to the
	// doc without thinking about the migration implications.
	c.Assert(unknown, gc.HasLen, 0)
}

func getExportedFields(arg interface{}) set.Strings {
	t := reflect.TypeOf(arg)
	result := set.NewStrings()

	count := t.NumField()
	for i := 0; i < count; i++ {
		f := t.Field(i)
		// empty PkgPath means exported field.
		// see https://golang.org/pkg/reflect/#StructField
		if f.PkgPath == "" {
			result.Add(f.Name)
		}
	}

	return result
}
<commit_msg>Mark sequences as done.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"reflect"

	"github.com/juju/utils/set"
	gc "gopkg.in/check.v1"
	"gopkg.in/juju/charm.v6-unstable"
)

type MigrationSuite struct{}

var _ = gc.Suite(&MigrationSuite{})

func (s *MigrationSuite) TestKnownCollections(c *gc.C) {
	completedCollections := set.NewStrings(
		annotationsC,
		modelsC,
		modelUsersC,
		modelUserLastConnectionC,
		settingsC,
		sequenceC,
		statusesC,

		// machine
		instanceDataC,
		machinesC,
		openedPortsC,

		// service / unit
		servicesC,
		unitsC,
		meterStatusC, // red / green status for metrics of units

		// settings reference counts are only used for services
		settingsrefsC,

		// relation
		relationsC,
		relationScopesC,
	)

	ignoredCollections := set.NewStrings(
		// We don't export the controller model at this stage.
		controllersC,
		// Users aren't migrated.
		usersC,
		userLastLoginC,
		// userenvnameC is just to provide a unique key constraint.
		usermodelnameC,
		// Metrics aren't migrated.
		metricsC,
		// leaseC is deprecated in favour of leasesC.
		leaseC,
		// Backup and restore information is not migrated.
		restoreInfoC,
		// upgradeInfoC is used to coordinate upgrades and schema migrations,
		// and aren't needed for model migrations.
		upgradeInfoC,
		// Not exported, but the tools will possibly need to be either bundled
		// with the representation or sent separately.
		toolsmetadataC,
		// Transaction stuff.
		"txns",
		"txns.log",

		// We don't import any of the migration collections.
		modelMigrationsC,
		modelMigrationStatusC,
		modelMigrationsActiveC,

		// The container ref document is primarily there to keep track
		// of a particular machine's containers. The migration format
		// uses object containment for this purpose.
		containerRefsC,
		// The min units collection is only used to trigger a watcher
		// in order to have the service add or remove units if the minimum
		// number of units is changed. The Service doc has all we need
		// for migratino.
		minUnitsC,
		// This is a transitory collection of units that need to be assigned
		// to machines.
		assignUnitC,
	)

	// THIS SET WILL BE REMOVED WHEN MIGRATIONS ARE COMPLETE
	todoCollections := set.NewStrings(
		// model
		blocksC,
		cleanupsC,
		cloudimagemetadataC,

		// machine
		rebootC,

		// service / unit
		charmsC,
		leasesC,
		"payloads",
		"resources",
		endpointBindingsC,

		// storage
		blockDevicesC,
		filesystemsC,
		filesystemAttachmentsC,
		storageInstancesC,
		storageAttachmentsC,
		storageConstraintsC,
		volumesC,
		volumeAttachmentsC,

		// network
		ipaddressesC,
		networksC,
		networkInterfacesC,
		requestedNetworksC,
		subnetsC,
		spacesC,

		// actions
		actionsC,
		actionNotificationsC,
		actionresultsC,

		// done as part of machines/services/units
		constraintsC,
		statusesHistoryC,

		// uncategorised
		metricsManagerC, // should really be copied across
	)

	envCollections := set.NewStrings()
	for name := range allCollections() {
		envCollections.Add(name)
	}

	known := completedCollections.Union(ignoredCollections)

	remainder := envCollections.Difference(known)
	remainder = remainder.Difference(todoCollections)

	// If this test fails, it means that a new collection has been added
	// but migrations for it has not been done. This is a Bad Thing™.
	c.Assert(remainder, gc.HasLen, 0)
}

func (s *MigrationSuite) TestModelDocFields(c *gc.C) {
	fields := set.NewStrings(
		// UUID and Mame are constructed from the model config.
		"UUID",
		"Name",
		// Life will always be alive, or we won't be migrating.
		"Life",
		"Owner",
		"LatestAvailableTools",
		// ServerUUID is recreated when the new model is created in the
		// new controller (yay name changes).
		"ServerUUID",
		// Both of the times for dying and death are empty as the model
		// is alive.
		"TimeOfDying",
		"TimeOfDeath",
	)
	s.AssertExportedFields(c, modelDoc{}, fields)
}

func (s *MigrationSuite) TestEnvUserDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ID is the same as UserName (but lowercased)
		"ID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Tracked fields:
		"UserName",
		"DisplayName",
		"CreatedBy",
		"DateCreated",
		"ReadOnly",
	)
	s.AssertExportedFields(c, modelUserDoc{}, fields)
}

func (s *MigrationSuite) TestEnvUserLastConnectionDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ID is the same as UserName (but lowercased)
		"ID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// UserName is captured in the migration.User.
		"UserName",
		"LastConnection",
	)
	s.AssertExportedFields(c, modelUserLastConnectionDoc{}, fields)
}

func (s *MigrationSuite) TestMachineDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + machine id
		"DocID",
		// ID is the machine id
		"Id",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Life is always alive, confirmed by export precheck.
		"Life",

		"Addresses",
		"ContainerType",
		"Jobs",
		"MachineAddresses",
		"Nonce",
		"PasswordHash",
		"Placement",
		"PreferredPrivateAddress",
		"PreferredPublicAddress",
		"Series",
		"SupportedContainers",
		"SupportedContainersKnown",
		"Tools",

		// Ignored at this stage, could be an issue if mongo 3.0 isn't
		// available.
		"StopMongoUntilVersion",
	)
	todo := set.NewStrings(
		"Principals",
		"Volumes",
		"NoVote",
		"Clean",
		"Filesystems",
		"HasVote",
	)
	s.AssertExportedFields(c, machineDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestInstanceDataFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + machine id
		"DocID",
		"MachineId",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",

		"InstanceId",
		"Status",
		"Arch",
		"Mem",
		"RootDisk",
		"CpuCores",
		"CpuPower",
		"Tags",
		"AvailZone",
	)
	s.AssertExportedFields(c, instanceData{}, fields)
}

func (s *MigrationSuite) TestServiceDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + name
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Always alive, not explicitly exported.
		"Life",
		// OwnerTag is deprecated and should be deleted.
		"OwnerTag",
		// TxnRevno is mgo internals and should not be migrated.
		"TxnRevno",

		"Name",
		"Series",
		"Subordinate",
		"CharmURL",
		"ForceCharm",
		"Exposed",
		"MinUnits",
		"MetricCredentials",
		// UnitCount is handled by the number of units for the exported service.
		"UnitCount",
	)
	todo := set.NewStrings(
		"RelationCount",
	)
	s.AssertExportedFields(c, serviceDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestSettingsRefsDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",

		"RefCount",
	)
	s.AssertExportedFields(c, settingsRefsDoc{}, fields)
}

func (s *MigrationSuite) TestUnitDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		"Name",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Service is implicit in the migration structure through containment.
		"Service",
		// Series and CharmURL also come from the service.
		"Series",
		"CharmURL",
		"Principal",
		"Subordinates",
		"MachineId",
		// Resolved is not migrated as we check that all is good before we start.
		"Resolved",
		"Tools",
		// Life isn't migrated as we only migrate live things.
		"Life",
		// TxnRevno isn't migrated.
		"TxnRevno",
		"PasswordHash",
		// Obsolete and not migrated.
		"Ports",
		"PublicAddress",
		"PrivateAddress",
	)
	todo := set.NewStrings(
		"StorageAttachmentCount",
	)

	s.AssertExportedFields(c, unitDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestPortsDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// MachineId is implicit in the migration structure through containment.
		"MachineID",
		"NetworkName",
		"Ports",
		// TxnRevno isn't migrated.
		"TxnRevno",
	)
	s.AssertExportedFields(c, portsDoc{}, fields)
}

func (s *MigrationSuite) TestMeterStatusDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		"Code",
		"Info",
	)
	s.AssertExportedFields(c, meterStatusDoc{}, fields)
}

func (s *MigrationSuite) TestRelationDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		"Key",
		"Id",
		"Endpoints",
		// Life isn't exported, only alive.
		"Life",
		// UnitCount isn't explicitly exported, but defined by the stored
		// unit settings data for the relation endpoint.
		"UnitCount",
	)
	s.AssertExportedFields(c, relationDoc{}, fields)
	// We also need to check the Endpoint and nested charm.Relation field.
	endpointFields := set.NewStrings("ServiceName", "Relation")
	s.AssertExportedFields(c, Endpoint{}, endpointFields)
	charmRelationFields := set.NewStrings(
		"Name",
		"Role",
		"Interface",
		"Optional",
		"Limit",
		"Scope",
	)
	s.AssertExportedFields(c, charm.Relation{}, charmRelationFields)
}

func (s *MigrationSuite) TestRelationScopeDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		"Key",
		// Departing isn't exported as we only deal with live, stable systems.
		"Departing",
	)
	s.AssertExportedFields(c, relationScopeDoc{}, fields)
}

func (s *MigrationSuite) TestAnnatatorDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		"GlobalKey",
		"Tag",
		"Annotations",
	)
	s.AssertExportedFields(c, annotatorDoc{}, fields)
}

func (s *MigrationSuite) TestSequenceDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		"DocID",
		"Name",
		"Counter",
	)
	s.AssertExportedFields(c, sequenceDoc{}, fields)
}

func (s *MigrationSuite) AssertExportedFields(c *gc.C, doc interface{}, fields set.Strings) {
	expected := getExportedFields(doc)
	unknown := expected.Difference(fields)
	// If this test fails, it means that extra fields have been added to the
	// doc without thinking about the migration implications.
	c.Assert(unknown, gc.HasLen, 0)
}

func getExportedFields(arg interface{}) set.Strings {
	t := reflect.TypeOf(arg)
	result := set.NewStrings()

	count := t.NumField()
	for i := 0; i < count; i++ {
		f := t.Field(i)
		// empty PkgPath means exported field.
		// see https://golang.org/pkg/reflect/#StructField
		if f.PkgPath == "" {
			result.Add(f.Name)
		}
	}

	return result
}
<|endoftext|>