<commit_before>package annotation

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/m-lab/etl/web100"

	"github.com/m-lab/etl/metrics"
	"github.com/prometheus/client_golang/prometheus"
)

var IPAnnotationEnabled = false

func init() {
	getFlagValues()
}

func getFlagValues() {
	// Check for ANNOTATE_IP = 'true'
	flag, ok := os.LookupEnv("ANNOTATE_IP")
	if ok {
		IPAnnotationEnabled, _ = strconv.ParseBool(flag)
		// If parse fails, then ipAnn will be set to false.
	}
}

// EnableAnnotation is used only for testing.  Should be placed in whitebox _test file.
func EnableAnnotation() {
	os.Setenv("ANNOTATE_IP", "True")
	getFlagValues()
}

// The GeolocationIP struct contains all the information needed for the
// geolocation data that will be inserted into big query. The fiels are
// capitalized for exporting, although the originals in the DB schema
// are not.
type GeolocationIP struct {
	Continent_code string  `json:"continent_code, string,omitempty"` // Gives a shorthand for the continent
	Country_code   string  `json:"country_code, string,omitempty"`   // Gives a shorthand for the country
	Country_code3  string  `json:"country_code3, string,omitempty"`  // Gives a shorthand for the country
	Country_name   string  `json:"country_name, string,omitempty"`   // Name of the country
	Region         string  `json:"region, string,omitempty"`         // Region or State within the country
	Metro_code     int64   `json:"metro_code, integer,omitempty"`    // Metro code within the country
	City           string  `json:"city, string,omitempty"`           // City within the region
	Area_code      int64   `json:"area_code, integer,omitempty"`     // Area code, similar to metro code
	Postal_code    string  `json:"postal_code, string,omitempty"`    // Postal code, again similar to metro
	Latitude       float64 `json:"latitude, float"`                  // Latitude
	Longitude      float64 `json:"longitude, float"`                 // Longitude

}

// The struct that will hold the IP/ASN data when it gets added to the
// schema. Currently empty and unused.
type IPASNData struct{}

// The main struct for the geo metadata, which holds pointers to the
// Geolocation data and the IP/ASN data. This is what we parse the JSON
// response from the annotator into.
type GeoData struct {
	Geo *GeolocationIP // Holds the geolocation data
	ASN *IPASNData     // Holds the IP/ASN data
}

// The RequestData schema is the schema for the json that we will send
// down the pipe to the annotation service.
type RequestData struct {
	IP        string    // Holds the IP from an incoming request
	IPFormat  int       // Holds the ip format, 4 or 6
	Timestamp time.Time // Holds the timestamp from an incoming request
}

// AnnotatorURL holds the https address of the annotator.
// TODO(gfr) See if there is a better way of determining
// where to send the request (there almost certainly is)
var AnnotatorURL = "https://annotator-dot-" +
	os.Getenv("GCLOUD_PROJECT") +
	".appspot.com"

// BaseURL provides the base URL for single annotation requests
var BaseURL = AnnotatorURL + "/annotate?"

// BatchURL provides the base URL for batch annotation requests
var BatchURL = AnnotatorURL + "/batch_annotate"

// FetchGeoAnnotations takes a slice of strings
// containing ip addresses, a timestamp, and a slice of pointers to
// the GeolocationIP structs that correspond to the ip addresses. A
// precondition assumed by this function is that both slices are the
// same length. It will then make a call to the batch annotator, using
// the ip addresses and the timestamp. Then, it uses that data to fill
// in the structs pointed to by the slice of GeolocationIP pointers.
func FetchGeoAnnotations(ips []string, timestamp time.Time, geoDest []*GeolocationIP) {
	reqData := make([]RequestData, 0, len(ips))
	for _, ip := range ips {
		if ip == "" {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Empty IP Address!!!"}).Inc()
			continue
		}
		ip, _ = web100.NormalizeIPv6(ip)
		reqData = append(reqData, RequestData{ip, 0, timestamp})
	}
	annotationData := GetBatchGeoData(BatchURL, reqData)
	timeString := strconv.FormatInt(timestamp.Unix(), 36)
	for index, ip := range ips {
		data, ok := annotationData[ip+timeString]
		if !ok || data.Geo == nil {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Missing or empty data for IP Address!!!"}).Inc()
			continue
		}
		*geoDest[index] = *data.Geo

	}
}

// GetAndInsertGeolocationIPStruct takes a NON-NIL pointer to a
// pre-allocated GeolocationIP struct, an IP address, and a
// timestamp. It will connect to the annotation service, get the
// geo data, and insert the geo data into the reigion pointed to by
// the GeolocationIP pointer.
func GetAndInsertGeolocationIPStruct(geo *GeolocationIP, ip string, timestamp time.Time) {
	url := BaseURL + "ip_addr=" + url.QueryEscape(ip) +
		"&since_epoch=" + strconv.FormatInt(timestamp.Unix(), 10)
	annotationData := GetGeoData(url)
	if annotationData != nil && annotationData.Geo != nil {
		*geo = *annotationData.Geo
	}
}

// GetGeoData combines the functionality of QueryAnnotationService and
// ParseJSONGeoDataResponse to query the annotator service and return
// the corresponding GeoData if it can, or a nil pointer if it
// encounters any error and cannot get the data for any reason
func GetGeoData(url string) *GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := QueryAnnotationService(url)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := ParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		return nil
	}
	return geoDataFromResponse
}

// QueryAnnotationService will connect to the annotation service and
// copy the body of a valid response to a byte slice and return it to a
// user, returning an error if any occurs
func QueryAnnotationService(url string) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()
	// Make the actual request
	resp, err := http.Get(url)

	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Request to Annotator failed"}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Webserver gave non-ok response"}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// ParseJSONGeoDataResponse takes a byte slice containing the test of
// the JSON from the annotator service and parses it into a GeoData
// struct, for easy manipulation. It returns a pointer to the struct on
// success and an error if an error occurs.
func ParseJSONGeoDataResponse(jsonBuffer []byte) (*GeoData, error) {
	parsedJSON := &GeoData{}
	err := json.Unmarshal(jsonBuffer, parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}

// GetBatchGeoData combines the functionality of
// BatchQueryAnnotationService and BatchParseJSONGeoDataResponse to
// query the annotator service and return the corresponding map of
// ip-timestamp strings to GeoData structs, or a nil map if it
// encounters any error and cannot get the data for any reason
// TODO - dedup common code in GetGeoData
func GetBatchGeoData(url string, data []RequestData) map[string]GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := BatchQueryAnnotationService(url, data)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := BatchParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		log.Printf("%+v\n", data)
		return nil
	}
	return geoDataFromResponse
}

// BatchQueryAnnotationService takes a url to POST the request to and
// a slice of RequestDatas to be sent in the body in a JSON
// format. It will copy the response into a []byte and return it to
// the user, returning an error if any occurs
// TODO(gfr) Should pass the annotator's request context through and use it here.
func BatchQueryAnnotationService(url string, data []RequestData) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()

	encodedData, err := json.Marshal(data)
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Couldn't Marshal Data"}).Inc()
		return nil, err
	}

	var netClient = &http.Client{
		// Median response time is < 10 msec, but 99th percentile is 0.6 seconds.
		Timeout: 2 * time.Second,
	}

	// Make the actual request
	resp, err := netClient.Post(url, "raw", bytes.NewReader(encodedData))
	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": err.Error()}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": http.StatusText(resp.StatusCode)}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// BatchParseJSONGeoDataResponse takes a byte slice containing the
// text of the JSON from the annoator service's batch request endpoint
// and parses it into a map of strings to GeoData structs, for
// easy manipulation. It returns a pointer to the struct on success
// and an error if one occurs.
// TODO - is there duplicate code with ParseJSON... ?
func BatchParseJSONGeoDataResponse(jsonBuffer []byte) (map[string]GeoData, error) {
	parsedJSON := make(map[string]GeoData)
	err := json.Unmarshal(jsonBuffer, &parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}
<commit_msg>Fix annotation json tags<commit_after>package annotation

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/m-lab/etl/web100"

	"github.com/m-lab/etl/metrics"
	"github.com/prometheus/client_golang/prometheus"
)

var IPAnnotationEnabled = false

func init() {
	getFlagValues()
}

func getFlagValues() {
	// Check for ANNOTATE_IP = 'true'
	flag, ok := os.LookupEnv("ANNOTATE_IP")
	if ok {
		IPAnnotationEnabled, _ = strconv.ParseBool(flag)
		// If parse fails, then ipAnn will be set to false.
	}
}

// EnableAnnotation is used only for testing.  Should be placed in whitebox _test file.
func EnableAnnotation() {
	os.Setenv("ANNOTATE_IP", "True")
	getFlagValues()
}

// The GeolocationIP struct contains all the information needed for the
// geolocation data that will be inserted into big query. The fiels are
// capitalized for exporting, although the originals in the DB schema
// are not.
type GeolocationIP struct {
	Continent_code string  `json:"continent_code,string,omitempty"` // Gives a shorthand for the continent
	Country_code   string  `json:"country_code,string,omitempty"`   // Gives a shorthand for the country
	Country_code3  string  `json:"country_code3,string,omitempty"`  // Gives a shorthand for the country
	Country_name   string  `json:"country_name,string,omitempty"`   // Name of the country
	Region         string  `json:"region,string,omitempty"`         // Region or State within the country
	Metro_code     int64   `json:"metro_code,integer,omitempty"`    // Metro code within the country
	City           string  `json:"city,string,omitempty"`           // City within the region
	Area_code      int64   `json:"area_code,integer,omitempty"`     // Area code, similar to metro code
	Postal_code    string  `json:"postal_code,string,omitempty"`    // Postal code, again similar to metro
	Latitude       float64 `json:"latitude,float"`                  // Latitude
	Longitude      float64 `json:"longitude,float"`                 // Longitude

}

// The struct that will hold the IP/ASN data when it gets added to the
// schema. Currently empty and unused.
type IPASNData struct{}

// The main struct for the geo metadata, which holds pointers to the
// Geolocation data and the IP/ASN data. This is what we parse the JSON
// response from the annotator into.
type GeoData struct {
	Geo *GeolocationIP // Holds the geolocation data
	ASN *IPASNData     // Holds the IP/ASN data
}

// The RequestData schema is the schema for the json that we will send
// down the pipe to the annotation service.
type RequestData struct {
	IP        string    // Holds the IP from an incoming request
	IPFormat  int       // Holds the ip format, 4 or 6
	Timestamp time.Time // Holds the timestamp from an incoming request
}

// AnnotatorURL holds the https address of the annotator.
// TODO(gfr) See if there is a better way of determining
// where to send the request (there almost certainly is)
var AnnotatorURL = "https://annotator-dot-" +
	os.Getenv("GCLOUD_PROJECT") +
	".appspot.com"

// BaseURL provides the base URL for single annotation requests
var BaseURL = AnnotatorURL + "/annotate?"

// BatchURL provides the base URL for batch annotation requests
var BatchURL = AnnotatorURL + "/batch_annotate"

// FetchGeoAnnotations takes a slice of strings
// containing ip addresses, a timestamp, and a slice of pointers to
// the GeolocationIP structs that correspond to the ip addresses. A
// precondition assumed by this function is that both slices are the
// same length. It will then make a call to the batch annotator, using
// the ip addresses and the timestamp. Then, it uses that data to fill
// in the structs pointed to by the slice of GeolocationIP pointers.
func FetchGeoAnnotations(ips []string, timestamp time.Time, geoDest []*GeolocationIP) {
	reqData := make([]RequestData, 0, len(ips))
	for _, ip := range ips {
		if ip == "" {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Empty IP Address!!!"}).Inc()
			continue
		}
		ip, _ = web100.NormalizeIPv6(ip)
		reqData = append(reqData, RequestData{ip, 0, timestamp})
	}
	annotationData := GetBatchGeoData(BatchURL, reqData)
	timeString := strconv.FormatInt(timestamp.Unix(), 36)
	for index, ip := range ips {
		data, ok := annotationData[ip+timeString]
		if !ok || data.Geo == nil {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Missing or empty data for IP Address!!!"}).Inc()
			continue
		}
		*geoDest[index] = *data.Geo

	}
}

// GetAndInsertGeolocationIPStruct takes a NON-NIL pointer to a
// pre-allocated GeolocationIP struct, an IP address, and a
// timestamp. It will connect to the annotation service, get the
// geo data, and insert the geo data into the reigion pointed to by
// the GeolocationIP pointer.
func GetAndInsertGeolocationIPStruct(geo *GeolocationIP, ip string, timestamp time.Time) {
	url := BaseURL + "ip_addr=" + url.QueryEscape(ip) +
		"&since_epoch=" + strconv.FormatInt(timestamp.Unix(), 10)
	annotationData := GetGeoData(url)
	if annotationData != nil && annotationData.Geo != nil {
		*geo = *annotationData.Geo
	}
}

