<commit_before>// Copyright (c) 2016 Intel Corporation
//
// 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 main

import (
	"bytes"
	"encoding/json"
	"github.com/01org/ciao/payloads"
	"io/ioutil"
	"net/http"
	"testing"
)

func testCreateServer(t *testing.T, n int) payloads.ComputeServers {
	tenant, err := context.ds.GetTenant(computeTestUser)
	if err != nil {
		t.Fatal(err)
	}

	// get a valid workload ID
	wls, err := context.ds.GetWorkloads()
	if err != nil {
		t.Fatal(err)
	}

	if len(wls) == 0 {
		t.Fatal("No valid workloads")
	}

	url := computeURL + "/v2.1/" + tenant.ID + "/servers"

	var server payloads.ComputeCreateServer
	server.Server.MaxInstances = n
	server.Server.Workload = wls[0].ID

	b, err := json.Marshal(server)
	if err != nil {
		t.Fatal(err)
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(b))
	req.Header.Set("X-Auth-Token", "imavalidtoken")
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusAccepted {
		t.Fatalf("expected response code: %d, got %d", http.StatusAccepted, resp.StatusCode)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	var servers payloads.ComputeServers

	err = json.Unmarshal(body, &servers)
	if err != nil {
		t.Fatal(err)
	}

	if servers.TotalServers != n {
		t.Fatal("Not enough servers returned")
	}

	return servers
}

func testListServerDetailsTenant(t *testing.T, tenantID string) payloads.ComputeServers {
	url := computeURL + "/v2.1/" + tenantID + "/servers/detail"
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("X-Auth-Token", "imavalidtoken")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatalf("expected response code: %d, got %d", http.StatusOK, resp.StatusCode)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	var s payloads.ComputeServers
	err = json.Unmarshal(body, &s)
	if err != nil {
		t.Fatal(err)
	}

	return s
}

func TestCreateSingleServer(t *testing.T) {
	_ = testCreateServer(t, 1)
}

func TestListServerDetailsTenant(t *testing.T) {
	tenant, err := context.ds.GetTenant(computeTestUser)
	if err != nil {
		t.Fatal(err)
	}

	servers := testCreateServer(t, 1)
	if servers.TotalServers != 1 {
		t.Fatal(err)
	}

	s := testListServerDetailsTenant(t, tenant.ID)

	if s.TotalServers < 1 {
		t.Fatal("Not enough servers returned")
	}
}

func TestListServerDetailsWorkload(t *testing.T) {
	// get a valid workload ID
	wls, err := context.ds.GetWorkloads()
	if err != nil {
		t.Fatal(err)
	}

	if len(wls) == 0 {
		t.Fatal("No valid workloads")
	}

	servers := testCreateServer(t, 10)
	if servers.TotalServers != 10 {
		t.Fatal("failed to create enough servers")
	}

	url := computeURL + "/v2.1/flavors/" + wls[0].ID + "/servers/detail"
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("X-Auth-Token", "imavalidtoken")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := ioutil.ReadAll(resp.Body)
		t.Fatal(string(body))
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	var s payloads.ComputeServers
	err = json.Unmarshal(body, &s)
	if err != nil {
		t.Fatal(err)
	}

	if s.TotalServers < 10 {
		t.Fatal("Did not return correct number of servers")
	}
}
<commit_msg>ciao-controller: compute_test: add test for showServerDetails<commit_after>// Copyright (c) 2016 Intel Corporation
//
// 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 main

import (
	"bytes"
	"encoding/json"
	"github.com/01org/ciao/payloads"
	"io/ioutil"
	"net/http"
	"reflect"
	"testing"
)

func testCreateServer(t *testing.T, n int) payloads.ComputeServers {
	tenant, err := context.ds.GetTenant(computeTestUser)
	if err != nil {
		t.Fatal(err)
	}

	// get a valid workload ID
	wls, err := context.ds.GetWorkloads()
	if err != nil {
		t.Fatal(err)
	}

	if len(wls) == 0 {
		t.Fatal("No valid workloads")
	}

	url := computeURL + "/v2.1/" + tenant.ID + "/servers"

	var server payloads.ComputeCreateServer
	server.Server.MaxInstances = n
	server.Server.Workload = wls[0].ID

	b, err := json.Marshal(server)
	if err != nil {
		t.Fatal(err)
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(b))
	req.Header.Set("X-Auth-Token", "imavalidtoken")
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusAccepted {
		t.Fatalf("expected response code: %d, got %d", http.StatusAccepted, resp.StatusCode)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	var servers payloads.ComputeServers

	err = json.Unmarshal(body, &servers)
	if err != nil {
		t.Fatal(err)
	}

	if servers.TotalServers != n {
		t.Fatal("Not enough servers returned")
	}

	return servers
}

func testListServerDetailsTenant(t *testing.T, tenantID string) payloads.ComputeServers {
	url := computeURL + "/v2.1/" + tenantID + "/servers/detail"
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("X-Auth-Token", "imavalidtoken")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatalf("expected response code: %d, got %d", http.StatusOK, resp.StatusCode)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	var s payloads.ComputeServers
	err = json.Unmarshal(body, &s)
	if err != nil {
		t.Fatal(err)
	}

	return s
}

func TestCreateSingleServer(t *testing.T) {
	_ = testCreateServer(t, 1)
}

func TestListServerDetailsTenant(t *testing.T) {
	tenant, err := context.ds.GetTenant(computeTestUser)
	if err != nil {
		t.Fatal(err)
	}

	servers := testCreateServer(t, 1)
	if servers.TotalServers != 1 {
		t.Fatal(err)
	}

	s := testListServerDetailsTenant(t, tenant.ID)

	if s.TotalServers < 1 {
		t.Fatal("Not enough servers returned")
	}
}

func TestListServerDetailsWorkload(t *testing.T) {
	// get a valid workload ID
	wls, err := context.ds.GetWorkloads()
	if err != nil {
		t.Fatal(err)
	}

	if len(wls) == 0 {
		t.Fatal("No valid workloads")
	}

	servers := testCreateServer(t, 10)
	if servers.TotalServers != 10 {
		t.Fatal("failed to create enough servers")
	}

	url := computeURL + "/v2.1/flavors/" + wls[0].ID + "/servers/detail"
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("X-Auth-Token", "imavalidtoken")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := ioutil.ReadAll(resp.Body)
		t.Fatal(string(body))
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	var s payloads.ComputeServers
	err = json.Unmarshal(body, &s)
	if err != nil {
		t.Fatal(err)
	}

	if s.TotalServers < 10 {
		t.Fatal("Did not return correct number of servers")
	}
}

func TestShowServerDetails(t *testing.T) {
	tenant, err := context.ds.GetTenant(computeTestUser)
	if err != nil {
		t.Fatal(err)
	}

	tURL := computeURL + "/v2.1/" + tenant.ID + "/servers/"

	servers := testCreateServer(t, 1)
	if servers.TotalServers != 1 {
		t.Fatal(err)
	}

	s := testListServerDetailsTenant(t, tenant.ID)

	if s.TotalServers < 1 {
		t.Fatal("Not enough servers returned")
	}

	for _, s1 := range s.Servers {
		URL := tURL + s1.ID

		req, err := http.NewRequest("GET", URL, nil)
		req.Header.Set("X-Auth-Token", "imavalidtoken")

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		defer resp.Body.Close()

		if resp.StatusCode != http.StatusOK {
			t.Fatalf("expected response code: %d, got %d", http.StatusOK, resp.StatusCode)
		}

		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			t.Fatal(err)
		}

		var s2 payloads.ComputeServer
		err = json.Unmarshal(body, &s2)
		if err != nil {
			t.Fatal(err)
		}

		if reflect.DeepEqual(s1, s2.Server) == false {
			t.Fatal("Server details not correct")
		}
	}
}
<|endoftext|>


<commit_before>package post

import (
	"database/sql"
	"fmt"
	"log"
	"math"

	"github.com/rwcarlsen/cyan/query"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"CREATE TABLE IF NOT EXISTS AgentExit (SimId BLOB,AgentId INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Compositions (SimId BLOB,QualId INTEGER,NucId INTEGER, MassFrac REAL);",
		"CREATE TABLE IF NOT EXISTS Products (SimId BLOB,QualId INTEGER,Quality TEXT);",
		"CREATE TABLE IF NOT EXISTS Resources (SimId INTEGER,ResourceId INTEGER,ObjId INTEGER,Type TEXT,TimeCreated INTEGER,Quantity REAL,Units TEXT,QualId INTEGER,Parent1 INTEGER,Parent2 INTEGER);",
		"CREATE TABLE IF NOT EXISTS ResCreators (SimId INTEGER,ResourceId INTEGER,AgentId INTEGER);",
		"CREATE TABLE IF NOT EXISTS Agents (SimId BLOB,AgentId INTEGER,Kind TEXT,Spec TEXT,Prototype TEXT,ParentId INTEGER,Lifetime INTEGER,EnterTime INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Inventories (SimId BLOB,ResourceId INTEGER,AgentId INTEGER,StartTime INTEGER,EndTime INTEGER,QualId INTEGER,Quantity REAL);",
		"CREATE TABLE IF NOT EXISTS TimeList (SimId BLOB, Time INTEGER);",
		"CREATE TABLE IF NOT EXISTS Transactions (SimId BLOB, TransactionId INTEGER, SenderId INTEGER, ReceiverId INTEGER, ResourceId INTEGER, Commodity TEXT, Time INTEGER);",
		query.Index("TimeList", "Time"),
		query.Index("Resources", "SimId", "ResourceId", "QualId"),
		query.Index("Compositions", "SimId", "QualId", "NucId"),
		query.Index("Transactions", "SimId", "ResourceId"),
		query.Index("Transactions", "TransactionId"),
		query.Index("ResCreators", "SimId", "ResourceId"),
	}
	postExecStmts = []string{
		query.Index("Agents", "SimId", "Prototype"),
		query.Index("Inventories", "SimId", "AgentId"),
		query.Index("Inventories", "SimId", "ResourceId", "StartTime"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ResourceId,TimeCreated,QualId,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverId, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceId = ? AND tr.SimId = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ResourceId,res.TimeCreated,rc.AgentId,res.QualId,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ResourceId = rc.ResourceId
				  WHERE res.SimId = ? AND rc.SimId = ?;`
)

func Process(db *sql.DB) (simids [][]byte, err error) {
	err = Prepare(db)
	if err != nil {
		return nil, err
	}
	defer Finish(db)

	simids, err = GetSimIds(db)
	if err != nil {
		return nil, err
	}

	for _, id := range simids {
		ctx := NewContext(db, id)
		if err := ctx.WalkAll(); err != nil {
			log.Print(err)
		}
	}
	return simids, nil
}

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(db *sql.DB) (err error) {
	for _, s := range preExecStmts {
		if _, err := db.Exec(s); err != nil {
			log.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(db *sql.DB) (err error) {
	for _, s := range postExecStmts {
		if _, err := db.Exec(s); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	QualId    int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sql.DB
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       []byte
	Log         *log.Logger
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sql.Stmt
	dumpStmt    *sql.Stmt
	ownerStmt   *sql.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(db *sql.DB, simid []byte) *Context {
	return &Context{
		DB:    db,
		Simid: simid,
		Log:   log.New(NullWriter{}, "", 0),
	}
}

type AlreadyPostErr []byte

func (s AlreadyPostErr) Error() string {
	return fmt.Sprintf("SimId %x is already post processed", []byte(s))
}

func IsAlreadyPostErr(err error) bool {
	_, ok := err.(AlreadyPostErr)
	return ok
}

func (c *Context) init() {
	// skip if the post processing already exists for this simid in the db
	dummy := 0
	err := c.QueryRow("SELECT AgentId FROM Agents WHERE SimId = ? LIMIT 1", c.Simid).Scan(&dummy)
	if err == nil {
		panic(AlreadyPostErr(c.Simid))
	} else if err != sql.ErrNoRows {
		panicif(err)
	}

	tx, err := c.Begin()
	panicif(err)

	// build Agents table
	sql := `INSERT INTO Agents
				SELECT n.SimId,n.AgentId,n.Kind,n.Spec,n.Prototype,n.ParentId,n.Lifetime,n.EnterTime,x.ExitTime
				FROM
					AgentEntry AS n
					LEFT JOIN AgentExit AS x ON n.AgentId = x.AgentId AND n.SimId = x.SimId
					WHERE n.SimId = ?;`
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// build TimeList table
	sql = "SELECT Duration FROM Info WHERE SimId = ?;"
	rows, err := tx.Query(sql, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		var dur int
		panicif(rows.Scan(&dur))
		for i := 0; i < dur; i++ {
			_, err := tx.Exec("INSERT INTO TimeList VALUES (?, ?);", c.Simid, i)
			panicif(err)
		}
	}
	panicif(rows.Err())

	// create temp res table without simid
	c.Log.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + fmt.Sprintf("%x", c.Simid)
	_, err = tx.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql = "CREATE TABLE " + c.tmpResTbl + " AS SELECT ResourceId,TimeCreated,Parent1,Parent2,QualId,Quantity FROM Resources WHERE SimId = ?;"
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.Log.Println("Indexing temporary resource table...")
	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	tx.Commit()

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			if er, ok := r.(error); ok {
				err = er
			} else {
				err = fmt.Errorf("%v", r)
			}
		}
	}()

	c.Log.Printf("--- Building inventories for simid %x ---\n", c.Simid)
	c.init()

	c.Log.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	c.Log.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		c.Log.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	c.Log.Println("Dropping temporary resource table...")
	_, err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimId = ?"
	row := c.QueryRow(sql, c.Simid)

	n := 0
	err := row.Scan(&n)
	panicif(err)

	roots = make([]*Node, 0, n)
	rows, err := c.Query(rootsSql, c.Simid, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.QualId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	panicif(rows.Err())
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	rows, err := c.tmpResStmt.Query(node.ResId, node.ResId)
	panicif(err)
	defer rows.Close()

	for rows.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&child.ResId, &child.StartTime, &child.QualId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	panicif(rows.Err())
	rows.Close() // Close is idempotent - and this function is recursive.

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				QualId:    node.QualId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	rows, err := c.ownerStmt.Query(id, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	panicif(rows.Err())
	return owners, times
}

func (c *Context) dumpNodes() {
	c.Log.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	tx, err := c.Begin()
	panicif(err)
	stmt := tx.Stmt(c.dumpStmt)

	for _, n := range c.nodes {
		if n.EndTime > n.StartTime {
			_, err = stmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.QualId, n.Quantity)
			panicif(err)
		}
	}

	err = tx.Commit()
	panicif(err)
	c.nodes = c.nodes[:0]
}
<commit_msg>add another index to agents table and inventories table<commit_after>package post

import (
	"database/sql"
	"fmt"
	"log"
	"math"

	"github.com/rwcarlsen/cyan/query"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"CREATE TABLE IF NOT EXISTS AgentExit (SimId BLOB,AgentId INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Compositions (SimId BLOB,QualId INTEGER,NucId INTEGER, MassFrac REAL);",
		"CREATE TABLE IF NOT EXISTS Products (SimId BLOB,QualId INTEGER,Quality TEXT);",
		"CREATE TABLE IF NOT EXISTS Resources (SimId INTEGER,ResourceId INTEGER,ObjId INTEGER,Type TEXT,TimeCreated INTEGER,Quantity REAL,Units TEXT,QualId INTEGER,Parent1 INTEGER,Parent2 INTEGER);",
		"CREATE TABLE IF NOT EXISTS ResCreators (SimId INTEGER,ResourceId INTEGER,AgentId INTEGER);",
		"CREATE TABLE IF NOT EXISTS Agents (SimId BLOB,AgentId INTEGER,Kind TEXT,Spec TEXT,Prototype TEXT,ParentId INTEGER,Lifetime INTEGER,EnterTime INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Inventories (SimId BLOB,ResourceId INTEGER,AgentId INTEGER,StartTime INTEGER,EndTime INTEGER,QualId INTEGER,Quantity REAL);",
		"CREATE TABLE IF NOT EXISTS TimeList (SimId BLOB, Time INTEGER);",
		"CREATE TABLE IF NOT EXISTS Transactions (SimId BLOB, TransactionId INTEGER, SenderId INTEGER, ReceiverId INTEGER, ResourceId INTEGER, Commodity TEXT, Time INTEGER);",
		query.Index("TimeList", "Time"),
		query.Index("Resources", "SimId", "ResourceId", "QualId"),
		query.Index("Compositions", "SimId", "QualId", "NucId"),
		query.Index("Transactions", "SimId", "ResourceId"),
		query.Index("Transactions", "TransactionId"),
		query.Index("ResCreators", "SimId", "ResourceId"),
	}
	postExecStmts = []string{
		query.Index("Agents", "SimId", "Prototype"),
		query.Index("Agents", "SimId", "AgentId", "Prototype"),
		query.Index("Inventories", "SimId", "AgentId", "StartTime", "EndTime", "Quantity"),
		query.Index("Inventories", "SimId", "ResourceId", "StartTime"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ResourceId,TimeCreated,QualId,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverId, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceId = ? AND tr.SimId = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ResourceId,res.TimeCreated,rc.AgentId,res.QualId,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ResourceId = rc.ResourceId
				  WHERE res.SimId = ? AND rc.SimId = ?;`
)

func Process(db *sql.DB) (simids [][]byte, err error) {
	err = Prepare(db)
	if err != nil {
		return nil, err
	}
	defer Finish(db)

	simids, err = GetSimIds(db)
	if err != nil {
		return nil, err
	}

	for _, id := range simids {
		ctx := NewContext(db, id)
		if err := ctx.WalkAll(); err != nil {
			log.Print(err)
		}
	}
	return simids, nil
}

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(db *sql.DB) (err error) {
	for _, s := range preExecStmts {
		if _, err := db.Exec(s); err != nil {
			log.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(db *sql.DB) (err error) {
	for _, s := range postExecStmts {
		if _, err := db.Exec(s); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	QualId    int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sql.DB
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       []byte
	Log         *log.Logger
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sql.Stmt
	dumpStmt    *sql.Stmt
	ownerStmt   *sql.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(db *sql.DB, simid []byte) *Context {
	return &Context{
		DB:    db,
		Simid: simid,
		Log:   log.New(NullWriter{}, "", 0),
	}
}

type AlreadyPostErr []byte

func (s AlreadyPostErr) Error() string {
	return fmt.Sprintf("SimId %x is already post processed", []byte(s))
}

func IsAlreadyPostErr(err error) bool {
	_, ok := err.(AlreadyPostErr)
	return ok
}

func (c *Context) init() {
	// skip if the post processing already exists for this simid in the db
	dummy := 0
	err := c.QueryRow("SELECT AgentId FROM Agents WHERE SimId = ? LIMIT 1", c.Simid).Scan(&dummy)
	if err == nil {
		panic(AlreadyPostErr(c.Simid))
	} else if err != sql.ErrNoRows {
		panicif(err)
	}

	tx, err := c.Begin()
	panicif(err)

	// build Agents table
	sql := `INSERT INTO Agents
				SELECT n.SimId,n.AgentId,n.Kind,n.Spec,n.Prototype,n.ParentId,n.Lifetime,n.EnterTime,x.ExitTime
				FROM
					AgentEntry AS n
					LEFT JOIN AgentExit AS x ON n.AgentId = x.AgentId AND n.SimId = x.SimId
					WHERE n.SimId = ?;`
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// build TimeList table
	sql = "SELECT Duration FROM Info WHERE SimId = ?;"
	rows, err := tx.Query(sql, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		var dur int
		panicif(rows.Scan(&dur))
		for i := 0; i < dur; i++ {
			_, err := tx.Exec("INSERT INTO TimeList VALUES (?, ?);", c.Simid, i)
			panicif(err)
		}
	}
	panicif(rows.Err())

	// create temp res table without simid
	c.Log.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + fmt.Sprintf("%x", c.Simid)
	_, err = tx.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql = "CREATE TABLE " + c.tmpResTbl + " AS SELECT ResourceId,TimeCreated,Parent1,Parent2,QualId,Quantity FROM Resources WHERE SimId = ?;"
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.Log.Println("Indexing temporary resource table...")
	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	tx.Commit()

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			if er, ok := r.(error); ok {
				err = er
			} else {
				err = fmt.Errorf("%v", r)
			}
		}
	}()

	c.Log.Printf("--- Building inventories for simid %x ---\n", c.Simid)
	c.init()

	c.Log.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	c.Log.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		c.Log.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	c.Log.Println("Dropping temporary resource table...")
	_, err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimId = ?"
	row := c.QueryRow(sql, c.Simid)

	n := 0
	err := row.Scan(&n)
	panicif(err)

	roots = make([]*Node, 0, n)
	rows, err := c.Query(rootsSql, c.Simid, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.QualId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	panicif(rows.Err())
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	rows, err := c.tmpResStmt.Query(node.ResId, node.ResId)
	panicif(err)
	defer rows.Close()

	for rows.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&child.ResId, &child.StartTime, &child.QualId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	panicif(rows.Err())
	rows.Close() // Close is idempotent - and this function is recursive.

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				QualId:    node.QualId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	rows, err := c.ownerStmt.Query(id, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	panicif(rows.Err())
	return owners, times
}

func (c *Context) dumpNodes() {
	c.Log.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	tx, err := c.Begin()
	panicif(err)
	stmt := tx.Stmt(c.dumpStmt)

	for _, n := range c.nodes {
		if n.EndTime > n.StartTime {
			_, err = stmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.QualId, n.Quantity)
			panicif(err)
		}
	}

	err = tx.Commit()
	panicif(err)
	c.nodes = c.nodes[:0]
}
<|endoftext|>


<commit_before>package handlers

import (
	"fmt"

	"github.com/databus23/requestutil"
	"github.com/ghodss/yaml"
	"github.com/go-openapi/runtime/middleware"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	certutil "k8s.io/client-go/util/cert"

	"github.com/sapcc/kubernikus/pkg/api"
	"github.com/sapcc/kubernikus/pkg/api/models"
	"github.com/sapcc/kubernikus/pkg/api/rest/operations"
	"github.com/sapcc/kubernikus/pkg/client/kubernetes"
	"github.com/sapcc/kubernikus/pkg/util"
)

func NewGetClusterCredentials(rt *api.Runtime) operations.GetClusterCredentialsHandler {
	return &getClusterCredentials{rt}
}

type getClusterCredentials struct {
	*api.Runtime
}

func (d *getClusterCredentials) Handle(params operations.GetClusterCredentialsParams, principal *models.Principal) middleware.Responder {

	kluster, err := d.Klusters.Klusters(d.Namespace).Get(qualifiedName(params.Name, principal.Account))
	if err != nil {
		if apierrors.IsNotFound(err) {
			return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 404, "Kluster not found")
		}
		return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 500, err.Error())
	}
	secret, err := util.KlusterSecret(d.Kubernetes, kluster)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 404, "Secret not found")
		}
		return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 500, err.Error())
	}

	factory := util.NewCertificateFactory(kluster, &secret.Certificates, "")

	var organizations []string
	for _, role := range principal.Roles {
		organizations = append(organizations, "os:"+role)
	}

	cert, err := factory.UserCert(principal, fmt.Sprintf("%s://%s", requestutil.Scheme(params.HTTPRequest), requestutil.HostWithPort(params.HTTPRequest)))
	if err != nil {
		return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 500, "Failed to issue cert: %s", err)
	}
	config := kubernetes.NewClientConfigV1(
		params.Name,
		fmt.Sprintf("%v@%v", principal.Name, params.Name),
		kluster.Status.Apiserver,
		certutil.EncodePrivateKeyPEM(cert.PrivateKey),
		certutil.EncodeCertPEM(cert.Certificate),
		[]byte(secret.TLSCACertificate),
		"",
	)

	kubeconfig, err := yaml.Marshal(config)
	if err != nil {
		return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 500, "Failed to generate YAML document: %s", err)
	}

	credentials := models.Credentials{
		Kubeconfig: string(kubeconfig),
	}

	return operations.NewGetClusterCredentialsOK().WithPayload(&credentials)
}
<commit_msg>credential api: embed cluster version in cluster name (#461)<commit_after>package handlers

import (
	"fmt"

	"github.com/databus23/requestutil"
	"github.com/ghodss/yaml"
	"github.com/go-openapi/runtime/middleware"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	certutil "k8s.io/client-go/util/cert"

	"github.com/sapcc/kubernikus/pkg/api"
	"github.com/sapcc/kubernikus/pkg/api/models"
	"github.com/sapcc/kubernikus/pkg/api/rest/operations"
	"github.com/sapcc/kubernikus/pkg/client/kubernetes"
	"github.com/sapcc/kubernikus/pkg/util"
)

func NewGetClusterCredentials(rt *api.Runtime) operations.GetClusterCredentialsHandler {
	return &getClusterCredentials{rt}
}

type getClusterCredentials struct {
	*api.Runtime
}

func (d *getClusterCredentials) Handle(params operations.GetClusterCredentialsParams, principal *models.Principal) middleware.Responder {

	kluster, err := d.Klusters.Klusters(d.Namespace).Get(qualifiedName(params.Name, principal.Account))
	if err != nil {
		if apierrors.IsNotFound(err) {
			return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 404, "Kluster not found")
		}
		return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 500, err.Error())
	}
	secret, err := util.KlusterSecret(d.Kubernetes, kluster)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 404, "Secret not found")
		}
		return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 500, err.Error())
	}

	factory := util.NewCertificateFactory(kluster, &secret.Certificates, "")

	var organizations []string
	for _, role := range principal.Roles {
		organizations = append(organizations, "os:"+role)
	}

	cert, err := factory.UserCert(principal, fmt.Sprintf("%s://%s", requestutil.Scheme(params.HTTPRequest), requestutil.HostWithPort(params.HTTPRequest)))
	if err != nil {
		return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 500, "Failed to issue cert: %s", err)
	}
	config := kubernetes.NewClientConfigV1(
		params.Name,
		fmt.Sprintf("%v@%v", principal.Name, params.Name),
		kluster.Status.Apiserver,
		certutil.EncodePrivateKeyPEM(cert.PrivateKey),
		certutil.EncodeCertPEM(cert.Certificate),
		[]byte(secret.TLSCACertificate),
		"",
	)
	//embed version number in cluster name
	if kluster.Status.ApiserverVersion != "" {
		config.Clusters[0].Name += "@" + kluster.Status.ApiserverVersion
		config.Contexts[0].Context.Cluster = config.Clusters[0].Name
	}

	kubeconfig, err := yaml.Marshal(config)
	if err != nil {
		return NewErrorResponse(&operations.GetClusterCredentialsDefault{}, 500, "Failed to generate YAML document: %s", err)
	}

	credentials := models.Credentials{
		Kubeconfig: string(kubeconfig),
	}

	return operations.NewGetClusterCredentialsOK().WithPayload(&credentials)
}
<|endoftext|>