// GetGeoData combines the functionality of QueryAnnotationService and
// ParseJSONGeoDataResponse to query the annotator service and return
// the corresponding GeoData if it can, or a nil pointer if it
// encounters any error and cannot get the data for any reason
func GetGeoData(url string) *GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := QueryAnnotationService(url)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := ParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		return nil
	}
	return geoDataFromResponse
}

// QueryAnnotationService will connect to the annotation service and
// copy the body of a valid response to a byte slice and return it to a
// user, returning an error if any occurs
func QueryAnnotationService(url string) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()
	// Make the actual request
	resp, err := http.Get(url)

	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Request to Annotator failed"}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Webserver gave non-ok response"}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// ParseJSONGeoDataResponse takes a byte slice containing the test of
// the JSON from the annotator service and parses it into a GeoData
// struct, for easy manipulation. It returns a pointer to the struct on
// success and an error if an error occurs.
func ParseJSONGeoDataResponse(jsonBuffer []byte) (*GeoData, error) {
	parsedJSON := &GeoData{}
	err := json.Unmarshal(jsonBuffer, parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}

// GetBatchGeoData combines the functionality of
// BatchQueryAnnotationService and BatchParseJSONGeoDataResponse to
// query the annotator service and return the corresponding map of
// ip-timestamp strings to GeoData structs, or a nil map if it
// encounters any error and cannot get the data for any reason
// TODO - dedup common code in GetGeoData
func GetBatchGeoData(url string, data []RequestData) map[string]GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := BatchQueryAnnotationService(url, data)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := BatchParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		log.Printf("%+v\n", data)
		return nil
	}
	return geoDataFromResponse
}

// BatchQueryAnnotationService takes a url to POST the request to and
// a slice of RequestDatas to be sent in the body in a JSON
// format. It will copy the response into a []byte and return it to
// the user, returning an error if any occurs
// TODO(gfr) Should pass the annotator's request context through and use it here.
func BatchQueryAnnotationService(url string, data []RequestData) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()

	encodedData, err := json.Marshal(data)
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Couldn't Marshal Data"}).Inc()
		return nil, err
	}

	var netClient = &http.Client{
		// Median response time is < 10 msec, but 99th percentile is 0.6 seconds.
		Timeout: 2 * time.Second,
	}

	// Make the actual request
	resp, err := netClient.Post(url, "raw", bytes.NewReader(encodedData))
	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": err.Error()}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": http.StatusText(resp.StatusCode)}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// BatchParseJSONGeoDataResponse takes a byte slice containing the
// text of the JSON from the annoator service's batch request endpoint
// and parses it into a map of strings to GeoData structs, for
// easy manipulation. It returns a pointer to the struct on success
// and an error if one occurs.
// TODO - is there duplicate code with ParseJSON... ?
func BatchParseJSONGeoDataResponse(jsonBuffer []byte) (map[string]GeoData, error) {
	parsedJSON := make(map[string]GeoData)
	err := json.Unmarshal(jsonBuffer, &parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}
<|endoftext|>


<commit_before>package dbutil

const dbconversion = `
-- Properties Needed for golang template
-- DBVersion, string, the current DB version or "00000000" if none
-- DBType, string, postgres
-- DroppingTables, boolean, should we drop old tables?

{{/* These variables need to be defined for all databases */}}

-- {{.DBType}} specific template features
-- Primary Key Expression: {{.pkey_exp}}

{{if eq .DBVersion "00000000"}}

CREATE TABLE PhoneCarrier (
	    id {{.pkey_exp}},
	    name VARCHAR UNIQUE NOT NULL,
	    emaildomain VARCHAR UNIQUE NOT NULL);

CREATE TABLE IF NOT EXISTS Users (
    Id {{.pkey_exp}},
	Name VARCHAR UNIQUE NOT NULL,
	Email VARCHAR UNIQUE NOT NULL,
	Password VARCHAR NOT NULL,
	PasswordSalt VARCHAR NOT NULL,
	PasswordHashScheme VARCHAR NOT NULL,
	Admin BOOLEAN DEFAULT FALSE,
	Phone VARCHAR DEFAULT '',
	PhoneCarrier INTEGER DEFAULT 0,
	UploadLimit_Items INTEGER DEFAULT 24000,
	ProcessingLimit_S INTEGER DEFAULT 86400,
	StorageLimit_Gb INTEGER DEFAULT 4,
	CreateTime INTEGER DEFAULT 0,
	ModifyTime INTEGER DEFAULT 0,
	UserGroup INTEGER DEFAULT 0,
	FOREIGN KEY(PhoneCarrier) REFERENCES PhoneCarrier(Id) ON DELETE SET NULL);

CREATE UNIQUE INDEX UserNameIndex ON Users (Name);

CREATE TABLE IF NOT EXISTS Device (
    Id {{.pkey_exp}},
	Name VARCHAR NOT NULL,
	ApiKey VARCHAR UNIQUE NOT NULL,
	Enabled BOOLEAN DEFAULT TRUE,
	Icon_PngB64 VARCHAR DEFAULT '',
	Shortname VARCHAR DEFAULT '',
	Superdevice BOOLEAN DEFAULT FALSE,
	OwnerId INTEGER,

	CanWrite BOOLEAN DEFAULT TRUE,
	CanWriteAnywhere BOOLEAN DEFAULT TRUE,
	UserProxy BOOLEAN DEFAULT FALSE,

	UNIQUE(Name, OwnerId),
	FOREIGN KEY(OwnerId) REFERENCES Users(Id) ON DELETE CASCADE
	);

CREATE UNIQUE INDEX DeviceNameIndex ON Device (Name);
CREATE UNIQUE INDEX DeviceAPIIndex ON Device (ApiKey);
CREATE INDEX DeviceOwnerIndex ON Device (OwnerId);

CREATE TABLE Stream (
    Id {{.pkey_exp}},
    Name VARCHAR NOT NULL,
    Active BOOLEAN DEFAULT TRUE,
    Public BOOLEAN DEFAULT FALSE,
    Type VARCHAR NOT NULL,
    OwnerId INTEGER,
    Ephemeral BOOLEAN DEFAULT FALSE,
    Output BOOLEAN DEFAULT FALSE,
    UNIQUE(Name, OwnerId),
    FOREIGN KEY(OwnerId) REFERENCES Device(Id) ON DELETE CASCADE
    );

CREATE INDEX StreamNameIndex ON Stream (Name);
CREATE INDEX StreamOwnerIndex ON Stream (OwnerId);

CREATE TABLE IF NOT EXISTS datastream (
	StreamId BIGINT NOT NULL,
	Substream VARCHAR,
	EndTime DOUBLE PRECISION,
	EndIndex BIGINT,
	Version INTEGER,
	Data BYTEA,
	UNIQUE (StreamId, Substream, EndIndex),
	PRIMARY KEY (StreamId, Substream, EndIndex)
);

--Index creation should only be run once.
CREATE INDEX datastreamtime ON datastream (StreamId,Substream,EndTime ASC);

{{end}}

{{/*========================================================================*/}}
{{/*Changes: updates to all tables in the database for newer schemas*/}}
{{/*========================================================================*/}}

{{if lt .DBVersion "20150328"}}

-- POSTGRES 9.1 should work with If Not Exists

-- This table won't exist for the first one.
CREATE TABLE IF NOT EXISTS StreamdbMeta (
     Key VARCHAR UNIQUE NOT NULL,
     Value VARCHAR NOT NULL);

CREATE INDEX sdb_meta ON StreamdbMeta (Key);

-- If we use this format date, we can just test < using lexocographic comparisons
INSERT INTO StreamdbMeta VALUES ('DBVersion', '20150328');

-- Rename the tables to their temporary alternatives
-- This keeps them from cascading deletes until the end.

ALTER TABLE PhoneCarrier RENAME TO PhoneCarriers;
ALTER TABLE Users RENAME TO Users20150328;
ALTER TABLE Device RENAME TO Devices20150328;
ALTER TABLE Stream RENAME TO Streams20150328;

-- We don't use the INE because we need to fail if these exist

CREATE TABLE Users (
    UserId {{.pkey_exp}},
	Name VARCHAR UNIQUE NOT NULL,
	Email VARCHAR UNIQUE NOT NULL,

	Password VARCHAR NOT NULL,
	PasswordSalt VARCHAR NOT NULL,
	PasswordHashScheme VARCHAR NOT NULL,

	Admin BOOLEAN DEFAULT FALSE,

	UploadLimit_Items INTEGER DEFAULT 24000,
	ProcessingLimit_S INTEGER DEFAULT 86400,
	StorageLimit_Gb INTEGER DEFAULT 4);

CREATE TABLE Devices (
    DeviceId {{.pkey_exp}},
    Name VARCHAR NOT NULL,
    Nickname VARCHAR DEFAULT '',
    UserId INTEGER,
    ApiKey VARCHAR UNIQUE NOT NULL,
    Enabled BOOLEAN DEFAULT TRUE,
    IsAdmin BOOLEAN DEFAULT FALSE,
    CanWrite BOOLEAN DEFAULT TRUE,
    CanWriteAnywhere BOOLEAN DEFAULT FALSE,
    CanActAsUser BOOLEAN DEFAULT FALSE,
    IsVisible BOOLEAN DEFAULT TRUE,
    UserEditable BOOLEAN DEFAULT TRUE,
    UNIQUE(UserId, Name),
    FOREIGN KEY(UserId) REFERENCES Users(UserId) ON DELETE CASCADE);

CREATE TABLE Streams (
    StreamId {{.pkey_exp}},
    Name VARCHAR NOT NULL,
    Nickname VARCHAR NOT NULL DEFAULT '',
    Type VARCHAR NOT NULL,
    DeviceId INTEGER,
    Ephemeral BOOLEAN DEFAULT FALSE,
    Downlink BOOLEAN DEFAULT FALSE,
    UNIQUE(Name, DeviceId),
    FOREIGN KEY(DeviceId) REFERENCES Devices(DeviceId) ON DELETE CASCADE);

CREATE TABLE UserKeyValues (
    UserId INTEGER,
    Key VARCHAR NOT NULL,
    Value VARCHAR NOT NULL DEFAULT '',
    FOREIGN KEY(UserId) REFERENCES Users(UserId) ON DELETE CASCADE,
    UNIQUE(UserId, Key),
    PRIMARY KEY (UserId, Key)
);

CREATE TABLE DeviceKeyValues (
    DeviceId INTEGER,
    Key VARCHAR NOT NULL,
    Value VARCHAR NOT NULL DEFAULT '',
    FOREIGN KEY(DeviceId) REFERENCES Devices(DeviceId) ON DELETE CASCADE,
    UNIQUE(DeviceId, Key),
    PRIMARY KEY (DeviceId, Key)
);

CREATE TABLE StreamKeyValues (
    StreamId INTEGER,
    Key VARCHAR NOT NULL,
    Value VARCHAR NOT NULL DEFAULT '',
    FOREIGN KEY(StreamId) REFERENCES Streams(StreamId) ON DELETE CASCADE,
    UNIQUE(StreamId, Key),
    PRIMARY KEY (StreamId, Key)
);

{{ if eq .DBType "postgres"}}

CREATE FUNCTION AddUserdev20150328Func() RETURNS TRIGGER AS $_$
BEGIN
	INSERT INTO Devices (Name, UserId, ApiKey, CanActAsUser, UserEditable, IsAdmin) VALUES ('user', NEW.UserId, NEW.Name || '-' || NEW.PasswordSalt, TRUE, FALSE, NEW.Admin);
    RETURN NEW;
END $_$ LANGUAGE 'plpgsql';

CREATE TRIGGER AddUserdev20150328 AFTER INSERT ON Users FOR EACH ROW
    EXECUTE PROCEDURE AddUserdev20150328Func();

CREATE FUNCTION ModifyUserdev20150328Func() RETURNS TRIGGER AS $_$
BEGIN
	UPDATE Devices SET IsAdmin = NEW.Admin WHERE UserId = NEW.UserId AND IsAdmin = TRUE;
	UPDATE Devices SET IsAdmin = NEW.Admin WHERE UserId = NEW.UserId AND Name = 'user';
    RETURN NEW;
END $_$ LANGUAGE 'plpgsql';

CREATE TRIGGER ModifyUserdev20150328 AFTER UPDATE ON Users FOR EACH ROW
    EXECUTE PROCEDURE ModifyUserdev20150328Func();

{{end}}

-- Construct Indexes

CREATE UNIQUE INDEX UserNameIndex20150328 ON Users (Name);
CREATE INDEX DeviceNameIndex20150328 ON Devices (Name);
CREATE UNIQUE INDEX DeviceAPIIndex20150328 ON Devices (ApiKey);
CREATE INDEX DeviceOwnerIndex20150328 ON Devices (UserId);
CREATE INDEX StreamNameIndex20150328 ON Streams (Name);
CREATE INDEX StreamOwnerIndex20150328 ON Streams (DeviceId);

-- Transfer Data

INSERT INTO Users SELECT
    Id,
    Name,
    Email,
    Password,
    PasswordSalt,
    PasswordHashScheme,
    Admin,
    UploadLimit_Items,
    ProcessingLimit_S,
    StorageLimit_Gb FROM Users20150328;

INSERT INTO Devices (DeviceId, Name, Nickname,  UserId,  ApiKey, Enabled, IsAdmin,   CanWrite, CanWriteAnywhere, CanActAsUser)
              SELECT Id,       Name, Shortname, OwnerId, ApiKey, Enabled, Superdevice, CanWrite, CanWriteAnywhere, UserProxy FROM Devices20150328;

INSERT INTO Streams (StreamId, Name, Type, DeviceId, Ephemeral, Downlink) SELECT Id, Name, Type, OwnerId, Ephemeral, Output FROM Streams20150328;

-- Insert Data

-- Default Carriers
-- do one by one because some databases don't like multiple
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('Alltel US', '@message.alltel.com');
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('AT&T US', '@txt.att.net');
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('Nextel US', '@messaging.nextel.com');
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('T-mobile US', '@tmomail.net');
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('Verizon US', '@vtext.com');

{{if .DroppingTables}}
-- Comment these out if there's an issue.
DROP TABLE Streams20150328;
DROP TABLE Devices20150328;
DROP TABLE Users20150328;
{{end}}

{{end}}

{{/*========================================================================

Changelog: 2015062826

* drop the unused phone carriers table
* drop the key/value tables (we don't want to be storing random people's data)
* create a log stream when the user is created
* create a table for defunct stream ids. These need to be manually scanned and
  deleted from redis.

========================================================================*/}}

{{if lt .DBVersion "2015062826"}}

UPDATE StreamdbMeta SET Value = '2015062826' WHERE Key = 'DBVersion';

DROP TABLE PhoneCarriers;
DROP TABLE StreamKeyValues;
DROP TABLE UserKeyValues;
DROP TABLE DeviceKeyValues;

-- Create a new trigger for inserting a device that is for a user.

DROP TRIGGER AddUserdev20150328 ON Users;

CREATE FUNCTION initial_user_setup() RETURNS TRIGGER AS $_$
DECLARE
	var_deviceid INTEGER;
BEGIN
	INSERT INTO Devices (Name, UserId, ApiKey, CanActAsUser, UserEditable, IsAdmin)
	    VALUES ('user', NEW.UserId, NEW.Name || '-' || NEW.PasswordSalt, TRUE, FALSE, NEW.Admin);

	SELECT DeviceId INTO var_deviceid FROM Devices
	    WHERE UserId = NEW.UserId AND Name = 'user';

	INSERT INTO Streams (Name, Type, DeviceId)
		VALUES ('log',
			'{"type": "object", "properties": {"cmd": {"type": "string"},"arg": {"type": "string"}},"required": ["cmd","arg"]}',
			var_deviceid);

	RETURN NEW;
END $_$ LANGUAGE 'plpgsql';

CREATE TRIGGER initialize_user AFTER INSERT ON Users FOR EACH ROW
    EXECUTE PROCEDURE initial_user_setup();

-- Now update the old user Devices

INSERT INTO Streams (Name, Type, DeviceId)
    SELECT 'log' as "Name",
		'{"type": "object", "properties": {"cmd": {"type": "string"},"arg": {"type": "string"}},"required": ["cmd","arg"]}' as "Type",
		DeviceId FROM Devices d WHERE d.Name = 'user';

-- Update users to add a nickname
ALTER TABLE Users
	ADD COLUMN nickname VARCHAR;

-- Update tables to use a random id

CREATE FUNCTION permuteQPR(x BIGINT) RETURNS INTEGER AS $$
DECLARE
	prime 	INTEGER;
	residue	BIGINT;
BEGIN
	-- see:
	-- http://preshing.com/20121224/how-to-generate-a-sequence-of-unique-random-integers/
	-- for more information on these calculations
	prime = 2147483423; -- congruence to 3 % 4 holds to ensure 1:1 mapping

	IF x >= prime THEN
		RETURN x; -- last 5 digits map to themselves
	ELSE
		residue = (x * x) % prime;
		IF residue <= prime / 2 THEN
			RETURN residue;
		ELSE
			RETURN (prime - residue);
		END IF;
	END IF;
END
$$ LANGUAGE 'plpgsql';

CREATE FUNCTION id_scramble(id INTEGER) RETURNS INTEGER AS $$
DECLARE
	xor   	INTEGER;
	add   	INTEGER;
BEGIN
	xor   = 0x5bf03635;
	add   = 0xDEADBEEF;
	RETURN permuteQPR(permuteQPR(id) + add) # xor;
END
$$ LANGUAGE 'plpgsql';

CREATE FUNCTION userid_scramble()
RETURNS trigger AS $$
BEGIN
	NEW.UserId = id_scramble(NEW.UserId);
	RETURN NEW;
END$$ LANGUAGE 'plpgsql';

CREATE TRIGGER userid_scramble_trigger
	BEFORE INSERT ON Users
	FOR EACH ROW
	EXECUTE PROCEDURE userid_scramble();

CREATE FUNCTION deviceid_scramble()
RETURNS trigger AS $$
BEGIN
	NEW.DeviceId = id_scramble(NEW.DeviceId);
	RETURN NEW;
END$$ LANGUAGE 'plpgsql';

CREATE TRIGGER deviceid_scramble_trigger
	BEFORE INSERT ON Devices
	FOR EACH ROW
	EXECUTE PROCEDURE deviceid_scramble();

CREATE FUNCTION streamid_scramble()
RETURNS trigger AS $$
BEGIN
	NEW.StreamId = id_scramble(NEW.StreamId);
	RETURN NEW;
END$$ LANGUAGE 'plpgsql';

CREATE TRIGGER streamid_scramble_trigger
	BEFORE INSERT ON Streams
	FOR EACH ROW
	EXECUTE PROCEDURE streamid_scramble();

{{/* End 2015062826 */}}
{{end}}

`
<commit_msg>Added metadevice to conversionscript<commit_after>package dbutil

const dbconversion = `
-- Properties Needed for golang template
-- DBVersion, string, the current DB version or "00000000" if none
-- DBType, string, postgres
-- DroppingTables, boolean, should we drop old tables?

{{/* These variables need to be defined for all databases */}}

-- {{.DBType}} specific template features
-- Primary Key Expression: {{.pkey_exp}}

{{if eq .DBVersion "00000000"}}

CREATE TABLE PhoneCarrier (
	    id {{.pkey_exp}},
	    name VARCHAR UNIQUE NOT NULL,
	    emaildomain VARCHAR UNIQUE NOT NULL);

CREATE TABLE IF NOT EXISTS Users (
    Id {{.pkey_exp}},
	Name VARCHAR UNIQUE NOT NULL,
	Email VARCHAR UNIQUE NOT NULL,
	Password VARCHAR NOT NULL,
	PasswordSalt VARCHAR NOT NULL,
	PasswordHashScheme VARCHAR NOT NULL,
	Admin BOOLEAN DEFAULT FALSE,
	Phone VARCHAR DEFAULT '',
	PhoneCarrier INTEGER DEFAULT 0,
	UploadLimit_Items INTEGER DEFAULT 24000,
	ProcessingLimit_S INTEGER DEFAULT 86400,
	StorageLimit_Gb INTEGER DEFAULT 4,
	CreateTime INTEGER DEFAULT 0,
	ModifyTime INTEGER DEFAULT 0,
	UserGroup INTEGER DEFAULT 0,
	FOREIGN KEY(PhoneCarrier) REFERENCES PhoneCarrier(Id) ON DELETE SET NULL);

CREATE UNIQUE INDEX UserNameIndex ON Users (Name);

CREATE TABLE IF NOT EXISTS Device (
    Id {{.pkey_exp}},
	Name VARCHAR NOT NULL,
	ApiKey VARCHAR UNIQUE NOT NULL,
	Enabled BOOLEAN DEFAULT TRUE,
	Icon_PngB64 VARCHAR DEFAULT '',
	Shortname VARCHAR DEFAULT '',
	Superdevice BOOLEAN DEFAULT FALSE,
	OwnerId INTEGER,

	CanWrite BOOLEAN DEFAULT TRUE,
	CanWriteAnywhere BOOLEAN DEFAULT TRUE,
	UserProxy BOOLEAN DEFAULT FALSE,

	UNIQUE(Name, OwnerId),
	FOREIGN KEY(OwnerId) REFERENCES Users(Id) ON DELETE CASCADE
	);

CREATE UNIQUE INDEX DeviceNameIndex ON Device (Name);
CREATE UNIQUE INDEX DeviceAPIIndex ON Device (ApiKey);
CREATE INDEX DeviceOwnerIndex ON Device (OwnerId);

CREATE TABLE Stream (
    Id {{.pkey_exp}},
    Name VARCHAR NOT NULL,
    Active BOOLEAN DEFAULT TRUE,
    Public BOOLEAN DEFAULT FALSE,
    Type VARCHAR NOT NULL,
    OwnerId INTEGER,
    Ephemeral BOOLEAN DEFAULT FALSE,
    Output BOOLEAN DEFAULT FALSE,
    UNIQUE(Name, OwnerId),
    FOREIGN KEY(OwnerId) REFERENCES Device(Id) ON DELETE CASCADE
    );

CREATE INDEX StreamNameIndex ON Stream (Name);
CREATE INDEX StreamOwnerIndex ON Stream (OwnerId);

CREATE TABLE IF NOT EXISTS datastream (
	StreamId BIGINT NOT NULL,
	Substream VARCHAR,
	EndTime DOUBLE PRECISION,
	EndIndex BIGINT,
	Version INTEGER,
	Data BYTEA,
	UNIQUE (StreamId, Substream, EndIndex),
	PRIMARY KEY (StreamId, Substream, EndIndex)
);

--Index creation should only be run once.
CREATE INDEX datastreamtime ON datastream (StreamId,Substream,EndTime ASC);

{{end}}

{{/*========================================================================*/}}
{{/*Changes: updates to all tables in the database for newer schemas*/}}
{{/*========================================================================*/}}

{{if lt .DBVersion "20150328"}}

-- POSTGRES 9.1 should work with If Not Exists

-- This table won't exist for the first one.
CREATE TABLE IF NOT EXISTS StreamdbMeta (
     Key VARCHAR UNIQUE NOT NULL,
     Value VARCHAR NOT NULL);

CREATE INDEX sdb_meta ON StreamdbMeta (Key);

-- If we use this format date, we can just test < using lexocographic comparisons
INSERT INTO StreamdbMeta VALUES ('DBVersion', '20150328');

-- Rename the tables to their temporary alternatives
-- This keeps them from cascading deletes until the end.

ALTER TABLE PhoneCarrier RENAME TO PhoneCarriers;
ALTER TABLE Users RENAME TO Users20150328;
ALTER TABLE Device RENAME TO Devices20150328;
ALTER TABLE Stream RENAME TO Streams20150328;

-- We don't use the INE because we need to fail if these exist

CREATE TABLE Users (
    UserId {{.pkey_exp}},
	Name VARCHAR UNIQUE NOT NULL,
	Email VARCHAR UNIQUE NOT NULL,

	Password VARCHAR NOT NULL,
	PasswordSalt VARCHAR NOT NULL,
	PasswordHashScheme VARCHAR NOT NULL,

	Admin BOOLEAN DEFAULT FALSE,

	UploadLimit_Items INTEGER DEFAULT 24000,
	ProcessingLimit_S INTEGER DEFAULT 86400,
	StorageLimit_Gb INTEGER DEFAULT 4);

CREATE TABLE Devices (
    DeviceId {{.pkey_exp}},
    Name VARCHAR NOT NULL,
    Nickname VARCHAR DEFAULT '',
    UserId INTEGER,
    ApiKey VARCHAR UNIQUE NOT NULL,
    Enabled BOOLEAN DEFAULT TRUE,
    IsAdmin BOOLEAN DEFAULT FALSE,
    CanWrite BOOLEAN DEFAULT TRUE,
    CanWriteAnywhere BOOLEAN DEFAULT FALSE,
    CanActAsUser BOOLEAN DEFAULT FALSE,
    IsVisible BOOLEAN DEFAULT TRUE,
    UserEditable BOOLEAN DEFAULT TRUE,
    UNIQUE(UserId, Name),
    FOREIGN KEY(UserId) REFERENCES Users(UserId) ON DELETE CASCADE);

CREATE TABLE Streams (
    StreamId {{.pkey_exp}},
    Name VARCHAR NOT NULL,
    Nickname VARCHAR NOT NULL DEFAULT '',
    Type VARCHAR NOT NULL,
    DeviceId INTEGER,
    Ephemeral BOOLEAN DEFAULT FALSE,
    Downlink BOOLEAN DEFAULT FALSE,
    UNIQUE(Name, DeviceId),
    FOREIGN KEY(DeviceId) REFERENCES Devices(DeviceId) ON DELETE CASCADE);

CREATE TABLE UserKeyValues (
    UserId INTEGER,
    Key VARCHAR NOT NULL,
    Value VARCHAR NOT NULL DEFAULT '',
    FOREIGN KEY(UserId) REFERENCES Users(UserId) ON DELETE CASCADE,
    UNIQUE(UserId, Key),
    PRIMARY KEY (UserId, Key)
);

CREATE TABLE DeviceKeyValues (
    DeviceId INTEGER,
    Key VARCHAR NOT NULL,
    Value VARCHAR NOT NULL DEFAULT '',
    FOREIGN KEY(DeviceId) REFERENCES Devices(DeviceId) ON DELETE CASCADE,
    UNIQUE(DeviceId, Key),
    PRIMARY KEY (DeviceId, Key)
);

CREATE TABLE StreamKeyValues (
    StreamId INTEGER,
    Key VARCHAR NOT NULL,
    Value VARCHAR NOT NULL DEFAULT '',
    FOREIGN KEY(StreamId) REFERENCES Streams(StreamId) ON DELETE CASCADE,
    UNIQUE(StreamId, Key),
    PRIMARY KEY (StreamId, Key)
);

{{ if eq .DBType "postgres"}}

CREATE FUNCTION AddUserdev20150328Func() RETURNS TRIGGER AS $_$
BEGIN
	INSERT INTO Devices (Name, UserId, ApiKey, CanActAsUser, UserEditable, IsAdmin) VALUES ('user', NEW.UserId, NEW.Name || '-' || NEW.PasswordSalt, TRUE, FALSE, NEW.Admin);
	INSERT INTO Devices (Name, UserId, ApiKey, CanActAsUser, UserEditable, IsAdmin) VALUES ('meta', NEW.UserId, '', TRUE, FALSE, FALSE);
    RETURN NEW;
END $_$ LANGUAGE 'plpgsql';

CREATE TRIGGER AddUserdev20150328 AFTER INSERT ON Users FOR EACH ROW
    EXECUTE PROCEDURE AddUserdev20150328Func();

CREATE FUNCTION ModifyUserdev20150328Func() RETURNS TRIGGER AS $_$
BEGIN
	UPDATE Devices SET IsAdmin = NEW.Admin WHERE UserId = NEW.UserId AND IsAdmin = TRUE;
	UPDATE Devices SET IsAdmin = NEW.Admin WHERE UserId = NEW.UserId AND Name = 'user';
    RETURN NEW;
END $_$ LANGUAGE 'plpgsql';

CREATE TRIGGER ModifyUserdev20150328 AFTER UPDATE ON Users FOR EACH ROW
    EXECUTE PROCEDURE ModifyUserdev20150328Func();

{{end}}

-- Construct Indexes

CREATE UNIQUE INDEX UserNameIndex20150328 ON Users (Name);
CREATE INDEX DeviceNameIndex20150328 ON Devices (Name);
CREATE UNIQUE INDEX DeviceAPIIndex20150328 ON Devices (ApiKey);
CREATE INDEX DeviceOwnerIndex20150328 ON Devices (UserId);
CREATE INDEX StreamNameIndex20150328 ON Streams (Name);
CREATE INDEX StreamOwnerIndex20150328 ON Streams (DeviceId);

-- Transfer Data

INSERT INTO Users SELECT
    Id,
    Name,
    Email,
    Password,
    PasswordSalt,
    PasswordHashScheme,
    Admin,
    UploadLimit_Items,
    ProcessingLimit_S,
    StorageLimit_Gb FROM Users20150328;

INSERT INTO Devices (DeviceId, Name, Nickname,  UserId,  ApiKey, Enabled, IsAdmin,   CanWrite, CanWriteAnywhere, CanActAsUser)
              SELECT Id,       Name, Shortname, OwnerId, ApiKey, Enabled, Superdevice, CanWrite, CanWriteAnywhere, UserProxy FROM Devices20150328;

INSERT INTO Streams (StreamId, Name, Type, DeviceId, Ephemeral, Downlink) SELECT Id, Name, Type, OwnerId, Ephemeral, Output FROM Streams20150328;

-- Insert Data

-- Default Carriers
-- do one by one because some databases don't like multiple
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('Alltel US', '@message.alltel.com');
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('AT&T US', '@txt.att.net');
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('Nextel US', '@messaging.nextel.com');
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('T-mobile US', '@tmomail.net');
INSERT INTO PhoneCarriers (name, emaildomain) VALUES ('Verizon US', '@vtext.com');

{{if .DroppingTables}}
-- Comment these out if there's an issue.
DROP TABLE Streams20150328;
DROP TABLE Devices20150328;
DROP TABLE Users20150328;
{{end}}

{{end}}

{{/*========================================================================

Changelog: 2015062826

* drop the unused phone carriers table
* drop the key/value tables (we don't want to be storing random people's data)
* create a log stream when the user is created
* create a table for defunct stream ids. These need to be manually scanned and
  deleted from redis.

========================================================================*/}}

{{if lt .DBVersion "2015062826"}}

UPDATE StreamdbMeta SET Value = '2015062826' WHERE Key = 'DBVersion';

DROP TABLE PhoneCarriers;
DROP TABLE StreamKeyValues;
DROP TABLE UserKeyValues;
DROP TABLE DeviceKeyValues;

-- Create a new trigger for inserting a device that is for a user.

DROP TRIGGER AddUserdev20150328 ON Users;

CREATE FUNCTION initial_user_setup() RETURNS TRIGGER AS $_$
DECLARE
	var_deviceid INTEGER;
BEGIN
	INSERT INTO Devices (Name, UserId, ApiKey, CanActAsUser, UserEditable, IsAdmin)
	    VALUES ('user', NEW.UserId, NEW.Name || '-' || NEW.PasswordSalt, TRUE, FALSE, NEW.Admin);

	SELECT DeviceId INTO var_deviceid FROM Devices
	    WHERE UserId = NEW.UserId AND Name = 'user';

	INSERT INTO Streams (Name, Type, DeviceId)
		VALUES ('log',
			'{"type": "object", "properties": {"cmd": {"type": "string"},"arg": {"type": "string"}},"required": ["cmd","arg"]}',
			var_deviceid);

	RETURN NEW;
END $_$ LANGUAGE 'plpgsql';

CREATE TRIGGER initialize_user AFTER INSERT ON Users FOR EACH ROW
    EXECUTE PROCEDURE initial_user_setup();

-- Now update the old user Devices

INSERT INTO Streams (Name, Type, DeviceId)
    SELECT 'log' as "Name",
		'{"type": "object", "properties": {"cmd": {"type": "string"},"arg": {"type": "string"}},"required": ["cmd","arg"]}' as "Type",
		DeviceId FROM Devices d WHERE d.Name = 'user';

-- Update users to add a nickname
ALTER TABLE Users
	ADD COLUMN nickname VARCHAR;

-- Update tables to use a random id

CREATE FUNCTION permuteQPR(x BIGINT) RETURNS INTEGER AS $$
DECLARE
	prime 	INTEGER;
	residue	BIGINT;
BEGIN
	-- see:
	-- http://preshing.com/20121224/how-to-generate-a-sequence-of-unique-random-integers/
	-- for more information on these calculations
	prime = 2147483423; -- congruence to 3 % 4 holds to ensure 1:1 mapping

	IF x >= prime THEN
		RETURN x; -- last 5 digits map to themselves
	ELSE
		residue = (x * x) % prime;
		IF residue <= prime / 2 THEN
			RETURN residue;
		ELSE
			RETURN (prime - residue);
		END IF;
	END IF;
END
$$ LANGUAGE 'plpgsql';

CREATE FUNCTION id_scramble(id INTEGER) RETURNS INTEGER AS $$
DECLARE
	xor   	INTEGER;
	add   	INTEGER;
BEGIN
	xor   = 0x5bf03635;
	add   = 0xDEADBEEF;
	RETURN permuteQPR(permuteQPR(id) + add) # xor;
END
$$ LANGUAGE 'plpgsql';

CREATE FUNCTION userid_scramble()
RETURNS trigger AS $$
BEGIN
	NEW.UserId = id_scramble(NEW.UserId);
	RETURN NEW;
END$$ LANGUAGE 'plpgsql';

CREATE TRIGGER userid_scramble_trigger
	BEFORE INSERT ON Users
	FOR EACH ROW
	EXECUTE PROCEDURE userid_scramble();

CREATE FUNCTION deviceid_scramble()
RETURNS trigger AS $$
BEGIN
	NEW.DeviceId = id_scramble(NEW.DeviceId);
	RETURN NEW;
END$$ LANGUAGE 'plpgsql';

CREATE TRIGGER deviceid_scramble_trigger
	BEFORE INSERT ON Devices
	FOR EACH ROW
	EXECUTE PROCEDURE deviceid_scramble();

CREATE FUNCTION streamid_scramble()
RETURNS trigger AS $$
BEGIN
	NEW.StreamId = id_scramble(NEW.StreamId);
	RETURN NEW;
END$$ LANGUAGE 'plpgsql';

CREATE TRIGGER streamid_scramble_trigger
	BEFORE INSERT ON Streams
	FOR EACH ROW
	EXECUTE PROCEDURE streamid_scramble();

{{/* End 2015062826 */}}
{{end}}

`
<|endoftext|>


<commit_before>// Package scipipe is a library for writing scientific workflows (sometimes
// also called "pipelines") of shell commands that depend on each other, in the
// Go programming languages. It was initially designed for problems in
// cheminformatics and bioinformatics, but should apply equally well to any
// domain involving complex pipelines of interdependent shell commands.
package scipipe

import (
	"sync"
)

// ----------------------------------------------------------------------------
// Workflow
// ----------------------------------------------------------------------------

// Workflow is the centerpiece of the functionality in SciPipe, and is a
// container for a pipeline of processes making up a workflow. It has various
// methods for coordination the execution of the pipeline as a whole, such as
// keeping track of the maxiumum number of concurrent tasks, as well as helper
// methods for creating new processes, that automatically gets plugged in to the
// workflow on creation
type Workflow struct {
	name              string
	procs             map[string]WorkflowProcess
	concurrentTasks   chan struct{}
	concurrentTasksMx sync.Mutex
	sink              *Sink
	driver            WorkflowProcess
}

// NewWorkflow returns a new Workflow
func NewWorkflow(name string, maxConcurrentTasks int) *Workflow {
	InitLogInfo()
	sink := NewSink(name + "_default_sink")
	return &Workflow{
		name:            name,
		procs:           map[string]WorkflowProcess{},
		concurrentTasks: make(chan struct{}, maxConcurrentTasks),
		sink:            sink,
		driver:          sink,
	}
}

// WorkflowProcess is an interface for processes to be handled by Workflow
type WorkflowProcess interface {
	Name() string
	InPorts() map[string]*InPort
	OutPorts() map[string]*OutPort
	ParamInPorts() map[string]*ParamInPort
	ParamOutPorts() map[string]*ParamOutPort
	Connected() bool
	Run()
}

// EmptyWorkflowProcess implements the WorkflowProcess interface in the simplest
// possible way, by just returning default values. It is supposed to be used to
// avoid re-implementing such "default" behaviour where needed, in new
// components
type EmptyWorkflowProcess struct{}

// InPorts returns an empty map of InPorts
func (ep *EmptyWorkflowProcess) InPorts() map[string]*InPort {
	return map[string]*InPort{}
}

// OutPorts returns an empty map of OutPorts
func (ep *EmptyWorkflowProcess) OutPorts() map[string]*OutPort {
	return map[string]*OutPort{}
}

// ParamInPorts returns an empty map of ParamInPorts
func (ep *EmptyWorkflowProcess) ParamInPorts() map[string]*ParamInPort {
	return map[string]*ParamInPort{}
}

// ParamOutPorts returns an empty map of ParamOutPorts
func (ep *EmptyWorkflowProcess) ParamOutPorts() map[string]*ParamOutPort {
	return map[string]*ParamOutPort{}
}

// AddProc adds a Process to the workflow, to be run when the workflow runs
func (wf *Workflow) AddProc(proc WorkflowProcess) {
	if wf.procs[proc.Name()] != nil {
		Error.Fatalf(wf.name+" workflow: A process with name '%s' already exists in the workflow! Use a more unique name!\n", proc.Name())
	}
	wf.procs[proc.Name()] = proc
}

// AddProcs takes one or many Processes and adds them to the workflow, to be run
// when the workflow runs.
func (wf *Workflow) AddProcs(procs ...WorkflowProcess) {
	for _, proc := range procs {
		wf.AddProc(proc)
	}
}

// NewProc returns a new process based on a commandPattern (See the
// documentation for scipipe.NewProcess for more details about the pattern) and
// connects the process to the workflow
func (wf *Workflow) NewProc(procName string, commandPattern string) *Process {
	proc := NewProc(wf, procName, commandPattern)
	return proc
}

// Proc returns the process with name procName from the workflow
func (wf *Workflow) Proc(procName string) WorkflowProcess {
	return wf.procs[procName]
}

// Procs returns a map of all processes keyed by their names in the workflow
func (wf *Workflow) Procs() map[string]WorkflowProcess {
	return wf.procs
}

// Sink returns the sink process of the workflow
func (wf *Workflow) Sink() *Sink {
	return wf.sink
}

// SetSink sets the sink of the workflow to the provided sink process
func (wf *Workflow) SetSink(sink *Sink) {
	if wf.sink.Connected() {
		Error.Fatalln("Trying to replace a sink which is already connected. Are you combining SetSink() with ConnectFinalOutPort()? That is not allowed!")
	}
	wf.sink = sink
}

// IncConcurrentTasks increases the conter for how many concurrent tasks are
// currently running in the workflow
func (wf *Workflow) IncConcurrentTasks(slots int) {
	// We must lock so that multiple processes don't end up with partially "filled slots"
	wf.concurrentTasksMx.Lock()
	for i := 0; i < slots; i++ {
		wf.concurrentTasks <- struct{}{}
		Debug.Println("Increased concurrent tasks")
	}
	wf.concurrentTasksMx.Unlock()
}

// DecConcurrentTasks decreases the conter for how many concurrent tasks are
// currently running in the workflow
func (wf *Workflow) DecConcurrentTasks(slots int) {
	for i := 0; i < slots; i++ {
		<-wf.concurrentTasks
		Debug.Println("Decreased concurrent tasks")
	}
}

func (wf *Workflow) readyToRun(procs map[string]WorkflowProcess) bool {
	if len(procs) == 0 {
		Error.Println(wf.name + ": The workflow is empty. Did you forget to add the processes to it?")
		return false
	}
	if wf.sink == nil {
		Error.Println(wf.name + ": sink is nil!")
		return false
	}
	for _, proc := range procs {
		if !proc.Connected() {
			Error.Println(wf.name + ": Not everything connected. Workflow shutting down.")
			return false
		}
	}
	return true
}

// runProcs runs a specified set of processes only
func (wf *Workflow) runProcs(procs map[string]WorkflowProcess) {
	wf.reconnectDeadEndConnections(procs)

	if !wf.readyToRun(procs) {
		Error.Fatalln("Workflow not ready to run, due to previously reported errors, so exiting.")
	}

	for _, proc := range procs {
		Debug.Printf(wf.name+": Starting process %s in new go-routine", proc.Name())
		go proc.Run()
	}

	Debug.Printf(wf.name + ": Starting driver process in main go-routine")
	wf.driver.Run()
}

// reconnectDeadEndConnections disonnects connections to processes which are not
// in the set of processes to be run, and, if an out-port for a process that is supposed to be
// run becomes disconnected, it is connected to the sink instead
func (wf *Workflow) reconnectDeadEndConnections(procs map[string]WorkflowProcess) {
	foundNewDriverProc := false

	for _, proc := range procs {
		// OutPorts
		for _, opt := range proc.OutPorts() {
			for iptName, ipt := range opt.RemotePorts {
				// If the remotely connected process is not among the ones to run ...
				if ipt.Process == nil {
					Debug.Printf("Disconnecting in-port %s from out-port %s", ipt.Name(), opt.Name())
					opt.Disconnect(iptName)
				} else if _, ok := procs[ipt.Process.Name()]; !ok {
					Debug.Printf("Disconnecting in-port %s from out-port %s", ipt.Name(), opt.Name())
					opt.Disconnect(iptName)
				}
			}
			if !opt.Connected() {
				Debug.Printf("Connecting disconnected out-port %s of process %s to workflow sink", opt.Name(), opt.Process.Name())
				wf.sink.Connect(opt)
			}
		}

		// ParamOutPorts
		for _, pop := range proc.ParamOutPorts() {
			for rppName, rpp := range pop.RemotePorts {
				// If the remotely connected process is not among the ones to run ...
				if rpp.Process == nil {
					Debug.Printf("Disconnecting in-port %s from out-port %s", rpp.Name(), pop.Name())
					pop.Disconnect(rppName)
				} else if _, ok := procs[rpp.Process.Name()]; !ok {
					Debug.Printf("Disconnecting in-port %s from out-port %s", rpp.Name(), pop.Name())
					pop.Disconnect(rppName)
				}
			}
			if !pop.Connected() {
				Debug.Printf("Connecting disconnected out-port %s of process %s to workflow sink", pop.Name(), pop.Process.Name())
				wf.sink.ConnectParam(pop)
			}
		}

		if len(proc.OutPorts()) == 0 && len(proc.ParamOutPorts()) == 0 {
			if foundNewDriverProc {
				Error.Fatalf("Found more than one process without out-ports nor out-param ports. Cannot use both as drivers (One of them being '%s'). Adapt your workflow accordingly.", proc.Name())
			}
			foundNewDriverProc = true
			wf.driver = proc
		}
	}
}

// RunProcsByName runs a specified set of processes only, specified by their
// names as strings
func (wf *Workflow) RunProcsByName(procNames ...string) {
	procs := map[string]WorkflowProcess{}
	for _, procName := range procNames {
		procs[procName] = wf.Proc(procName)
	}
	wf.runProcs(procs)
}

// RunToProcName runs all processes upstream of, and including, the process with
// name finalProcName
func (wf *Workflow) RunToProcName(finalProcName string) {
	if _, ok := wf.Procs()[finalProcName]; !ok {
		Error.Fatalf("No process with name '%s' in %s workflow. Please check the name supplied to RunToProcName() in your workflow, and try again", finalProcName, wf.Name())
	}
	wf.RunToProc(wf.Proc(finalProcName))
}

// RunToProc runs all processes upstream of, and including, the finalProc
func (wf *Workflow) RunToProc(finalProc WorkflowProcess) {
	procsToRun := upstreamProcsForProc(finalProc)
	procsToRun[finalProc.Name()] = finalProc
	wf.runProcs(procsToRun)
}

// upstreamProcsForProc returns all processes it is connected to, either
// directly or indirectly, via its in-ports and param-in-ports
func upstreamProcsForProc(proc WorkflowProcess) map[string]WorkflowProcess {
	procs := map[string]WorkflowProcess{}
	for _, inp := range proc.InPorts() {
		for _, rpt := range inp.RemotePorts {
			procs[rpt.Process.Name()] = rpt.Process
			mergeWFMaps(procs, upstreamProcsForProc(rpt.Process))
		}
	}
	for _, pip := range proc.ParamInPorts() {
		for _, rpp := range pip.RemotePorts {
			procs[rpp.Process.Name()] = rpp.Process
			mergeWFMaps(procs, upstreamProcsForProc(rpp.Process))
		}
	}
	return procs
}

func mergeWFMaps(a map[string]WorkflowProcess, b map[string]WorkflowProcess) map[string]WorkflowProcess {
	for k, v := range b {
		a[k] = v
	}
	return a
}

// Run runs the workflow
func (wf *Workflow) Run() {
	wf.runProcs(wf.procs)
}
<commit_msg>Add Name() method to workflow<commit_after>// Package scipipe is a library for writing scientific workflows (sometimes
// also called "pipelines") of shell commands that depend on each other, in the
// Go programming languages. It was initially designed for problems in
// cheminformatics and bioinformatics, but should apply equally well to any
// domain involving complex pipelines of interdependent shell commands.
package scipipe

import (
	"sync"
)

// ----------------------------------------------------------------------------
// Workflow
// ----------------------------------------------------------------------------

// Workflow is the centerpiece of the functionality in SciPipe, and is a
// container for a pipeline of processes making up a workflow. It has various
// methods for coordination the execution of the pipeline as a whole, such as
// keeping track of the maxiumum number of concurrent tasks, as well as helper
// methods for creating new processes, that automatically gets plugged in to the
// workflow on creation
type Workflow struct {
	name              string
	procs             map[string]WorkflowProcess
	concurrentTasks   chan struct{}
	concurrentTasksMx sync.Mutex
	sink              *Sink
	driver            WorkflowProcess
}

// NewWorkflow returns a new Workflow
func NewWorkflow(name string, maxConcurrentTasks int) *Workflow {
	InitLogInfo()
	sink := NewSink(name + "_default_sink")
	return &Workflow{
		name:            name,
		procs:           map[string]WorkflowProcess{},
		concurrentTasks: make(chan struct{}, maxConcurrentTasks),
		sink:            sink,
		driver:          sink,
	}
}

// WorkflowProcess is an interface for processes to be handled by Workflow
type WorkflowProcess interface {
	Name() string
	InPorts() map[string]*InPort
	OutPorts() map[string]*OutPort
	ParamInPorts() map[string]*ParamInPort
	ParamOutPorts() map[string]*ParamOutPort
	Connected() bool
	Run()
}

// EmptyWorkflowProcess implements the WorkflowProcess interface in the simplest
// possible way, by just returning default values. It is supposed to be used to
// avoid re-implementing such "default" behaviour where needed, in new
// components
type EmptyWorkflowProcess struct{}

// InPorts returns an empty map of InPorts
func (ep *EmptyWorkflowProcess) InPorts() map[string]*InPort {
	return map[string]*InPort{}
}

// OutPorts returns an empty map of OutPorts
func (ep *EmptyWorkflowProcess) OutPorts() map[string]*OutPort {
	return map[string]*OutPort{}
}

// ParamInPorts returns an empty map of ParamInPorts
func (ep *EmptyWorkflowProcess) ParamInPorts() map[string]*ParamInPort {
	return map[string]*ParamInPort{}
}

// ParamOutPorts returns an empty map of ParamOutPorts
func (ep *EmptyWorkflowProcess) ParamOutPorts() map[string]*ParamOutPort {
	return map[string]*ParamOutPort{}
}

func (wf *Workflow) Name() string {
	return wf.name
}

// AddProc adds a Process to the workflow, to be run when the workflow runs
func (wf *Workflow) AddProc(proc WorkflowProcess) {
	if wf.procs[proc.Name()] != nil {
		Error.Fatalf(wf.name+" workflow: A process with name '%s' already exists in the workflow! Use a more unique name!\n", proc.Name())
	}
	wf.procs[proc.Name()] = proc
}

// AddProcs takes one or many Processes and adds them to the workflow, to be run
// when the workflow runs.
func (wf *Workflow) AddProcs(procs ...WorkflowProcess) {
	for _, proc := range procs {
		wf.AddProc(proc)
	}
}

// NewProc returns a new process based on a commandPattern (See the
// documentation for scipipe.NewProcess for more details about the pattern) and
// connects the process to the workflow
func (wf *Workflow) NewProc(procName string, commandPattern string) *Process {
	proc := NewProc(wf, procName, commandPattern)
	return proc
}

// Proc returns the process with name procName from the workflow
func (wf *Workflow) Proc(procName string) WorkflowProcess {
	return wf.procs[procName]
}

// Procs returns a map of all processes keyed by their names in the workflow
func (wf *Workflow) Procs() map[string]WorkflowProcess {
	return wf.procs
}

// Sink returns the sink process of the workflow
func (wf *Workflow) Sink() *Sink {
	return wf.sink
}

// SetSink sets the sink of the workflow to the provided sink process
func (wf *Workflow) SetSink(sink *Sink) {
	if wf.sink.Connected() {
		Error.Fatalln("Trying to replace a sink which is already connected. Are you combining SetSink() with ConnectFinalOutPort()? That is not allowed!")
	}
	wf.sink = sink
}

// IncConcurrentTasks increases the conter for how many concurrent tasks are
// currently running in the workflow
func (wf *Workflow) IncConcurrentTasks(slots int) {
	// We must lock so that multiple processes don't end up with partially "filled slots"
	wf.concurrentTasksMx.Lock()
	for i := 0; i < slots; i++ {
		wf.concurrentTasks <- struct{}{}
		Debug.Println("Increased concurrent tasks")
	}
	wf.concurrentTasksMx.Unlock()
}

// DecConcurrentTasks decreases the conter for how many concurrent tasks are
// currently running in the workflow
func (wf *Workflow) DecConcurrentTasks(slots int) {
	for i := 0; i < slots; i++ {
		<-wf.concurrentTasks
		Debug.Println("Decreased concurrent tasks")
	}
}

func (wf *Workflow) readyToRun(procs map[string]WorkflowProcess) bool {
	if len(procs) == 0 {
		Error.Println(wf.name + ": The workflow is empty. Did you forget to add the processes to it?")
		return false
	}
	if wf.sink == nil {
		Error.Println(wf.name + ": sink is nil!")
		return false
	}
	for _, proc := range procs {
		if !proc.Connected() {
			Error.Println(wf.name + ": Not everything connected. Workflow shutting down.")
			return false
		}
	}
	return true
}

// runProcs runs a specified set of processes only
func (wf *Workflow) runProcs(procs map[string]WorkflowProcess) {
	wf.reconnectDeadEndConnections(procs)

	if !wf.readyToRun(procs) {
		Error.Fatalln("Workflow not ready to run, due to previously reported errors, so exiting.")
	}

	for _, proc := range procs {
		Debug.Printf(wf.name+": Starting process %s in new go-routine", proc.Name())
		go proc.Run()
	}

	Debug.Printf(wf.name + ": Starting driver process in main go-routine")
	wf.driver.Run()
}

// reconnectDeadEndConnections disonnects connections to processes which are not
// in the set of processes to be run, and, if an out-port for a process that is supposed to be
// run becomes disconnected, it is connected to the sink instead
func (wf *Workflow) reconnectDeadEndConnections(procs map[string]WorkflowProcess) {
	foundNewDriverProc := false

	for _, proc := range procs {
		// OutPorts
		for _, opt := range proc.OutPorts() {
			for iptName, ipt := range opt.RemotePorts {
				// If the remotely connected process is not among the ones to run ...
				if ipt.Process == nil {
					Debug.Printf("Disconnecting in-port %s from out-port %s", ipt.Name(), opt.Name())
					opt.Disconnect(iptName)
				} else if _, ok := procs[ipt.Process.Name()]; !ok {
					Debug.Printf("Disconnecting in-port %s from out-port %s", ipt.Name(), opt.Name())
					opt.Disconnect(iptName)
				}
			}
			if !opt.Connected() {
				Debug.Printf("Connecting disconnected out-port %s of process %s to workflow sink", opt.Name(), opt.Process.Name())
				wf.sink.Connect(opt)
			}
		}

		// ParamOutPorts
		for _, pop := range proc.ParamOutPorts() {
			for rppName, rpp := range pop.RemotePorts {
				// If the remotely connected process is not among the ones to run ...
				if rpp.Process == nil {
					Debug.Printf("Disconnecting in-port %s from out-port %s", rpp.Name(), pop.Name())
					pop.Disconnect(rppName)
				} else if _, ok := procs[rpp.Process.Name()]; !ok {
					Debug.Printf("Disconnecting in-port %s from out-port %s", rpp.Name(), pop.Name())
					pop.Disconnect(rppName)
				}
			}
			if !pop.Connected() {
				Debug.Printf("Connecting disconnected out-port %s of process %s to workflow sink", pop.Name(), pop.Process.Name())
				wf.sink.ConnectParam(pop)
			}
		}

		if len(proc.OutPorts()) == 0 && len(proc.ParamOutPorts()) == 0 {
			if foundNewDriverProc {
				Error.Fatalf("Found more than one process without out-ports nor out-param ports. Cannot use both as drivers (One of them being '%s'). Adapt your workflow accordingly.", proc.Name())
			}
			foundNewDriverProc = true
			wf.driver = proc
		}
	}
}

// RunProcsByName runs a specified set of processes only, specified by their
// names as strings
func (wf *Workflow) RunProcsByName(procNames ...string) {
	procs := map[string]WorkflowProcess{}
	for _, procName := range procNames {
		procs[procName] = wf.Proc(procName)
	}
	wf.runProcs(procs)
}

// RunToProcName runs all processes upstream of, and including, the process with
// name finalProcName
func (wf *Workflow) RunToProcName(finalProcName string) {
	if _, ok := wf.Procs()[finalProcName]; !ok {
		Error.Fatalf("No process with name '%s' in %s workflow. Please check the name supplied to RunToProcName() in your workflow, and try again", finalProcName, wf.Name())
	}
	wf.RunToProc(wf.Proc(finalProcName))
}

// RunToProc runs all processes upstream of, and including, the finalProc
func (wf *Workflow) RunToProc(finalProc WorkflowProcess) {
	procsToRun := upstreamProcsForProc(finalProc)
	procsToRun[finalProc.Name()] = finalProc
	wf.runProcs(procsToRun)
}

// upstreamProcsForProc returns all processes it is connected to, either
// directly or indirectly, via its in-ports and param-in-ports
func upstreamProcsForProc(proc WorkflowProcess) map[string]WorkflowProcess {
	procs := map[string]WorkflowProcess{}
	for _, inp := range proc.InPorts() {
		for _, rpt := range inp.RemotePorts {
			procs[rpt.Process.Name()] = rpt.Process
			mergeWFMaps(procs, upstreamProcsForProc(rpt.Process))
		}
	}
	for _, pip := range proc.ParamInPorts() {
		for _, rpp := range pip.RemotePorts {
			procs[rpp.Process.Name()] = rpp.Process
			mergeWFMaps(procs, upstreamProcsForProc(rpp.Process))
		}
	}
	return procs
}

func mergeWFMaps(a map[string]WorkflowProcess, b map[string]WorkflowProcess) map[string]WorkflowProcess {
	for k, v := range b {
		a[k] = v
	}
	return a
}

// Run runs the workflow
func (wf *Workflow) Run() {
	wf.runProcs(wf.procs)
}
<|endoftext|>


<commit_before>//
// Copyright 2022 The GUAC 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.

// Note: All this code here is temporary and will change often. This module
// must be a leaf in the dependency tree!

package graphdb

import (
	"github.com/neo4j/neo4j-go-driver/v4/neo4j"
)

// Authentification token for connecting to the graph database. Only needed
// during initial connection setup.
// Use the `CreateAuthToken...` functions to create a token.
type AuthToken = neo4j.AuthToken

// Creates a simple authentication token with username, password and authentication realm.
func CreateAuthTokenWithUsernameAndPassword(username string, password string, realm string) AuthToken {
	return neo4j.BasicAuth(username, password, realm)
}

// Creates an empty authentication token.
// Use only for testing!
func CreateAuthTokenForTesting() AuthToken {
	return neo4j.NoAuth()
}

// Client for connecting to the graph database.
// TODO(mihaimaruseac): Switch to v5 and `...WithContext` API when v5 is released.
type Client = neo4j.Driver

// Creates a new connection to the graph database given by `uri`, performing
// authentication via `authToken`.
func NewGraphClient(uri string, authToken AuthToken) (Client, error) {
	// TODO(mihaimaruseac): Allow configuration to control internal
	// attributes of the connection (e.g., max connection pool size, etc.)
	driver, err := neo4j.NewDriver(uri, authToken)
	if err != nil {
		return nil, err
	}

	if err = driver.VerifyConnectivity(); err != nil {
		driver.Close()
		return nil, err
	}
	return driver, nil
}

// A transaction for the database
type Transaction = neo4j.Transaction
// Work done inside a transaction for the database
type TransactionWork = neo4j.TransactionWork

// TODO(mihaimaruseac): Define queries needed for GUAC production

// Runs a simple write query against the graph database.
// Use only for testing!
func WriteQueryForTesting(client Client, query string, args map[string]interface{}) error {
	session := client.NewSession(neo4j.SessionConfig{})
	defer session.Close()

	_, err := session.WriteTransaction(
		func(tx Transaction) (interface{}, error) {
			return tx.Run(query, args)
		})
	return err
}

// Runs a simple read query against the graph database.
// Use only for testing!
// Returns the result as an interface to be handled by the caller (as records),
// but this is not optimal to use in production!
func ReadQueryForTesting(client Client, query string, args map[string]interface{}) ([][]interface{}, error) { ///*(interface{}, error) { //*/(neo4j.Result, error) {
	session := client.NewSession(neo4j.SessionConfig{})
	defer session.Close()

	// For this testing function, we just collect all values.
	values := make([][]interface{}, 0)

	_, err := session.ReadTransaction(
		func(tx Transaction) (interface{}, error) {
			records, err := tx.Run(query, args)
			if err != nil {
				return nil, err
			}
			// Since `records` is valid only while `tx` is in
			// scope, we have to process all data here.
			for records.Next() {
				record := records.Record()
				values = append(values, record.Values)
			}
			if err = records.Err(); err != nil {
				return nil, err
			}
			return values, err
		})

	if err != nil {
		return nil, err
	}
	return values, err
}

// Clears the entire database.
// Use only for testing!
// Very slow on large amounts of data!
func ClearDBForTesting(client Client) error {
	session := client.NewSession(neo4j.SessionConfig{})
	defer session.Close()

	_, err := session.WriteTransaction(
		func(tx Transaction) (interface{}, error) {
			return tx.Run("MATCH (n) DETACH DELETE n", nil)
		})
	return err
}
<commit_msg>Fix typo, fix godoc, separate testing utils<commit_after>//
// Copyright 2022 The GUAC 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.

// Note: All this code here is temporary and will change often. This module
// must be a leaf in the dependency tree!

package graphdb

import (
	"github.com/neo4j/neo4j-go-driver/v4/neo4j"
)

// AuthToken is the authentication token needed for connecting to the graph
// database. Use the `CreateAuthToken...` functions to create a token.
type AuthToken = neo4j.AuthToken

// CreateAuthTokenWithUsernameAndPassword creates a simple authentication token
// with username, password and authentication realm. This is the method to call
// in most scenarios when you need an `AuthToken`.
func CreateAuthTokenWithUsernameAndPassword(username string, password string, realm string) AuthToken {
	return neo4j.BasicAuth(username, password, realm)
}

// Client represents a client to the graph database.
// TODO(mihaimaruseac): Switch to v5 and `...WithContext` API when v5 is released.
type Client = neo4j.Driver

// NewGraphClient creates a new connection to the graph database given by
// `uri`, performing authentication via `authToken`.
func NewGraphClient(uri string, authToken AuthToken) (Client, error) {
	// TODO(mihaimaruseac): Allow configuration to control internal
	// attributes of the connection (e.g., max connection pool size, etc.)
	driver, err := neo4j.NewDriver(uri, authToken)
	if err != nil {
		return nil, err
	}

	if err = driver.VerifyConnectivity(); err != nil {
		driver.Close()
		return nil, err
	}
	return driver, nil
}

// Transaction is a transaction in the database
type Transaction = neo4j.Transaction
// TransactionWork represents work done inside a `Transaction`.
type TransactionWork = neo4j.TransactionWork

// TODO(mihaimaruseac): Define queries needed for GUAC production
<|endoftext|>


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

package create

import (
	"fmt"

	"github.com/spf13/cobra"

	batchv1 "k8s.io/api/batch/v1"
	batchv1beta1 "k8s.io/api/batch/v1beta1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/resource"
	batchv1beta1client "k8s.io/client-go/kubernetes/typed/batch/v1beta1"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util/templates"
)

var (
	cronjobLong = templates.LongDesc(`
		Create a cronjob with the specified name.`)

	cronjobExample = templates.Examples(`
		# Create a cronjob
		kubectl create cronjob my-job --image=busybox

		# Create a cronjob with command
		kubectl create cronjob my-job --image=busybox -- date

		# Create a cronjob with schedule
		kubectl create cronjob test-job --image=busybox --schedule="*/1 * * * *"`)
)

type CreateCronJobOptions struct {
	PrintFlags *genericclioptions.PrintFlags

	PrintObj func(obj runtime.Object) error

	Name     string
	Image    string
	Schedule string
	Command  []string
	Restart  string

	Namespace string
	Client    batchv1beta1client.BatchV1beta1Interface
	DryRun    bool
	Builder   *resource.Builder
	Cmd       *cobra.Command

	genericclioptions.IOStreams
}

func NewCreateCronJobOptions(ioStreams genericclioptions.IOStreams) *CreateCronJobOptions {
	return &CreateCronJobOptions{
		PrintFlags: genericclioptions.NewPrintFlags("created").WithTypeSetter(scheme.Scheme),
		IOStreams:  ioStreams,
	}
}

// NewCmdCreateCronJob is a command to to create CronJobs.
func NewCmdCreateCronJob(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewCreateCronJobOptions(ioStreams)
	cmd := &cobra.Command{
		Use:     "cronjob NAME --image=image --schedule='0/5 * * * ?' -- [COMMAND] [args...]",
		Aliases: []string{"cj"},
		Short:   cronjobLong,
		Long:    cronjobLong,
		Example: cronjobExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.Run())
		},
	}

	o.PrintFlags.AddFlags(cmd)

	cmdutil.AddApplyAnnotationFlags(cmd)
	cmdutil.AddValidateFlags(cmd)
	cmdutil.AddDryRunFlag(cmd)
	cmd.Flags().StringVar(&o.Image, "image", o.Image, "Image name to run.")
	cmd.Flags().StringVar(&o.Schedule, "schedule", o.Schedule, "A schedule in the Cron format the job should be run with.")
	cmd.Flags().StringVar(&o.Restart, "restart", o.Restart, "job's restart policy. supported values: OnFailure, Never")

	return cmd
}