<commit_before>/*
Copyright The Kubernetes 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.
*/

// Code generated by client-gen. DO NOT EDIT.

package versioned

import (
	discovery "k8s.io/client-go/discovery"
	rest "k8s.io/client-go/rest"
	flowcontrol "k8s.io/client-go/util/flowcontrol"
	metricsv1alpha1 "k8s.io/metrics/pkg/client/clientset/versioned/typed/metrics/v1alpha1"
	metricsv1beta1 "k8s.io/metrics/pkg/client/clientset/versioned/typed/metrics/v1beta1"
)

type Interface interface {
	Discovery() discovery.DiscoveryInterface
	MetricsV1alpha1() metricsv1alpha1.MetricsV1alpha1Interface
	MetricsV1beta1() metricsv1beta1.MetricsV1beta1Interface
}

// Clientset contains the clients for groups. Each group has exactly one
// version included in a Clientset.
type Clientset struct {
	*discovery.DiscoveryClient
	metricsV1alpha1 *metricsv1alpha1.MetricsV1alpha1Client
	metricsV1beta1  *metricsv1beta1.MetricsV1beta1Client
}

// MetricsV1alpha1 retrieves the MetricsV1alpha1Client
func (c *Clientset) MetricsV1alpha1() metricsv1alpha1.MetricsV1alpha1Interface {
	return c.metricsV1alpha1
}

// MetricsV1beta1 retrieves the MetricsV1beta1Client
func (c *Clientset) MetricsV1beta1() metricsv1beta1.MetricsV1beta1Interface {
	return c.metricsV1beta1
}

// Discovery retrieves the DiscoveryClient
func (c *Clientset) Discovery() discovery.DiscoveryInterface {
	if c == nil {
		return nil
	}
	return c.DiscoveryClient
}

// NewForConfig creates a new Clientset for the given config.
func NewForConfig(c *rest.Config) (*Clientset, error) {
	configShallowCopy := *c
	if configShallowCopy.RateLimiter == nil && configShallowCopy.QPS > 0 {
		configShallowCopy.RateLimiter = flowcontrol.NewTokenBucketRateLimiter(configShallowCopy.QPS, configShallowCopy.Burst)
	}
	var cs Clientset
	var err error
	cs.metricsV1alpha1, err = metricsv1alpha1.NewForConfig(&configShallowCopy)
	if err != nil {
		return nil, err
	}
	cs.metricsV1beta1, err = metricsv1beta1.NewForConfig(&configShallowCopy)
	if err != nil {
		return nil, err
	}

	cs.DiscoveryClient, err = discovery.NewDiscoveryClientForConfig(&configShallowCopy)
	if err != nil {
		return nil, err
	}
	return &cs, nil
}

// NewForConfigOrDie creates a new Clientset for the given config and
// panics if there is an error in the config.
func NewForConfigOrDie(c *rest.Config) *Clientset {
	var cs Clientset
	cs.metricsV1alpha1 = metricsv1alpha1.NewForConfigOrDie(c)
	cs.metricsV1beta1 = metricsv1beta1.NewForConfigOrDie(c)

	cs.DiscoveryClient = discovery.NewDiscoveryClientForConfigOrDie(c)
	return &cs
}

// New creates a new Clientset for the given RESTClient.
func New(c rest.Interface) *Clientset {
	var cs Clientset
	cs.metricsV1alpha1 = metricsv1alpha1.New(c)
	cs.metricsV1beta1 = metricsv1beta1.New(c)

	cs.DiscoveryClient = discovery.NewDiscoveryClient(c)
	return &cs
}
<commit_msg>execute hack/update-codegen.sh, sync updates<commit_after>/*
Copyright The Kubernetes 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.
*/

// Code generated by client-gen. DO NOT EDIT.

package versioned

import (
	"fmt"

	discovery "k8s.io/client-go/discovery"
	rest "k8s.io/client-go/rest"
	flowcontrol "k8s.io/client-go/util/flowcontrol"
	metricsv1alpha1 "k8s.io/metrics/pkg/client/clientset/versioned/typed/metrics/v1alpha1"
	metricsv1beta1 "k8s.io/metrics/pkg/client/clientset/versioned/typed/metrics/v1beta1"
)

type Interface interface {
	Discovery() discovery.DiscoveryInterface
	MetricsV1alpha1() metricsv1alpha1.MetricsV1alpha1Interface
	MetricsV1beta1() metricsv1beta1.MetricsV1beta1Interface
}

// Clientset contains the clients for groups. Each group has exactly one
// version included in a Clientset.
type Clientset struct {
	*discovery.DiscoveryClient
	metricsV1alpha1 *metricsv1alpha1.MetricsV1alpha1Client
	metricsV1beta1  *metricsv1beta1.MetricsV1beta1Client
}

// MetricsV1alpha1 retrieves the MetricsV1alpha1Client
func (c *Clientset) MetricsV1alpha1() metricsv1alpha1.MetricsV1alpha1Interface {
	return c.metricsV1alpha1
}

// MetricsV1beta1 retrieves the MetricsV1beta1Client
func (c *Clientset) MetricsV1beta1() metricsv1beta1.MetricsV1beta1Interface {
	return c.metricsV1beta1
}

// Discovery retrieves the DiscoveryClient
func (c *Clientset) Discovery() discovery.DiscoveryInterface {
	if c == nil {
		return nil
	}
	return c.DiscoveryClient
}

// NewForConfig creates a new Clientset for the given config.
// If config's RateLimiter is not set and QPS and Burst are acceptable,
// NewForConfig will generate a rate-limiter in configShallowCopy.
func NewForConfig(c *rest.Config) (*Clientset, error) {
	configShallowCopy := *c
	if configShallowCopy.RateLimiter == nil && configShallowCopy.QPS > 0 {
		if configShallowCopy.Burst <= 0 {
			return nil, fmt.Errorf("Burst is required to be greater than 0 when RateLimiter is not set and QPS is set to greater than 0")
		}
		configShallowCopy.RateLimiter = flowcontrol.NewTokenBucketRateLimiter(configShallowCopy.QPS, configShallowCopy.Burst)
	}
	var cs Clientset
	var err error
	cs.metricsV1alpha1, err = metricsv1alpha1.NewForConfig(&configShallowCopy)
	if err != nil {
		return nil, err
	}
	cs.metricsV1beta1, err = metricsv1beta1.NewForConfig(&configShallowCopy)
	if err != nil {
		return nil, err
	}

	cs.DiscoveryClient, err = discovery.NewDiscoveryClientForConfig(&configShallowCopy)
	if err != nil {
		return nil, err
	}
	return &cs, nil
}

// NewForConfigOrDie creates a new Clientset for the given config and
// panics if there is an error in the config.
func NewForConfigOrDie(c *rest.Config) *Clientset {
	var cs Clientset
	cs.metricsV1alpha1 = metricsv1alpha1.NewForConfigOrDie(c)
	cs.metricsV1beta1 = metricsv1beta1.NewForConfigOrDie(c)

	cs.DiscoveryClient = discovery.NewDiscoveryClientForConfigOrDie(c)
	return &cs
}

// New creates a new Clientset for the given RESTClient.
func New(c rest.Interface) *Clientset {
	var cs Clientset
	cs.metricsV1alpha1 = metricsv1alpha1.New(c)
	cs.metricsV1beta1 = metricsv1beta1.New(c)

	cs.DiscoveryClient = discovery.NewDiscoveryClient(c)
	return &cs
}
<|endoftext|>


<commit_before>package pubsub

import (
	"fmt"
	"os"

	"cloud.google.com/go/pubsub"
)

// Option is the function signature required to be considered an pubsub.Option.
type Option func(*Transport) error

const (
	DefaultProjectEnvKey      = "GOOGLE_CLOUD_PROJECT"
	DefaultTopicEnvKey        = "PUBSUB_TOPIC"
	DefaultSubscriptionEnvKey = "PUBSUB_SUBSCRIPTION"
)

// WithEncoding sets the encoding for pubsub transport.
func WithEncoding(encoding Encoding) Option {
	return func(t *Transport) error {
		t.Encoding = encoding
		return nil
	}
}