func (o *CreateCronJobOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	name, err := NameFromCommandArgs(cmd, args)
	if err != nil {
		return err
	}
	o.Name = name
	if len(args) > 1 {
		o.Command = args[1:]
	}
	if len(o.Restart) == 0 {
		o.Restart = "OnFailure"
	}

	clientConfig, err := f.ToRESTConfig()
	if err != nil {
		return err
	}
	o.Client, err = batchv1beta1client.NewForConfig(clientConfig)
	if err != nil {
		return err
	}

	o.Namespace, _, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	o.Builder = f.NewBuilder()
	o.Cmd = cmd

	o.DryRun = cmdutil.GetDryRunFlag(cmd)
	if o.DryRun {
		o.PrintFlags.Complete("%s (dry run)")
	}
	printer, err := o.PrintFlags.ToPrinter()
	if err != nil {
		return err
	}
	o.PrintObj = func(obj runtime.Object) error {
		return printer.PrintObj(obj, o.Out)
	}

	return nil
}

func (o *CreateCronJobOptions) Validate() error {
	if len(o.Image) == 0 {
		return fmt.Errorf("--image must be specified")
	}
	if len(o.Schedule) == 0 {
		return fmt.Errorf("--schedule must be specified")
	}
	return nil
}

func (o *CreateCronJobOptions) Run() error {
	var cronjob *batchv1beta1.CronJob
	cronjob = o.createCronJob()

	if !o.DryRun {
		var err error
		cronjob, err = o.Client.CronJobs(o.Namespace).Create(cronjob)
		if err != nil {
			return fmt.Errorf("failed to create cronjob: %v", err)
		}
	}

	return o.PrintObj(cronjob)
}

func (o *CreateCronJobOptions) createCronJob() *batchv1beta1.CronJob {
	return &batchv1beta1.CronJob{
		TypeMeta: metav1.TypeMeta{APIVersion: batchv1beta1.SchemeGroupVersion.String(), Kind: "CronJob"},
		ObjectMeta: metav1.ObjectMeta{
			Name: o.Name,
		},
		Spec: batchv1beta1.CronJobSpec{
			Schedule: o.Schedule,
			JobTemplate: batchv1beta1.JobTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Name: o.Name,
				},
				Spec: batchv1.JobSpec{
					Template: corev1.PodTemplateSpec{
						Spec: corev1.PodSpec{
							Containers: []corev1.Container{
								{
									Name:    o.Name,
									Image:   o.Image,
									Command: o.Command,
								},
							},
							RestartPolicy: corev1.RestartPolicy(o.Restart),
						},
					},
				},
			},
		},
	}
}
<commit_msg>remove the repeat word in documents<commit_after>/*
Copyright 2018 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.
*/

package create

import (
	"fmt"

	"github.com/spf13/cobra"

	batchv1 "k8s.io/api/batch/v1"
	batchv1beta1 "k8s.io/api/batch/v1beta1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/resource"
	batchv1beta1client "k8s.io/client-go/kubernetes/typed/batch/v1beta1"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util/templates"
)