// WithDefaultEncodingSelector sets the encoding selection strategy for
// default encoding selections based on Event.
func WithDefaultEncodingSelector(fn EncodingSelector) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http default encoding selector option can not set nil transport")
		}
		if fn != nil {
			t.DefaultEncodingSelectionFn = fn
			return nil
		}
		return fmt.Errorf("pubsub fn for DefaultEncodingSelector was nil")
	}
}

// WithBinaryEncoding sets the encoding selection strategy for
// default encoding selections based on Event, the encoded event will be the
// given version in Binary form.
func WithBinaryEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("pubsub binary encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = DefaultBinaryEncodingSelectionStrategy
		return nil
	}
}

// WithStructuredEncoding sets the encoding selection strategy for
// default encoding selections based on Event, the encoded event will be the
// given version in Structured form.
func WithStructuredEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("pubsub structured encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = DefaultStructuredEncodingSelectionStrategy
		return nil
	}
}

// WithClient sets the pubsub client for pubsub transport. Use this for explicit
// auth setup. Otherwise the env var 'GOOGLE_APPLICATION_CREDENTIALS' is used.
// See https://cloud.google.com/docs/authentication/production for more details.
func WithClient(client *pubsub.Client) Option {
	return func(t *Transport) error {
		t.client = client
		return nil
	}
}

// WithProjectID sets the project ID for pubsub transport.
func WithProjectID(projectID string) Option {
	return func(t *Transport) error {
		t.projectID = projectID
		return nil
	}
}

// WithProjectIDFromEnv sets the project ID for pubsub transport from a
// given environment variable name.
func WithProjectIDFromEnv(key string) Option {
	return func(t *Transport) error {
		v := os.Getenv(key)
		if v == "" {
			return fmt.Errorf("unable to load project id, %q environment variable not set", key)
		}
		t.projectID = v
		return nil
	}
}

// WithProjectIDFromDefaultEnv sets the project ID for pubsub transport from
// the environment variable named 'GOOGLE_CLOUD_PROJECT'.
func WithProjectIDFromDefaultEnv() Option {
	return WithProjectIDFromEnv(DefaultProjectEnvKey)
}

// WithTopicID sets the topic ID for pubsub transport.
func WithTopicID(topicID string) Option {
	return func(t *Transport) error {
		t.topicID = topicID
		return nil
	}
}

// WithTopicIDFromEnv sets the topic ID for pubsub transport from a given
// environment variable name.
func WithTopicIDFromEnv(key string) Option {
	return func(t *Transport) error {
		v := os.Getenv(key)
		if v == "" {
			return fmt.Errorf("unable to load topic id, %q environment variable not set", key)
		}
		t.topicID = v
		return nil
	}
}

// WithTopicIDFromDefaultEnv sets the topic ID for pubsub transport from the
// environment variable named 'PUBSUB_TOPIC'.
func WithTopicIDFromDefaultEnv() Option {
	return WithTopicIDFromEnv(DefaultTopicEnvKey)
}

// WithSubscriptionID sets the subscription ID for pubsub transport.
// This option can be used multiple times.
func WithSubscriptionID(subscriptionID string) Option {
	return func(t *Transport) error {
		if t.subscriptions == nil {
			t.subscriptions = make([]subscriptionWithTopic, 0)
		}
		t.subscriptions = append(t.subscriptions, subscriptionWithTopic{
			subscriptionID: subscriptionID,
		})
		return nil
	}
}

// WithSubscriptionIDFromEnv sets the subscription ID for pubsub transport from
// a given environment variable name.
func WithSubscriptionIDFromEnv(key string) Option {
	return func(t *Transport) error {
		v := os.Getenv(key)
		if v == "" {
			return fmt.Errorf("unable to load subscription id, %q environment variable not set", key)
		}

		opt := WithSubscriptionID(v)
		return opt(t)
	}
}

// WithSubscriptionIDFromDefaultEnv sets the subscription ID for pubsub
// transport from the environment variable named 'PUBSUB_SUBSCRIPTION'.
func WithSubscriptionIDFromDefaultEnv() Option {
	return WithSubscriptionIDFromEnv(DefaultSubscriptionEnvKey)
}

// AllowCreateTopic sets if the transport can create a topic if it does not
// exist.
func AllowCreateTopic(allow bool) Option {
	return func(t *Transport) error {
		t.AllowCreateTopic = allow
		return nil
	}
}

// AllowCreateSubscription sets if the transport can create a subscription if
// it does not exist.
func AllowCreateSubscription(allow bool) Option {
	return func(t *Transport) error {
		t.AllowCreateSubscription = allow
		return nil
	}
}
<commit_msg>adding option to also set the topic id<commit_after>package pubsub

import (
	"fmt"
	"os"

	"cloud.google.com/go/pubsub"
)

// Option is the function signature required to be considered an pubsub.Option.
type Option func(*Transport) error

const (
	DefaultProjectEnvKey      = "GOOGLE_CLOUD_PROJECT"
	DefaultTopicEnvKey        = "PUBSUB_TOPIC"
	DefaultSubscriptionEnvKey = "PUBSUB_SUBSCRIPTION"
)

// WithEncoding sets the encoding for pubsub transport.
func WithEncoding(encoding Encoding) Option {
	return func(t *Transport) error {
		t.Encoding = encoding
		return nil
	}
}

// WithDefaultEncodingSelector sets the encoding selection strategy for
// default encoding selections based on Event.
func WithDefaultEncodingSelector(fn EncodingSelector) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http default encoding selector option can not set nil transport")
		}
		if fn != nil {
			t.DefaultEncodingSelectionFn = fn
			return nil
		}
		return fmt.Errorf("pubsub fn for DefaultEncodingSelector was nil")
	}
}

// WithBinaryEncoding sets the encoding selection strategy for
// default encoding selections based on Event, the encoded event will be the
// given version in Binary form.
func WithBinaryEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("pubsub binary encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = DefaultBinaryEncodingSelectionStrategy
		return nil
	}
}

// WithStructuredEncoding sets the encoding selection strategy for
// default encoding selections based on Event, the encoded event will be the
// given version in Structured form.
func WithStructuredEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("pubsub structured encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = DefaultStructuredEncodingSelectionStrategy
		return nil
	}
}

// WithClient sets the pubsub client for pubsub transport. Use this for explicit
// auth setup. Otherwise the env var 'GOOGLE_APPLICATION_CREDENTIALS' is used.
// See https://cloud.google.com/docs/authentication/production for more details.
func WithClient(client *pubsub.Client) Option {
	return func(t *Transport) error {
		t.client = client
		return nil
	}
}

// WithProjectID sets the project ID for pubsub transport.
func WithProjectID(projectID string) Option {
	return func(t *Transport) error {
		t.projectID = projectID
		return nil
	}
}

// WithProjectIDFromEnv sets the project ID for pubsub transport from a
// given environment variable name.
func WithProjectIDFromEnv(key string) Option {
	return func(t *Transport) error {
		v := os.Getenv(key)
		if v == "" {
			return fmt.Errorf("unable to load project id, %q environment variable not set", key)
		}
		t.projectID = v
		return nil
	}
}

// WithProjectIDFromDefaultEnv sets the project ID for pubsub transport from
// the environment variable named 'GOOGLE_CLOUD_PROJECT'.
func WithProjectIDFromDefaultEnv() Option {
	return WithProjectIDFromEnv(DefaultProjectEnvKey)
}

// WithTopicID sets the topic ID for pubsub transport.
func WithTopicID(topicID string) Option {
	return func(t *Transport) error {
		t.topicID = topicID
		return nil
	}
}

// WithTopicIDFromEnv sets the topic ID for pubsub transport from a given
// environment variable name.
func WithTopicIDFromEnv(key string) Option {
	return func(t *Transport) error {
		v := os.Getenv(key)
		if v == "" {
			return fmt.Errorf("unable to load topic id, %q environment variable not set", key)
		}
		t.topicID = v
		return nil
	}
}

// WithTopicIDFromDefaultEnv sets the topic ID for pubsub transport from the
// environment variable named 'PUBSUB_TOPIC'.
func WithTopicIDFromDefaultEnv() Option {
	return WithTopicIDFromEnv(DefaultTopicEnvKey)
}

// WithSubscriptionID sets the subscription ID for pubsub transport.
// This option can be used multiple times.
func WithSubscriptionID(subscriptionID string) Option {
	return func(t *Transport) error {
		if t.subscriptions == nil {
			t.subscriptions = make([]subscriptionWithTopic, 0)
		}
		t.subscriptions = append(t.subscriptions, subscriptionWithTopic{
			subscriptionID: subscriptionID,
		})
		return nil
	}
}

// WithSubscriptionAndTopicID sets the subscription and topic IDs for pubsub transport.
// This option can be used multiple times.
func WithSubscriptionAndTopicID(subscriptionID, topicID string) Option {
	return func(t *Transport) error {
		if t.subscriptions == nil {
			t.subscriptions = make([]subscriptionWithTopic, 0)
		}
		t.subscriptions = append(t.subscriptions, subscriptionWithTopic{
			subscriptionID: subscriptionID,
			topicID:        topicID,
		})
		return nil
	}
}

// WithSubscriptionIDFromEnv sets the subscription ID for pubsub transport from
// a given environment variable name.
func WithSubscriptionIDFromEnv(key string) Option {
	return func(t *Transport) error {
		v := os.Getenv(key)
		if v == "" {
			return fmt.Errorf("unable to load subscription id, %q environment variable not set", key)
		}

		opt := WithSubscriptionID(v)
		return opt(t)
	}
}

// WithSubscriptionIDFromDefaultEnv sets the subscription ID for pubsub
// transport from the environment variable named 'PUBSUB_SUBSCRIPTION'.
func WithSubscriptionIDFromDefaultEnv() Option {
	return WithSubscriptionIDFromEnv(DefaultSubscriptionEnvKey)
}

// AllowCreateTopic sets if the transport can create a topic if it does not
// exist.
func AllowCreateTopic(allow bool) Option {
	return func(t *Transport) error {
		t.AllowCreateTopic = allow
		return nil
	}
}

// AllowCreateSubscription sets if the transport can create a subscription if
// it does not exist.
func AllowCreateSubscription(allow bool) Option {
	return func(t *Transport) error {
		t.AllowCreateSubscription = allow
		return nil
	}
}
<|endoftext|>


<commit_before><commit_msg>Fix DataImportCron watch race (#2238)<commit_after><|endoftext|>


<commit_before>package socialmedia

import (
	"net/http"

	"github.com/vardius/go-api-boilerplate/pkg/common/http/response"
	"github.com/vardius/go-api-boilerplate/pkg/common/jwt"
	"github.com/vardius/go-api-boilerplate/pkg/common/security/identity"
	user_grpc_server "github.com/vardius/go-api-boilerplate/pkg/user/interfaces/grpc"
	user_proto "github.com/vardius/go-api-boilerplate/pkg/user/interfaces/proto"
)

type google struct {
	client user_proto.UserClient
	jwt    jwt.Jwt
}

func (g *google) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	accessToken := r.FormValue("accessToken")
	data, e := getProfile(accessToken, "https://www.googleapis.com/oauth2/v2/userinfo")
	if e != nil {
		response.WithError(r.Context(), response.HTTPError{
			Code:    http.StatusBadRequest,
			Error:   e,
			Message: "Invalid access token",
		})
		return
	}

	identity := &identity.Identity{}
	identity.FromGoogleData(data)

	token, e := g.jwt.Encode(identity)
	if e != nil {
		response.WithError(r.Context(), response.HTTPError{
			Code:    http.StatusInternalServerError,
			Error:   e,
			Message: "Generate token failure",
		})
		return
	}

	payload := &commandPayload{token, data}
	_, e = g.client.DispatchCommand(r.Context(), &user_proto.DispatchCommandRequest{
		Name:    user_grpc_server.RegisterUserWithGoogle,
		Payload: payload.toJSON(),
	})

	if e != nil {
		response.WithError(r.Context(), response.HTTPError{
			Code:    http.StatusBadRequest,
			Error:   e,
			Message: "Invalid request",
		})
		return
	}

	response.WithPayload(r.Context(), &responsePayload{token, identity})
	return
}

// NewGoogle creates google auth handler
func NewGoogle(c user_proto.UserClient, j jwt.Jwt) http.Handler {
	return &google{c, j}
}
<commit_msg>Update google.go<commit_after>package socialmedia

import (
	"net/http"

	"github.com/vardius/go-api-boilerplate/pkg/common/http/response"
	"github.com/vardius/go-api-boilerplate/pkg/common/jwt"
	"github.com/vardius/go-api-boilerplate/pkg/common/security/identity"
	user_grpc "github.com/vardius/go-api-boilerplate/pkg/user/interfaces/grpc"
	user_proto "github.com/vardius/go-api-boilerplate/pkg/user/interfaces/proto"
)

type google struct {
	client user_proto.UserClient
	jwt    jwt.Jwt
}

func (g *google) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	accessToken := r.FormValue("accessToken")
	data, e := getProfile(accessToken, "https://www.googleapis.com/oauth2/v2/userinfo")
	if e != nil {
		response.WithError(r.Context(), response.HTTPError{
			Code:    http.StatusBadRequest,
			Error:   e,
			Message: "Invalid access token",
		})
		return
	}

	identity := &identity.Identity{}
	identity.FromGoogleData(data)

	token, e := g.jwt.Encode(identity)
	if e != nil {
		response.WithError(r.Context(), response.HTTPError{
			Code:    http.StatusInternalServerError,
			Error:   e,
			Message: "Generate token failure",
		})
		return
	}

	payload := &commandPayload{token, data}
	_, e = g.client.DispatchCommand(r.Context(), &user_proto.DispatchCommandRequest{
		Name:    user_grpc.RegisterUserWithGoogle,
		Payload: payload.toJSON(),
	})

	if e != nil {
		response.WithError(r.Context(), response.HTTPError{
			Code:    http.StatusBadRequest,
			Error:   e,
			Message: "Invalid request",
		})
		return
	}

	response.WithPayload(r.Context(), &responsePayload{token, identity})
	return
}

// NewGoogle creates google auth handler
func NewGoogle(c user_proto.UserClient, j jwt.Jwt) http.Handler {
	return &google{c, j}
}
<|endoftext|>


<commit_before>// +build integration

package integration

import (
	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/swan/pkg/executor"
	"github.com/intelsdi-x/swan/pkg/workloads/memcached"
	. "github.com/smartystreets/goconvey/convey"
	"os"
	"testing"
)

const (
	netstatCommand       = "echo stats | nc -w 1 127.0.0.1 11211"
	defaultMemcachedPath = "workloads/data_caching/memcached/memcached-1.4.25/build/memcached"
	swanPkg              = "github.com/intelsdi-x/swan"
)

// TestMemcachedWithExecutor is an integration test with local executor.
// See README for setup items.
func TestMemcachedWithExecutor(t *testing.T) {
	log.SetLevel(log.ErrorLevel)

	// Get optional custom Memcached path from MEMCACHED_PATH.
	memcachedPath := os.Getenv("MEMCACHED_BIN")

	if memcachedPath == "" {
		// If custom path does not exists use default path for built memcached.
		memcachedPath = os.Getenv("GOPATH") + "/src/" + swanPkg + "/" + defaultMemcachedPath
	}

	Convey("While using Local Shell in Memcached launcher", t, func() {
		l := executor.NewLocal()
		memcachedLauncher := memcached.New(
			l, memcached.DefaultMemcachedConfig(memcachedPath))

		Convey("When memcached is launched", func() {
			task, err := memcachedLauncher.Launch()

			Convey("There should be no error", func() {
				task.Stop()

				So(err, ShouldBeNil)
			})

			Convey("Wait 1 second for memcached to init", func() {
				isTerminated := task.Wait(1000)

				Convey("Memcached should be still running", func() {
					task.Stop()

					So(isTerminated, ShouldBeFalse)
				})

				Convey("When we check the memcached endpoint for stats after 1 second", func() {

					netstatTask, netstatErr := l.Execute(netstatCommand)

					Convey("There should be no error", func() {
						task.Stop()
						netstatTask.Stop()

						So(netstatErr, ShouldBeNil)

					})

					Convey("When we wait for netstat ", func() {
						netstatTask.Wait(0)

						Convey("The netstat task should be terminated, the task status should be 0"+
							" and output resultes with a STAT information", func() {
							netstatTaskState, netstatTaskStatus := netstatTask.Status()

							task.Stop()
							netstatTask.Stop()

							So(netstatTaskState, ShouldEqual, executor.TERMINATED)
							So(netstatTaskStatus.ExitCode, ShouldEqual, 0)
							So(netstatTaskStatus.Stdout, ShouldStartWith, "STAT")

						})
					})
				})

				Convey("When we stop the memcached task", func() {
					err := task.Stop()

					Convey("There should be no error", func() {
						So(err, ShouldBeNil)
					})

					Convey("The task should be terminated and the task status should be -15 or 0", func() {
						taskState, taskStatus := task.Status()
						So(taskState, ShouldEqual, executor.TERMINATED)
						// Memcached on CentOS returns 0 (successful code) after SIGTERM.
						So(taskStatus.ExitCode, ShouldBeIn, -15, 0)
					})
				})
			})
		})
	})
}
<commit_msg>Issues addressed.<commit_after>// +build integration

package integration

import (
	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/swan/pkg/executor"
	"github.com/intelsdi-x/swan/pkg/workloads/memcached"
	. "github.com/smartystreets/goconvey/convey"
	"os"
	"strings"
	"testing"
)

const (
	netstatCommand       = "echo stats | nc -w 1 127.0.0.1 11211"
	defaultMemcachedPath = "workloads/data_caching/memcached/memcached-1.4.25/build/memcached"
	swanPkg              = "github.com/intelsdi-x/swan"
)

// TestMemcachedWithExecutor is an integration test with local executor.
// See README for setup items.
func TestMemcachedWithExecutor(t *testing.T) {
	log.SetLevel(log.ErrorLevel)

	// Get optional custom Memcached path from MEMCACHED_PATH.
	memcachedPath := os.Getenv("MEMCACHED_BIN")

	if memcachedPath == "" {
		// If custom path does not exists use default path for built memcached.
		pathSlice := []string{os.Getenv("GOPATH"), "src", swanPkg, defaultMemcachedPath}
		memcachedPath = strings.Join(pathSlice, "/")
	}

	Convey("While using Local Shell in Memcached launcher", t, func() {
		l := executor.NewLocal()
		memcachedLauncher := memcached.New(
			l, memcached.DefaultMemcachedConfig(memcachedPath))

		Convey("When memcached is launched", func() {
			task, err := memcachedLauncher.Launch()

			Convey("There should be no error", func() {
				task.Stop()

				So(err, ShouldBeNil)
			})

			Convey("Wait 1 second for memcached to init", func() {
				isTerminated := task.Wait(1000)

				Convey("Memcached should be still running", func() {
					task.Stop()

					So(isTerminated, ShouldBeFalse)
				})

				Convey("When we check the memcached endpoint for stats after 1 second", func() {

					netstatTask, netstatErr := l.Execute(netstatCommand)

					Convey("There should be no error", func() {
						task.Stop()
						netstatTask.Stop()

						So(netstatErr, ShouldBeNil)

					})

					Convey("When we wait for netstat ", func() {
						netstatTask.Wait(0)

						Convey("The netstat task should be terminated, the task status should be 0"+
							" and output resultes with a STAT information", func() {
							netstatTaskState, netstatTaskStatus := netstatTask.Status()

							task.Stop()
							netstatTask.Stop()

							So(netstatTaskState, ShouldEqual, executor.TERMINATED)
							So(netstatTaskStatus.ExitCode, ShouldEqual, 0)
							So(netstatTaskStatus.Stdout, ShouldStartWith, "STAT")

						})
					})
				})

				Convey("When we stop the memcached task", func() {
					err := task.Stop()

					Convey("There should be no error", func() {
						So(err, ShouldBeNil)
					})

					Convey("The task should be terminated and the task status should be -15 or 0", func() {
						taskState, taskStatus := task.Status()
						So(taskState, ShouldEqual, executor.TERMINATED)
						// Memcached on CentOS returns 0 (successful code) after SIGTERM.
						So(taskStatus.ExitCode, ShouldBeIn, -15, 0)
					})
				})
			})
		})
	})
}
<|endoftext|>


<commit_before>// Copyright 2014 The Serviced 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 service

import (
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/domain"
	"github.com/control-center/serviced/domain/servicedefinition"
	"github.com/zenoss/elastigo/search"
	"github.com/zenoss/glog"

	"errors"
	"fmt"
	"strings"
)

//NewStore creates a Service  store
func NewStore() *Store {
	return &Store{}
}

//Store type for interacting with Service persistent storage
type Store struct {
	ds datastore.DataStore
}

// Put adds or updates a Service
func (s *Store) Put(ctx datastore.Context, svc *Service) error {
	//No need to store ConfigFiles
	svc.ConfigFiles = make(map[string]servicedefinition.ConfigFile)

	return s.ds.Put(ctx, Key(svc.ID), svc)
}

// fillBuiltinMetrics adds internal metrics to the monitoring profile
func fillBuiltinMetrics(svc *Service) {
	if svc.MonitoringProfile.MetricConfigs == nil {
		builder, err := domain.NewMetricConfigBuilder("/metrics/api/performance/query", "POST")
		if err != nil {
			glog.Errorf("Could not create builder to add internal metrics: %s", err)
			return
		}
		config, err := builder.Config("metrics", "metrics", "metrics", "-1h")
		if err != nil {
			glog.Errorf("could not create metric config for internal metrics: %s", err)
		}
		svc.MonitoringProfile.MetricConfigs = []domain.MetricConfig{*config}
	}
	index, found := findInternalMetricConfig(svc)
	if !found {
		glog.Errorf("should have been able to find internal metrics config")
		return
	}
	config := &svc.MonitoringProfile.MetricConfigs[index]
	removeInternalMetrics(config)
	removeInternalGraphConfigs(svc)

	if len(svc.Startup) > 2 {
		addInternalMetrics(config)
		addInternalGraphConfigs(svc)
	}
}

var internalCounterStats = []string{
	"net.collisions", "net.multicast", "net.rx_bytes", "net.rx_compressed",
	"net.rx_crc_errors", "net.rx_dropped", "net.rx_errors", "net.rx_fifo_errors",
	"net.rx_frame_errors", "net.rx_length_errors", "net.rx_missed_errors",
	"net.rx_over_errors", "net.rx_packets", "net.tx_aborted_errors",
	"net.tx_bytes", "net.tx_carrier_errors", "net.tx_compressed",
	"net.tx_dropped", "net.tx_errors", "net.tx_fifo_errors",
	"net.tx_heartbeat_errors", "net.tx_packets", "net.tx_window_errors",
	"cgroup.cpuacct.system", "cgroup.cpuacct.user", "cgroup.memory.pgmajfault",
}
var internalGuageStats = []string{
	"cgroup.memory.totalrss", "cgroup.memory.cache", "net.rx_bytes", "net.rx_compressed",
}

func removeInternalGraphConfigs(svc *Service) {
	var configs []domain.GraphConfig
	for _, config := range svc.MonitoringProfile.GraphConfigs {
		if config.BuiltIn {
			continue
		}
		configs = append(configs, config)
	}
	svc.MonitoringProfile.GraphConfigs = configs
}