var (
	cronjobLong = templates.LongDesc(`
		Create a cronjob with the specified name.`)

	cronjobExample = templates.Examples(`
		# Create a cronjob
		kubectl create cronjob my-job --image=busybox

		# Create a cronjob with command
		kubectl create cronjob my-job --image=busybox -- date

		# Create a cronjob with schedule
		kubectl create cronjob test-job --image=busybox --schedule="*/1 * * * *"`)
)

type CreateCronJobOptions struct {
	PrintFlags *genericclioptions.PrintFlags

	PrintObj func(obj runtime.Object) error

	Name     string
	Image    string
	Schedule string
	Command  []string
	Restart  string

	Namespace string
	Client    batchv1beta1client.BatchV1beta1Interface
	DryRun    bool
	Builder   *resource.Builder
	Cmd       *cobra.Command

	genericclioptions.IOStreams
}

func NewCreateCronJobOptions(ioStreams genericclioptions.IOStreams) *CreateCronJobOptions {
	return &CreateCronJobOptions{
		PrintFlags: genericclioptions.NewPrintFlags("created").WithTypeSetter(scheme.Scheme),
		IOStreams:  ioStreams,
	}
}

// NewCmdCreateCronJob is a command to create CronJobs.
func NewCmdCreateCronJob(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewCreateCronJobOptions(ioStreams)
	cmd := &cobra.Command{
		Use:     "cronjob NAME --image=image --schedule='0/5 * * * ?' -- [COMMAND] [args...]",
		Aliases: []string{"cj"},
		Short:   cronjobLong,
		Long:    cronjobLong,
		Example: cronjobExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.Run())
		},
	}

	o.PrintFlags.AddFlags(cmd)

	cmdutil.AddApplyAnnotationFlags(cmd)
	cmdutil.AddValidateFlags(cmd)
	cmdutil.AddDryRunFlag(cmd)
	cmd.Flags().StringVar(&o.Image, "image", o.Image, "Image name to run.")
	cmd.Flags().StringVar(&o.Schedule, "schedule", o.Schedule, "A schedule in the Cron format the job should be run with.")
	cmd.Flags().StringVar(&o.Restart, "restart", o.Restart, "job's restart policy. supported values: OnFailure, Never")

	return cmd
}