func addInternalGraphConfigs(svc *Service) {

	tags := make(map[string][]string)
	tags["controlplane_service_id"] = []string{svc.ID}

	tRange := domain.GraphConfigRange{
		Start: "1h-ago",
		End:   "0s-ago",
	}
	zero := 0
	svc.MonitoringProfile.GraphConfigs = append(
		svc.MonitoringProfile.GraphConfigs,
		domain.GraphConfig{
			ID:          "internalusage",
			Name:        "CPU Usage",
			BuiltIn:     true,
			Format:      "%d",
			ReturnSet:   "EXACT",
			Type:        "area",
			Tags:        tags,
			YAxisLabel:  "% CPU Used",
			Description: "% CPU Used Over Last Hour",
			MinY:        &zero,
			Range:       &tRange,
			DataPoints: []domain.DataPoint{
				domain.DataPoint{
					Aggregator:   "avg",
					Format:       "%d",
					Legend:       "System",
					Metric:       "cgroup.cpuacct.system",
					MetricSource: "metrics",
					ID:           "cgroup.cpuacct.system",
					Name:         "System",
					Rate:         true,
					RateOptions: &domain.DataPointRateOptions{
						Counter: true,
					},
					Type: "area",
				},
				domain.DataPoint{
					Aggregator:   "avg",
					Format:       "%d",
					Legend:       "User",
					Metric:       "cgroup.cpuacct.user",
					MetricSource: "metrics",
					ID:           "cgroup.cpuacct.user",
					Name:         "User",
					Rate:         true,
					RateOptions: &domain.DataPointRateOptions{
						Counter: true,
					},
					Type: "area",
				},
			},
		},
	)

	// memory graph
	svc.MonitoringProfile.GraphConfigs = append(
		svc.MonitoringProfile.GraphConfigs,
		domain.GraphConfig{
			ID:          "internalMemoryUsage",
			Name:        "Memory Usage",
			BuiltIn:     true,
			Format:      "%6.2f",
			ReturnSet:   "EXACT",
			Type:        "area",
			Tags:        tags,
			YAxisLabel:  "GB",
			Description: "GB Memory Used Over Last Hour",
			MinY:        &zero,
			Range:       &tRange,
			DataPoints: []domain.DataPoint{
				domain.DataPoint{
					Aggregator:   "avg",
					Expression:   "rpn:1024,/,1024,/,1024,/",
					Fill:         true,
					Format:       "%6.2f",
					Legend:       "RSS",
					Metric:       "cgroup.memory.totalrss",
					MetricSource: "metrics",
					ID:           "cgroup.memory.totalrss",
					Name:         "Total RSS",
					Rate:         false,
					Type:         "area",
				},
				domain.DataPoint{
					Aggregator:   "avg",
					Expression:   "rpn:1024,/,1024,/,1024,/",
					Fill:         true,
					Format:       "%6.2f",
					Legend:       "Cache",
					Metric:       "cgroup.memory.cache",
					MetricSource: "metrics",
					ID:           "cgroup.memory.cache",
					Name:         "Cache",
					Rate:         false,
					Type:         "area",
				},
			},
		},
	)

	// network usage graph
	svc.MonitoringProfile.GraphConfigs = append(
		svc.MonitoringProfile.GraphConfigs,
		domain.GraphConfig{
			ID:          "internalNetworkUsage",
			Name:        "Network Usage",
			BuiltIn:     true,
			Format:      "%6.2f",
			ReturnSet:   "EXACT",
			Type:        "area",
			Tags:        tags,
			YAxisLabel:  "kbps",
			Range:       &tRange,
			Description: "kbps over last hour",
			DataPoints: []domain.DataPoint{
				domain.DataPoint{
					Aggregator:   "avg",
					Expression:   "rpn:8,/,1024,/",
					Fill:         true,
					Format:       "%6.2f",
					Legend:       "TX",
					Metric:       "net.tx_bytes",
					MetricSource: "metrics",
					ID:           "net.tx_bytes",
					Name:         "TX kbps",
					Rate:         true,
					Type:         "area",
				},
				domain.DataPoint{
					Aggregator:   "avg",
					Expression:   "rpn:8,/,1024,/",
					Fill:         true,
					Format:       "%6.2f",
					Legend:       "RX",
					Metric:       "net.rx_bytes",
					MetricSource: "metrics",
					ID:           "net.rx_bytes",
					Name:         "RX kbps",
					Rate:         true,
					Type:         "area",
				},
			},
		},
	)
}

// addInternalMetrics adds internal metrics to the config. It assumes that
// the current config does not container any internal metrics
func addInternalMetrics(config *domain.MetricConfig) {

	for _, metricName := range internalCounterStats {
		config.Metrics = append(config.Metrics,
			domain.Metric{
				ID:      metricName,
				Name:    metricName,
				Counter: true,
				BuiltIn: true,
			})

	}
	for _, metricName := range internalGuageStats {
		config.Metrics = append(config.Metrics,
			domain.Metric{
				ID:      metricName,
				Name:    metricName,
				Counter: false,
				BuiltIn: true,
			})

	}
}

func removeInternalMetrics(config *domain.MetricConfig) {
	// create an empty list of metrics
	var metrics []domain.Metric
	for _, metric := range config.Metrics {
		// and copy metrics, except built in ones
		if metric.BuiltIn {
			continue
		}
		metrics = append(metrics, metric)
	}
	config.Metrics = metrics
}

func findInternalMetricConfig(svc *Service) (index int, found bool) {
	// find the metric config
	for i := range svc.MonitoringProfile.MetricConfigs {
		if svc.MonitoringProfile.MetricConfigs[i].ID == "metrics" {
			return i, true
		}
	}
	builder, err := domain.NewMetricConfigBuilder("/metrics/api/performance/query", "POST")
	if err != nil {
		glog.Errorf("Could not create builder to add internal metrics: %s", err)
		return
	}
	config, err := builder.Config("metrics", "metrics", "metrics", "-1h")
	if err != nil {
		glog.Errorf("could not create metric config for internal metrics: %s", err)
	}
	svc.MonitoringProfile.MetricConfigs = append(
		svc.MonitoringProfile.MetricConfigs,
		*config)

	return len(svc.MonitoringProfile.MetricConfigs) - 1, true
}

// Get a Service by id. Return ErrNoSuchEntity if not found
func (s *Store) Get(ctx datastore.Context, id string) (*Service, error) {
	svc := &Service{}
	if err := s.ds.Get(ctx, Key(id), svc); err != nil {
		return nil, err
	}

	//Copy original config files
	fillConfig(svc)

	//Add builtin metrics
	fillBuiltinMetrics(svc)
	return svc, nil
}

// Delete removes the a Service if it exists
func (s *Store) Delete(ctx datastore.Context, id string) error {
	return s.ds.Delete(ctx, Key(id))
}

//GetServices returns all services
func (s *Store) GetServices(ctx datastore.Context) ([]*Service, error) {
	return query(ctx, "_exists_:ID")
}

//GetTaggedServices returns services with the given tags
func (s *Store) GetTaggedServices(ctx datastore.Context, tags ...string) ([]*Service, error) {
	if len(tags) == 0 {
		return nil, errors.New("empty tags not allowed")
	}
	qs := strings.Join(tags, " AND ")
	return query(ctx, qs)
}

//GetServicesByPool returns services with the given pool id
func (s *Store) GetServicesByPool(ctx datastore.Context, poolID string) ([]*Service, error) {
	id := strings.TrimSpace(poolID)
	if id == "" {
		return nil, errors.New("empty poolID not allowed")
	}
	queryString := fmt.Sprintf("PoolID:%s", id)
	return query(ctx, queryString)
}

//GetServicesByDeployment returns services with the given deployment id
func (s *Store) GetServicesByDeployment(ctx datastore.Context, deploymentID string) ([]*Service, error) {
	id := strings.TrimSpace(deploymentID)
	if id == "" {
		return nil, errors.New("empty deploymentID not allowed")
	}
	queryString := fmt.Sprintf("DeploymentID:%s", id)
	return query(ctx, queryString)
}

//GetChildServices returns services that are children of the given parent service id
func (s *Store) GetChildServices(ctx datastore.Context, parentID string) ([]*Service, error) {
	id := strings.TrimSpace(parentID)
	if id == "" {
		return nil, errors.New("empty parent service id not allowed")
	}

	queryString := fmt.Sprintf("ParentServiceID:%s", parentID)
	return query(ctx, queryString)
}

func query(ctx datastore.Context, query string) ([]*Service, error) {
	q := datastore.NewQuery(ctx)
	elasticQuery := search.Query().Search(query)
	search := search.Search("controlplane").Type(kind).Size("50000").Query(elasticQuery)
	results, err := q.Execute(search)
	if err != nil {
		return nil, err
	}
	return convert(results)
}

func fillConfig(svc *Service) {
	svc.ConfigFiles = make(map[string]servicedefinition.ConfigFile)
	for key, val := range svc.OriginalConfigs {
		svc.ConfigFiles[key] = val
	}
}

func convert(results datastore.Results) ([]*Service, error) {
	svcs := make([]*Service, results.Len())
	for idx := range svcs {
		var svc Service
		err := results.Get(idx, &svc)
		if err != nil {
			return nil, err
		}
		fillConfig(&svc)
		fillBuiltinMetrics(&svc)
		svcs[idx] = &svc
	}
	return svcs, nil
}

//Key creates a Key suitable for getting, putting and deleting Services
func Key(id string) datastore.Key {
	return datastore.NewKey(kind, id)
}

//confFileKey creates a Key suitable for getting, putting and deleting svcConfigFile
func confFileKey(id string) datastore.Key {
	return datastore.NewKey(confKind, id)
}

var (
	kind     = "service"
	confKind = "serviceconfig"
)
<commit_msg>Automated merge from v5.0.0b2<commit_after>// Copyright 2014 The Serviced 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 service

import (
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/domain"
	"github.com/control-center/serviced/domain/servicedefinition"
	"github.com/zenoss/elastigo/search"
	"github.com/zenoss/glog"

	"errors"
	"fmt"
	"strings"
)

//NewStore creates a Service  store
func NewStore() *Store {
	return &Store{}
}

//Store type for interacting with Service persistent storage
type Store struct {
	ds datastore.DataStore
}

// Put adds or updates a Service
func (s *Store) Put(ctx datastore.Context, svc *Service) error {
	//No need to store ConfigFiles
	svc.ConfigFiles = make(map[string]servicedefinition.ConfigFile)

	return s.ds.Put(ctx, Key(svc.ID), svc)
}

// fillBuiltinMetrics adds internal metrics to the monitoring profile
func fillBuiltinMetrics(svc *Service) {
	if svc.MonitoringProfile.MetricConfigs == nil {
		builder, err := domain.NewMetricConfigBuilder("/metrics/api/performance/query", "POST")
		if err != nil {
			glog.Errorf("Could not create builder to add internal metrics: %s", err)
			return
		}
		config, err := builder.Config("metrics", "metrics", "metrics", "-1h")
		if err != nil {
			glog.Errorf("could not create metric config for internal metrics: %s", err)
		}
		svc.MonitoringProfile.MetricConfigs = []domain.MetricConfig{*config}
	}
	index, found := findInternalMetricConfig(svc)
	if !found {
		glog.Errorf("should have been able to find internal metrics config")
		return
	}
	config := &svc.MonitoringProfile.MetricConfigs[index]
	removeInternalMetrics(config)
	removeInternalGraphConfigs(svc)

	if len(svc.Startup) > 2 {
		addInternalMetrics(config)
		addInternalGraphConfigs(svc)
	}
}

var internalCounterStats = []string{
	"net.collisions", "net.multicast", "net.rx_bytes", "net.rx_compressed",
	"net.rx_crc_errors", "net.rx_dropped", "net.rx_errors", "net.rx_fifo_errors",
	"net.rx_frame_errors", "net.rx_length_errors", "net.rx_missed_errors",
	"net.rx_over_errors", "net.rx_packets", "net.tx_aborted_errors",
	"net.tx_bytes", "net.tx_carrier_errors", "net.tx_compressed",
	"net.tx_dropped", "net.tx_errors", "net.tx_fifo_errors",
	"net.tx_heartbeat_errors", "net.tx_packets", "net.tx_window_errors",
	"cgroup.cpuacct.system", "cgroup.cpuacct.user", "cgroup.memory.pgmajfault",
}
var internalGuageStats = []string{
	"cgroup.memory.totalrss", "cgroup.memory.cache", "net.rx_bytes", "net.rx_compressed",
}

func removeInternalGraphConfigs(svc *Service) {
	var configs []domain.GraphConfig
	for _, config := range svc.MonitoringProfile.GraphConfigs {
		if config.BuiltIn {
			continue
		}
		configs = append(configs, config)
	}
	svc.MonitoringProfile.GraphConfigs = configs
}

func addInternalGraphConfigs(svc *Service) {

	tags := make(map[string][]string)
	tags["controlplane_service_id"] = []string{svc.ID}

	tRange := domain.GraphConfigRange{
		Start: "1h-ago",
		End:   "0s-ago",
	}
	zero := 0
	svc.MonitoringProfile.GraphConfigs = append(
		svc.MonitoringProfile.GraphConfigs,
		domain.GraphConfig{
			ID:          "internalusage",
			Name:        "CPU Usage",
			BuiltIn:     true,
			Format:      "%6.2f",
			ReturnSet:   "EXACT",
			Type:        "area",
			Tags:        tags,
			YAxisLabel:  "% CPU Used",
			Description: "% CPU Used Over Last Hour",
			MinY:        &zero,
			Range:       &tRange,
			DataPoints: []domain.DataPoint{
				domain.DataPoint{
					Aggregator:   "avg",
					Format:       "%6.2f",
					Legend:       "System",
					Metric:       "cgroup.cpuacct.system",
					MetricSource: "metrics",
					ID:           "cgroup.cpuacct.system",
					Name:         "System",
					Rate:         true,
					RateOptions: &domain.DataPointRateOptions{
						Counter: true,
					},
					Type: "area",
				},
				domain.DataPoint{
					Aggregator:   "avg",
					Format:       "%6.2f",
					Legend:       "User",
					Metric:       "cgroup.cpuacct.user",
					MetricSource: "metrics",
					ID:           "cgroup.cpuacct.user",
					Name:         "User",
					Rate:         true,
					RateOptions: &domain.DataPointRateOptions{
						Counter: true,
					},
					Type: "area",
				},
			},
		},
	)

	// memory graph
	svc.MonitoringProfile.GraphConfigs = append(
		svc.MonitoringProfile.GraphConfigs,
		domain.GraphConfig{
			ID:          "internalMemoryUsage",
			Name:        "Memory Usage",
			BuiltIn:     true,
			Format:      "%6.2f",
			ReturnSet:   "EXACT",
			Type:        "area",
			Tags:        tags,
			YAxisLabel:  "GB",
			Description: "GB Memory Used Over Last Hour",
			MinY:        &zero,
			Range:       &tRange,
			DataPoints: []domain.DataPoint{
				domain.DataPoint{
					Aggregator:   "avg",
					Expression:   "rpn:1024,/,1024,/,1024,/",
					Fill:         true,
					Format:       "%6.2f",
					Legend:       "RSS",
					Metric:       "cgroup.memory.totalrss",
					MetricSource: "metrics",
					ID:           "cgroup.memory.totalrss",
					Name:         "Total RSS",
					Rate:         false,
					Type:         "area",
				},
				domain.DataPoint{
					Aggregator:   "avg",
					Expression:   "rpn:1024,/,1024,/,1024,/",
					Fill:         true,
					Format:       "%6.2f",
					Legend:       "Cache",
					Metric:       "cgroup.memory.cache",
					MetricSource: "metrics",
					ID:           "cgroup.memory.cache",
					Name:         "Cache",
					Rate:         false,
					Type:         "area",
				},
			},
		},
	)

	// network usage graph
	svc.MonitoringProfile.GraphConfigs = append(
		svc.MonitoringProfile.GraphConfigs,
		domain.GraphConfig{
			ID:          "internalNetworkUsage",
			Name:        "Network Usage",
			BuiltIn:     true,
			Format:      "%6.2f",
			ReturnSet:   "EXACT",
			Type:        "area",
			Tags:        tags,
			YAxisLabel:  "kbps",
			Range:       &tRange,
			Description: "kbps over last hour",
			DataPoints: []domain.DataPoint{
				domain.DataPoint{
					Aggregator:   "avg",
					Expression:   "rpn:8,/,1024,/",
					Fill:         true,
					Format:       "%6.2f",
					Legend:       "TX",
					Metric:       "net.tx_bytes",
					MetricSource: "metrics",
					ID:           "net.tx_bytes",
					Name:         "TX kbps",
					Rate:         true,
					Type:         "area",
				},
				domain.DataPoint{
					Aggregator:   "avg",
					Expression:   "rpn:8,/,1024,/",
					Fill:         true,
					Format:       "%6.2f",
					Legend:       "RX",
					Metric:       "net.rx_bytes",
					MetricSource: "metrics",
					ID:           "net.rx_bytes",
					Name:         "RX kbps",
					Rate:         true,
					Type:         "area",
				},
			},
		},
	)
}

// addInternalMetrics adds internal metrics to the config. It assumes that
// the current config does not container any internal metrics
func addInternalMetrics(config *domain.MetricConfig) {

	for _, metricName := range internalCounterStats {
		config.Metrics = append(config.Metrics,
			domain.Metric{
				ID:      metricName,
				Name:    metricName,
				Counter: true,
				BuiltIn: true,
			})

	}
	for _, metricName := range internalGuageStats {
		config.Metrics = append(config.Metrics,
			domain.Metric{
				ID:      metricName,
				Name:    metricName,
				Counter: false,
				BuiltIn: true,
			})

	}
}

func removeInternalMetrics(config *domain.MetricConfig) {
	// create an empty list of metrics
	var metrics []domain.Metric
	for _, metric := range config.Metrics {
		// and copy metrics, except built in ones
		if metric.BuiltIn {
			continue
		}
		metrics = append(metrics, metric)
	}
	config.Metrics = metrics
}

func findInternalMetricConfig(svc *Service) (index int, found bool) {
	// find the metric config
	for i := range svc.MonitoringProfile.MetricConfigs {
		if svc.MonitoringProfile.MetricConfigs[i].ID == "metrics" {
			return i, true
		}
	}
	builder, err := domain.NewMetricConfigBuilder("/metrics/api/performance/query", "POST")
	if err != nil {
		glog.Errorf("Could not create builder to add internal metrics: %s", err)
		return
	}
	config, err := builder.Config("metrics", "metrics", "metrics", "-1h")
	if err != nil {
		glog.Errorf("could not create metric config for internal metrics: %s", err)
	}
	svc.MonitoringProfile.MetricConfigs = append(
		svc.MonitoringProfile.MetricConfigs,
		*config)

	return len(svc.MonitoringProfile.MetricConfigs) - 1, true
}

// Get a Service by id. Return ErrNoSuchEntity if not found
func (s *Store) Get(ctx datastore.Context, id string) (*Service, error) {
	svc := &Service{}
	if err := s.ds.Get(ctx, Key(id), svc); err != nil {
		return nil, err
	}

	//Copy original config files
	fillConfig(svc)

	//Add builtin metrics
	fillBuiltinMetrics(svc)
	return svc, nil
}

// Delete removes the a Service if it exists
func (s *Store) Delete(ctx datastore.Context, id string) error {
	return s.ds.Delete(ctx, Key(id))
}

//GetServices returns all services
func (s *Store) GetServices(ctx datastore.Context) ([]*Service, error) {
	return query(ctx, "_exists_:ID")
}

//GetTaggedServices returns services with the given tags
func (s *Store) GetTaggedServices(ctx datastore.Context, tags ...string) ([]*Service, error) {
	if len(tags) == 0 {
		return nil, errors.New("empty tags not allowed")
	}
	qs := strings.Join(tags, " AND ")
	return query(ctx, qs)
}

//GetServicesByPool returns services with the given pool id
func (s *Store) GetServicesByPool(ctx datastore.Context, poolID string) ([]*Service, error) {
	id := strings.TrimSpace(poolID)
	if id == "" {
		return nil, errors.New("empty poolID not allowed")
	}
	queryString := fmt.Sprintf("PoolID:%s", id)
	return query(ctx, queryString)
}

//GetServicesByDeployment returns services with the given deployment id
func (s *Store) GetServicesByDeployment(ctx datastore.Context, deploymentID string) ([]*Service, error) {
	id := strings.TrimSpace(deploymentID)
	if id == "" {
		return nil, errors.New("empty deploymentID not allowed")
	}
	queryString := fmt.Sprintf("DeploymentID:%s", id)
	return query(ctx, queryString)
}

//GetChildServices returns services that are children of the given parent service id
func (s *Store) GetChildServices(ctx datastore.Context, parentID string) ([]*Service, error) {
	id := strings.TrimSpace(parentID)
	if id == "" {
		return nil, errors.New("empty parent service id not allowed")
	}

	queryString := fmt.Sprintf("ParentServiceID:%s", parentID)
	return query(ctx, queryString)
}

func query(ctx datastore.Context, query string) ([]*Service, error) {
	q := datastore.NewQuery(ctx)
	elasticQuery := search.Query().Search(query)
	search := search.Search("controlplane").Type(kind).Size("50000").Query(elasticQuery)
	results, err := q.Execute(search)
	if err != nil {
		return nil, err
	}
	return convert(results)
}

func fillConfig(svc *Service) {
	svc.ConfigFiles = make(map[string]servicedefinition.ConfigFile)
	for key, val := range svc.OriginalConfigs {
		svc.ConfigFiles[key] = val
	}
}

func convert(results datastore.Results) ([]*Service, error) {
	svcs := make([]*Service, results.Len())
	for idx := range svcs {
		var svc Service
		err := results.Get(idx, &svc)
		if err != nil {
			return nil, err
		}
		fillConfig(&svc)
		fillBuiltinMetrics(&svc)
		svcs[idx] = &svc
	}
	return svcs, nil
}

//Key creates a Key suitable for getting, putting and deleting Services
func Key(id string) datastore.Key {
	return datastore.NewKey(kind, id)
}

//confFileKey creates a Key suitable for getting, putting and deleting svcConfigFile
func confFileKey(id string) datastore.Key {
	return datastore.NewKey(confKind, id)
}

var (
	kind     = "service"
	confKind = "serviceconfig"
)
<|endoftext|>


<commit_before>// Copyright 2020 The LUCI 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 lib

import (
	"context"
	"crypto"
	"fmt"
	"os"
	"path/filepath"
	"runtime"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/tree"
	repb "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
	"github.com/maruel/subcommands"
	"golang.org/x/sync/errgroup"

	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/data/caching/cache"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/isolated"
	"go.chromium.org/luci/common/system/filesystem"
	"go.chromium.org/luci/common/system/signals"
)