func (o *CreateCronJobOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	name, err := NameFromCommandArgs(cmd, args)
	if err != nil {
		return err
	}
	o.Name = name
	if len(args) > 1 {
		o.Command = args[1:]
	}
	if len(o.Restart) == 0 {
		o.Restart = "OnFailure"
	}

	clientConfig, err := f.ToRESTConfig()
	if err != nil {
		return err
	}
	o.Client, err = batchv1beta1client.NewForConfig(clientConfig)
	if err != nil {
		return err
	}

	o.Namespace, _, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	o.Builder = f.NewBuilder()
	o.Cmd = cmd

	o.DryRun = cmdutil.GetDryRunFlag(cmd)
	if o.DryRun {
		o.PrintFlags.Complete("%s (dry run)")
	}
	printer, err := o.PrintFlags.ToPrinter()
	if err != nil {
		return err
	}
	o.PrintObj = func(obj runtime.Object) error {
		return printer.PrintObj(obj, o.Out)
	}

	return nil
}

func (o *CreateCronJobOptions) Validate() error {
	if len(o.Image) == 0 {
		return fmt.Errorf("--image must be specified")
	}
	if len(o.Schedule) == 0 {
		return fmt.Errorf("--schedule must be specified")
	}
	return nil
}

func (o *CreateCronJobOptions) Run() error {
	var cronjob *batchv1beta1.CronJob
	cronjob = o.createCronJob()

	if !o.DryRun {
		var err error
		cronjob, err = o.Client.CronJobs(o.Namespace).Create(cronjob)
		if err != nil {
			return fmt.Errorf("failed to create cronjob: %v", err)
		}
	}

	return o.PrintObj(cronjob)
}

func (o *CreateCronJobOptions) createCronJob() *batchv1beta1.CronJob {
	return &batchv1beta1.CronJob{
		TypeMeta: metav1.TypeMeta{APIVersion: batchv1beta1.SchemeGroupVersion.String(), Kind: "CronJob"},
		ObjectMeta: metav1.ObjectMeta{
			Name: o.Name,
		},
		Spec: batchv1beta1.CronJobSpec{
			Schedule: o.Schedule,
			JobTemplate: batchv1beta1.JobTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Name: o.Name,
				},
				Spec: batchv1.JobSpec{
					Template: corev1.PodTemplateSpec{
						Spec: corev1.PodSpec{
							Containers: []corev1.Container{
								{
									Name:    o.Name,
									Image:   o.Image,
									Command: o.Command,
								},
							},
							RestartPolicy: corev1.RestartPolicy(o.Restart),
						},
					},
				},
			},
		},
	}
}
<|endoftext|>


<commit_before>//Copyright 2017 Huawei Technologies Co., Ltd
//
//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 store

import (
	"github.com/ServiceComb/service-center/pkg/util"
	"github.com/ServiceComb/service-center/server/core/backend"
	pb "github.com/ServiceComb/service-center/server/core/proto"
	"github.com/ServiceComb/service-center/server/infra/registry"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"golang.org/x/net/context"
	"strings"
	"sync"
	"time"
)

const (
	DEFAULT_MAX_EVENT_COUNT   = 1000
	DEFAULT_ADD_QUEUE_TIMEOUT = 5 * time.Second
)

var defaultRootKeys map[string]struct{}

func init() {
	defaultRootKeys = make(map[string]struct{}, len(defaultRootKeys))
	for _, root := range TypeRoots {
		defaultRootKeys[root] = struct{}{}
	}
}

type Indexer struct {
	BuildTimeout     time.Duration
	cacher           Cacher
	cacheType        StoreType
	prefixIndex      map[string]map[string]struct{}
	prefixLock       sync.RWMutex
	prefixBuildQueue chan *KvEvent
	goroutine        *util.GoRoutine
	ready            chan struct{}
	isClose          bool
}

func (i *Indexer) Search(ctx context.Context, opts ...registry.PluginOpOption) (*registry.PluginResponse, error) {
	op := registry.OpGet(opts...)

	key := util.BytesToStringWithNoCopy(op.Key)

	if op.Mode == registry.MODE_NO_CACHE ||
		op.Revision > 0 ||
		(op.Offset >= 0 && op.Limit > 0) {
		util.Logger().Debugf("search %s match special options, request etcd server, opts: %s",
			i.cacheType, op)
		return backend.Registry().Do(ctx, opts...)
	}

	if op.Prefix {
		resp, err := i.searchPrefixKeyWithCache(ctx, op)
		if err != nil {
			return nil, err
		}

		if len(resp.Kvs) > 0 || op.Mode == registry.MODE_CACHE {
			return resp, nil
		}

		util.Logger().Debugf("can not find any key from %s cache with prefix, request etcd server, key: %s",
			i.cacheType, key)
		return backend.Registry().Do(ctx, opts...)
	}

	resp := &registry.PluginResponse{
		Action:    op.Action,
		Count:     0,
		Revision:  i.Cache().Version(),
		Succeeded: true,
	}

	if op.CountOnly {
		if i.Cache().Have(key) {
			resp.Count = 1
			return resp, nil
		}
		if op.Mode == registry.MODE_CACHE {
			return resp, nil
		}

		util.Logger().Debugf("%s cache does not store this key, request etcd server, key: %s", i.cacheType, key)
		return backend.Registry().Do(ctx, opts...)
	}

	cacheData := i.Cache().Data(key)
	if cacheData == nil {
		if op.Mode == registry.MODE_CACHE {
			return resp, nil
		}

		util.Logger().Debugf("do not match any key in %s cache store, request etcd server, key: %s",
			i.cacheType, key)
		return backend.Registry().Do(ctx, opts...)
	}

	resp.Count = 1
	resp.Kvs = []*mvccpb.KeyValue{cacheData.(*mvccpb.KeyValue)}
	return resp, nil
}

func (i *Indexer) Cache() Cache {
	return i.cacher.Cache()
}

func (i *Indexer) searchPrefixKeyWithCache(ctx context.Context, op registry.PluginOp) (*registry.PluginResponse, error) {
	resp := &registry.PluginResponse{
		Action:    op.Action,
		Kvs:       []*mvccpb.KeyValue{},
		Count:     0,
		Revision:  i.Cache().Version(),
		Succeeded: true,
	}

	prefix := util.BytesToStringWithNoCopy(op.Key)

	i.prefixLock.RLock()
	resp.Count = int64(i.getPrefixKey(nil, prefix))
	if resp.Count == 0 || op.CountOnly {
		i.prefixLock.RUnlock()
		return resp, nil
	}

	t := time.Now()
	keys := make([]string, 0, resp.Count)
	i.getPrefixKey(&keys, prefix)
	i.prefixLock.RUnlock()

	kvs := make([]*mvccpb.KeyValue, resp.Count)
	idx := 0
	for _, key := range keys {
		c := i.Cache().Data(key) // TODO too slow when big data is requested
		if c == nil {
			// it means resp.Count is not equal to len(keys)
			util.Logger().Warnf(nil, "unexpected nil cache, maybe it is removed, key is %s", key)
			continue
		}
		kvs[idx] = c.(*mvccpb.KeyValue)
		idx++
	}
	util.LogNilOrWarnf(t, "too long to copy data from cache with prefix %s", prefix)

	resp.Kvs = kvs[:idx]
	return resp, nil
}