// CmdDownload returns an object for the `download` subcommand.
func CmdDownload() *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "download <options>...",
		ShortDesc: "download directory tree from a CAS server.",
		LongDesc: `Downloads directory tree from the CAS server.

Tree is referenced by their digest "<digest hash>/<size bytes>"`,
		CommandRun: func() subcommands.CommandRun {
			c := downloadRun{}
			c.Init()
			c.cachePolicies.AddFlags(&c.Flags)
			c.Flags.StringVar(&c.cacheDir, "cache-dir", "", "Cache directory to store downloaded files.")
			c.Flags.StringVar(&c.digest, "digest", "", `Digest of root directory proto "<digest hash>/<size bytes>".`)
			c.Flags.StringVar(&c.dir, "dir", "", "Directory to download tree.")
			return &c
		},
	}
}

type downloadRun struct {
	commonFlags
	digest string
	dir    string

	cacheDir      string
	cachePolicies cache.Policies
}

func (r *downloadRun) parse(a subcommands.Application, args []string) error {
	if err := r.commonFlags.Parse(); err != nil {
		return err
	}
	if len(args) != 0 {
		return errors.Reason("position arguments not expected").Err()
	}

	if r.cacheDir == "" && !r.cachePolicies.IsDefault() {
		return errors.New("cache-dir is necessary when cache-max-size, cache-max-items or cache-min-free-space are specified")
	}

	return nil
}

// doDownload downloads directory tree from the CAS server.
func (r *downloadRun) doDownload(ctx context.Context, args []string) error {
	ctx, cancel := context.WithCancel(ctx)
	signals.HandleInterrupt(cancel)

	d, err := digest.NewFromString(r.digest)
	if err != nil {
		return errors.Annotate(err, "failed to parse digest: %s", r.digest).Err()
	}

	c, err := r.casFlags.NewClient(ctx, true)
	if err != nil {
		return err
	}

	rootDir := &repb.Directory{}
	if err := c.ReadProto(ctx, d, rootDir); err != nil {
		return errors.Annotate(err, "failed to read root directory proto").Err()
	}

	dirs, err := c.GetDirectoryTree(ctx, d.ToProto())
	if err != nil {
		return errors.Annotate(err, "failed to call GetDirectoryTree").Err()
	}

	t := &repb.Tree{
		Root:     rootDir,
		Children: dirs,
	}

	outputs, err := tree.FlattenTree(t, r.dir)
	if err != nil {
		return errors.Annotate(err, "failed to call FlattenTree").Err()
	}

	to := make(map[digest.Digest]*tree.Output)

	var diskcache cache.Cache
	if r.cacheDir != "" {
		diskcache, err = cache.NewDisk(r.cachePolicies, r.cacheDir, crypto.SHA256)
		if err != nil {
			return errors.Annotate(err, "failed to create initialize cache").Err()
		}
		defer diskcache.Close()
	}

	// Files have the same digest are downloaded only once, so we need to
	// copy duplicates files later.
	var dups []*tree.Output

	for path, output := range outputs {
		if output.IsEmptyDirectory {
			if err := os.MkdirAll(path, 0o700); err != nil {
				return errors.Annotate(err, "failed to create directory").Err()
			}
			continue
		}

		if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
			return errors.Annotate(err, "failed to create directory").Err()
		}

		if output.SymlinkTarget != "" {
			if err := os.Symlink(output.SymlinkTarget, path); err != nil {
				return errors.Annotate(err, "failed to create symlink").Err()
			}
			continue
		}

		if diskcache != nil {
			mode := 0o600
			if output.IsExecutable {
				mode = 0o700
			}

			if err := diskcache.Hardlink(isolated.HexDigest(output.Digest.Hash), path, os.FileMode(mode)); err == nil {
				continue
			} else if !errors.Contains(err, os.ErrNotExist) {
				return err
			}
		}

		if _, ok := to[output.Digest]; ok {
			dups = append(dups, output)
		} else {
			to[output.Digest] = output
		}
	}

	if err := c.DownloadFiles(ctx, ".", to); err != nil {
		return errors.Annotate(err, "failed to download files").Err()
	}

	// DownloadFiles does not set desired file permission.
	for d, output := range to {
		mode := 0o600
		if output.IsExecutable {
			mode = 0o700
		}
		if err := os.Chmod(output.Path, os.FileMode(mode)); err != nil {
			return errors.Annotate(err, "failed to change mode").Err()
		}

		if diskcache != nil {
			if err := diskcache.AddFileWithoutValidation(
				isolated.HexDigest(d.Hash), output.Path); err != nil {
				return errors.Annotate(err, "failed to add cache; path=%s digest=%s", output.Path, d).Err()
			}
		}
	}

	eg, ctx := errgroup.WithContext(ctx)

	// limit the number of concurrent I/O threads.
	ch := make(chan struct{}, runtime.NumCPU())

	for _, dup := range dups {
		src := to[dup.Digest]
		dst := dup
		ch <- struct{}{}
		eg.Go(func() (err error) {
			defer func() { <-ch }()
			mode := 0o600
			if dst.IsExecutable {
				mode = 0o700
			}

			if err := filesystem.Copy(dst.Path, src.Path, os.FileMode(mode)); err != nil {
				return errors.Annotate(err, "failed to copy file from '%s' to '%s'", src.Path, dst.Path).Err()
			}

			return nil
		})
	}

	return eg.Wait()
}

func (r *downloadRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
	ctx := cli.GetContext(a, r, env)

	if err := r.parse(a, args); err != nil {
		errors.Log(ctx, err)
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}

	if err := r.doDownload(ctx, args); err != nil {
		errors.Log(ctx, err)
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}

	return 0
}
<commit_msg>cas: introduce logs in download<commit_after>// Copyright 2020 The LUCI 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 lib

import (
	"context"
	"crypto"
	"fmt"
	"os"
	"path/filepath"
	"runtime"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/tree"
	repb "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
	"github.com/maruel/subcommands"
	"golang.org/x/sync/errgroup"

	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/data/caching/cache"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/isolated"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/system/filesystem"
	"go.chromium.org/luci/common/system/signals"
)

// CmdDownload returns an object for the `download` subcommand.
func CmdDownload() *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "download <options>...",
		ShortDesc: "download directory tree from a CAS server.",
		LongDesc: `Downloads directory tree from the CAS server.

Tree is referenced by their digest "<digest hash>/<size bytes>"`,
		CommandRun: func() subcommands.CommandRun {
			c := downloadRun{}
			c.Init()
			c.cachePolicies.AddFlags(&c.Flags)
			c.Flags.StringVar(&c.cacheDir, "cache-dir", "", "Cache directory to store downloaded files.")
			c.Flags.StringVar(&c.digest, "digest", "", `Digest of root directory proto "<digest hash>/<size bytes>".`)
			c.Flags.StringVar(&c.dir, "dir", "", "Directory to download tree.")
			return &c
		},
	}
}

type downloadRun struct {
	commonFlags
	digest string
	dir    string

	cacheDir      string
	cachePolicies cache.Policies
}

func (r *downloadRun) parse(a subcommands.Application, args []string) error {
	if err := r.commonFlags.Parse(); err != nil {
		return err
	}
	if len(args) != 0 {
		return errors.Reason("position arguments not expected").Err()
	}

	if r.cacheDir == "" && !r.cachePolicies.IsDefault() {
		return errors.New("cache-dir is necessary when cache-max-size, cache-max-items or cache-min-free-space are specified")
	}

	return nil
}

// doDownload downloads directory tree from the CAS server.
func (r *downloadRun) doDownload(ctx context.Context, args []string) error {
	ctx, cancel := context.WithCancel(ctx)
	signals.HandleInterrupt(cancel)

	d, err := digest.NewFromString(r.digest)
	if err != nil {
		return errors.Annotate(err, "failed to parse digest: %s", r.digest).Err()
	}

	c, err := r.casFlags.NewClient(ctx, true)
	if err != nil {
		return err
	}

	rootDir := &repb.Directory{}
	if err := c.ReadProto(ctx, d, rootDir); err != nil {
		return errors.Annotate(err, "failed to read root directory proto").Err()
	}

	dirs, err := c.GetDirectoryTree(ctx, d.ToProto())
	if err != nil {
		return errors.Annotate(err, "failed to call GetDirectoryTree").Err()
	}
	logger := logging.Get(ctx)
	logger.Infof("finished GetDirectoryTree api call: %d", len(dirs))

	t := &repb.Tree{
		Root:     rootDir,
		Children: dirs,
	}

	outputs, err := tree.FlattenTree(t, r.dir)
	if err != nil {
		return errors.Annotate(err, "failed to call FlattenTree").Err()
	}

	to := make(map[digest.Digest]*tree.Output)

	var diskcache cache.Cache
	if r.cacheDir != "" {
		diskcache, err = cache.NewDisk(r.cachePolicies, r.cacheDir, crypto.SHA256)
		if err != nil {
			return errors.Annotate(err, "failed to create initialize cache").Err()
		}
		defer diskcache.Close()
	}

	// Files have the same digest are downloaded only once, so we need to
	// copy duplicates files later.
	var dups []*tree.Output

	for path, output := range outputs {
		if output.IsEmptyDirectory {
			if err := os.MkdirAll(path, 0o700); err != nil {
				return errors.Annotate(err, "failed to create directory").Err()
			}
			continue
		}

		if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
			return errors.Annotate(err, "failed to create directory").Err()
		}

		if output.SymlinkTarget != "" {
			if err := os.Symlink(output.SymlinkTarget, path); err != nil {
				return errors.Annotate(err, "failed to create symlink").Err()
			}
			continue
		}

		if diskcache != nil {
			mode := 0o600
			if output.IsExecutable {
				mode = 0o700
			}

			if err := diskcache.Hardlink(isolated.HexDigest(output.Digest.Hash), path, os.FileMode(mode)); err == nil {
				continue
			} else if !errors.Contains(err, os.ErrNotExist) {
				return err
			}
		}

		if _, ok := to[output.Digest]; ok {
			dups = append(dups, output)
		} else {
			to[output.Digest] = output
		}
	}
	logger.Infof("finished copy from cache (if any), dups: %d, to: %d", len(dirs), len(to))

	if err := c.DownloadFiles(ctx, ".", to); err != nil {
		return errors.Annotate(err, "failed to download files").Err()
	}

	logger.Infof("finished DownloadFiles api call")

	// DownloadFiles does not set desired file permission.
	for _, output := range to {
		mode := 0o600
		if output.IsExecutable {
			mode = 0o700
		}
		if err := os.Chmod(output.Path, os.FileMode(mode)); err != nil {
			return errors.Annotate(err, "failed to change mode").Err()
		}
	}

	logger.Infof("finished permission update")

	if diskcache != nil {
		for d, output := range to {
			if err := diskcache.AddFileWithoutValidation(
				isolated.HexDigest(d.Hash), output.Path); err != nil {
				return errors.Annotate(err, "failed to add cache; path=%s digest=%s", output.Path, d).Err()
			}
		}
		logger.Infof("finished cache addition")
	}

	eg, ctx := errgroup.WithContext(ctx)

	// limit the number of concurrent I/O threads.
	ch := make(chan struct{}, runtime.NumCPU())

	for _, dup := range dups {
		src := to[dup.Digest]
		dst := dup
		ch <- struct{}{}
		eg.Go(func() (err error) {
			defer func() { <-ch }()
			mode := 0o600
			if dst.IsExecutable {
				mode = 0o700
			}

			if err := filesystem.Copy(dst.Path, src.Path, os.FileMode(mode)); err != nil {
				return errors.Annotate(err, "failed to copy file from '%s' to '%s'", src.Path, dst.Path).Err()
			}

			return nil
		})
	}
	err = eg.Wait()

	logger.Infof("finished file copy")

	return err
}

func (r *downloadRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
	ctx := cli.GetContext(a, r, env)
	logging.Get(ctx).Infof("start command")
	if err := r.parse(a, args); err != nil {
		errors.Log(ctx, err)
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}

	if err := r.doDownload(ctx, args); err != nil {
		errors.Log(ctx, err)
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}

	return 0
}
<|endoftext|>