func (i *Indexer) OnCacheEvent(evt *KvEvent) {
	switch evt.Action {
	case pb.EVT_INIT, pb.EVT_CREATE, pb.EVT_DELETE:
	default:
		return
	}

	if i.isClose {
		return
	}
	defer util.RecoverAndReport()

	ctx, _ := context.WithTimeout(context.Background(), i.BuildTimeout)
	select {
	case <-ctx.Done():
		key := util.BytesToStringWithNoCopy(evt.KV.Key)
		util.Logger().Warnf(nil, "add event to build index queue timed out(%s), key is %s [%s] event",
			i.BuildTimeout, key, evt.Action)
	case i.prefixBuildQueue <- evt:
	}
}

func (i *Indexer) buildIndex() {
	i.goroutine.Do(func(stopCh <-chan struct{}) {
		util.SafeCloseChan(i.ready)
		for {
			select {
			case <-stopCh:
				return
			case evt, ok := <-i.prefixBuildQueue:
				if !ok {
					return
				}
				key := util.BytesToStringWithNoCopy(evt.KV.Key)
				prefix := key[:strings.LastIndex(key, "/")+1]

				i.prefixLock.Lock()
				switch evt.Action {
				case pb.EVT_DELETE:
					i.deletePrefixKey(prefix, key)
				default:
					i.addPrefixKey(prefix, key)
				}
				i.prefixLock.Unlock()

			}
		}
		util.Logger().Debugf("build %s index goroutine is stopped", i.cacheType)
	})
}

func (i *Indexer) getPrefixKey(arr *[]string, prefix string) (count int) {
	keysRef, ok := i.prefixIndex[prefix]
	if !ok {
		return 0
	}

	for key := range keysRef {
		var childs *[]string = nil
		if arr != nil {
			childs = &[]string{}
		}
		n := i.getPrefixKey(childs, key)
		if n == 0 {
			count += len(keysRef)
			if arr != nil {
				for k := range keysRef {
					*arr = append(*arr, k)
				}
			}
			break
		}
		count += n
		if arr != nil {
			*arr = append(*arr, *childs...)
		}
	}
	return count
}

func (i *Indexer) addPrefixKey(prefix, key string) {
	_, ok := defaultRootKeys[key]
	if ok {
		return
	}

	keys, ok := i.prefixIndex[prefix]
	if !ok {
		keys = make(map[string]struct{})
		i.prefixIndex[prefix] = keys
	} else if _, ok := keys[key]; ok {
		return
	}

	keys[key], key = struct{}{}, prefix
	prefix = key[:strings.LastIndex(key[:len(key)-1], "/")+1]

	i.addPrefixKey(prefix, key)
}

func (i *Indexer) deletePrefixKey(prefix, key string) {
	m, ok := i.prefixIndex[prefix]
	if !ok {
		return
	}

	for k := range i.prefixIndex[key] {
		i.deletePrefixKey(key, k)
	}
	delete(m, key)
}

func (i *Indexer) Run() {
	i.prefixLock.Lock()
	if !i.isClose {
		i.prefixLock.Unlock()
		return
	}
	i.isClose = false
	i.prefixLock.Unlock()

	i.buildIndex()

	i.cacher.Run()
}

func (i *Indexer) Stop() {
	i.prefixLock.Lock()
	if i.isClose {
		i.prefixLock.Unlock()
		return
	}
	i.isClose = true
	i.prefixLock.Unlock()

	i.cacher.Stop()

	i.goroutine.Close(true)

	close(i.prefixBuildQueue)

	util.SafeCloseChan(i.ready)
}

func (i *Indexer) Ready() <-chan struct{} {
	<-i.cacher.Ready()
	return i.ready
}

func NewCacheIndexer(t StoreType, cr Cacher) *Indexer {
	return &Indexer{
		BuildTimeout:     DEFAULT_ADD_QUEUE_TIMEOUT,
		cacher:           cr,
		cacheType:        t,
		prefixIndex:      make(map[string]map[string]struct{}, DEFAULT_MAX_EVENT_COUNT),
		prefixBuildQueue: make(chan *KvEvent, DEFAULT_MAX_EVENT_COUNT),
		goroutine:        util.NewGo(make(chan struct{})),
		ready:            make(chan struct{}),
		isClose:          true,
	}
}
<commit_msg>Bug fix: SC panic when received instance DELETE event.<commit_after>//Copyright 2017 Huawei Technologies Co., Ltd
//
//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 store

import (
	"github.com/ServiceComb/service-center/pkg/util"
	"github.com/ServiceComb/service-center/server/core/backend"
	pb "github.com/ServiceComb/service-center/server/core/proto"
	"github.com/ServiceComb/service-center/server/infra/registry"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"golang.org/x/net/context"
	"strings"
	"sync"
	"time"
)

const (
	DEFAULT_MAX_EVENT_COUNT   = 1000
	DEFAULT_ADD_QUEUE_TIMEOUT = 5 * time.Second
)

var defaultRootKeys map[string]struct{}

func init() {
	defaultRootKeys = make(map[string]struct{}, len(defaultRootKeys))
	for _, root := range TypeRoots {
		defaultRootKeys[root] = struct{}{}
	}
}

type Indexer struct {
	BuildTimeout     time.Duration
	cacher           Cacher
	cacheType        StoreType
	prefixIndex      map[string]map[string]struct{}
	prefixLock       sync.RWMutex
	prefixBuildQueue chan *KvEvent
	goroutine        *util.GoRoutine
	ready            chan struct{}
	isClose          bool
}

func (i *Indexer) Search(ctx context.Context, opts ...registry.PluginOpOption) (*registry.PluginResponse, error) {
	op := registry.OpGet(opts...)

	key := util.BytesToStringWithNoCopy(op.Key)

	if op.Mode == registry.MODE_NO_CACHE ||
		op.Revision > 0 ||
		(op.Offset >= 0 && op.Limit > 0) {
		util.Logger().Debugf("search %s match special options, request etcd server, opts: %s",
			i.cacheType, op)
		return backend.Registry().Do(ctx, opts...)
	}

	if op.Prefix {
		resp, err := i.searchPrefixKeyWithCache(ctx, op)
		if err != nil {
			return nil, err
		}

		if len(resp.Kvs) > 0 || op.Mode == registry.MODE_CACHE {
			return resp, nil
		}

		util.Logger().Debugf("can not find any key from %s cache with prefix, request etcd server, key: %s",
			i.cacheType, key)
		return backend.Registry().Do(ctx, opts...)
	}

	resp := &registry.PluginResponse{
		Action:    op.Action,
		Count:     0,
		Revision:  i.Cache().Version(),
		Succeeded: true,
	}

	if op.CountOnly {
		if i.Cache().Have(key) {
			resp.Count = 1
			return resp, nil
		}
		if op.Mode == registry.MODE_CACHE {
			return resp, nil
		}

		util.Logger().Debugf("%s cache does not store this key, request etcd server, key: %s", i.cacheType, key)
		return backend.Registry().Do(ctx, opts...)
	}

	cacheData := i.Cache().Data(key)
	if cacheData == nil {
		if op.Mode == registry.MODE_CACHE {
			return resp, nil
		}

		util.Logger().Debugf("do not match any key in %s cache store, request etcd server, key: %s",
			i.cacheType, key)
		return backend.Registry().Do(ctx, opts...)
	}

	resp.Count = 1
	resp.Kvs = []*mvccpb.KeyValue{cacheData.(*mvccpb.KeyValue)}
	return resp, nil
}

func (i *Indexer) Cache() Cache {
	return i.cacher.Cache()
}

func (i *Indexer) searchPrefixKeyWithCache(ctx context.Context, op registry.PluginOp) (*registry.PluginResponse, error) {
	resp := &registry.PluginResponse{
		Action:    op.Action,
		Kvs:       []*mvccpb.KeyValue{},
		Count:     0,
		Revision:  i.Cache().Version(),
		Succeeded: true,
	}

	prefix := util.BytesToStringWithNoCopy(op.Key)

	i.prefixLock.RLock()
	resp.Count = int64(i.getPrefixKey(nil, prefix))
	if resp.Count == 0 || op.CountOnly {
		i.prefixLock.RUnlock()
		return resp, nil
	}

	t := time.Now()
	keys := make([]string, 0, resp.Count)
	i.getPrefixKey(&keys, prefix)
	i.prefixLock.RUnlock()

	kvs := make([]*mvccpb.KeyValue, resp.Count)
	idx := 0
	for _, key := range keys {
		c := i.Cache().Data(key) // TODO too slow when big data is requested
		if c == nil {
			// it means resp.Count is not equal to len(keys)
			util.Logger().Warnf(nil, "unexpected nil cache, maybe it is removed, key is %s", key)
			continue
		}
		kvs[idx] = c.(*mvccpb.KeyValue)
		idx++
	}
	util.LogNilOrWarnf(t, "too long to copy data from cache with prefix %s", prefix)

	resp.Kvs = kvs[:idx]
	return resp, nil
}

func (i *Indexer) OnCacheEvent(evt *KvEvent) {
	switch evt.Action {
	case pb.EVT_INIT, pb.EVT_CREATE, pb.EVT_DELETE:
	default:
		return
	}

	if i.isClose {
		return
	}
	defer util.RecoverAndReport()

	ctx, _ := context.WithTimeout(context.Background(), i.BuildTimeout)
	select {
	case <-ctx.Done():
		key := util.BytesToStringWithNoCopy(evt.KV.Key)
		util.Logger().Warnf(nil, "add event to build index queue timed out(%s), key is %s [%s] event",
			i.BuildTimeout, key, evt.Action)
	case i.prefixBuildQueue <- evt:
	}
}

func (i *Indexer) buildIndex() {
	i.goroutine.Do(func(stopCh <-chan struct{}) {
		util.SafeCloseChan(i.ready)
		for {
			select {
			case <-stopCh:
				return
			case evt, ok := <-i.prefixBuildQueue:
				if !ok {
					return
				}
				key := util.BytesToStringWithNoCopy(evt.KV.Key)
				prefix := key[:strings.LastIndex(key[:len(key)-1], "/")+1]

				i.prefixLock.Lock()
				switch evt.Action {
				case pb.EVT_DELETE:
					i.deletePrefixKey(prefix, key)
				default:
					i.addPrefixKey(prefix, key)
				}
				i.prefixLock.Unlock()

			}
		}
		util.Logger().Debugf("build %s index goroutine is stopped", i.cacheType)
	})
}

func (i *Indexer) getPrefixKey(arr *[]string, prefix string) (count int) {
	keysRef, ok := i.prefixIndex[prefix]
	if !ok {
		return 0
	}

	for key := range keysRef {
		var childs *[]string = nil
		if arr != nil {
			childs = &[]string{}
		}
		n := i.getPrefixKey(childs, key)
		if n == 0 {
			count += len(keysRef)
			if arr != nil {
				for k := range keysRef {
					*arr = append(*arr, k)
				}
			}
			break
		}
		count += n
		if arr != nil {
			*arr = append(*arr, *childs...)
		}
	}
	return count
}

func (i *Indexer) addPrefixKey(prefix, key string) {
	_, ok := defaultRootKeys[key]
	if ok {
		return
	}

	keys, ok := i.prefixIndex[prefix]
	if !ok {
		keys = make(map[string]struct{})
		i.prefixIndex[prefix] = keys
	} else if _, ok := keys[key]; ok {
		return
	}

	keys[key], key = struct{}{}, prefix
	prefix = key[:strings.LastIndex(key[:len(key)-1], "/")+1]

	i.addPrefixKey(prefix, key)
}

func (i *Indexer) deletePrefixKey(prefix, key string) {
	m, ok := i.prefixIndex[prefix]
	if !ok {
		return
	}

	for k := range i.prefixIndex[key] {
		i.deletePrefixKey(key, k)
	}
	delete(m, key)
}

func (i *Indexer) Run() {
	i.prefixLock.Lock()
	if !i.isClose {
		i.prefixLock.Unlock()
		return
	}
	i.isClose = false
	i.prefixLock.Unlock()

	i.buildIndex()

	i.cacher.Run()
}

func (i *Indexer) Stop() {
	i.prefixLock.Lock()
	if i.isClose {
		i.prefixLock.Unlock()
		return
	}
	i.isClose = true
	i.prefixLock.Unlock()

	i.cacher.Stop()

	i.goroutine.Close(true)

	close(i.prefixBuildQueue)

	util.SafeCloseChan(i.ready)
}

func (i *Indexer) Ready() <-chan struct{} {
	<-i.cacher.Ready()
	return i.ready
}

func NewCacheIndexer(t StoreType, cr Cacher) *Indexer {
	return &Indexer{
		BuildTimeout:     DEFAULT_ADD_QUEUE_TIMEOUT,
		cacher:           cr,
		cacheType:        t,
		prefixIndex:      make(map[string]map[string]struct{}, DEFAULT_MAX_EVENT_COUNT),
		prefixBuildQueue: make(chan *KvEvent, DEFAULT_MAX_EVENT_COUNT),
		goroutine:        util.NewGo(make(chan struct{})),
		ready:            make(chan struct{}),
		isClose:          true,
	}
}
<|endoftext|>