<commit_before>package putio

import "fmt"

// File represents a Put.io file.
type File struct {
	ID                int    `json:"id"`
	Filename          string `json:"name"`
	Filesize          int64  `json:"size"`
	ContentType       string `json:"content_type"`
	CreatedAt         *Time  `json:"created_at"`
	FirstAccessedAt   *Time  `json:"first_accessed_at"`
	ParentID          int    `json:"parent_id"`
	Screenshot        string `json:"screenshot"`
	OpensubtitlesHash string `json:"opensubtitles_hash"`
	IsMP4Available    bool   `json:"is_mp4_available"`
	Icon              string `json:"icon"`
	CRC32             string `json:"crc32"`
	IsShared          bool   `json:"is_shared"`
}

// Upload represents a Put.io upload. If the uploaded file is a torrent file,
// Transfer field will represent the status of the transfer.
type Upload struct {
	File     *File     `json:"file"`
	Transfer *Transfer `json:"transfer"`
}

// Search represents a search response.
type Search struct {
	Files []File `json:"files"`
	Next  string `json:"next"`
}

// Transfer represents a Put.io transfer state.
type Transfer struct {
	Availability       string  `json:"availability"`
	CallbackURL        string  `json:"callback_url"`
	CreatedAt          *Time   `json:"created_at"`
	CreatedTorrent     bool    `json:"created_torrent"`
	ClientIP           string  `json:"client_ip"`
	CurrentRatio       float32 `json:"current_ratio"`
	DownloadSpeed      int     `json:"down_speed"`
	Downloaded         int     `json:"downloaded"`
	DownloadID         int     `json:"download_id"`
	ErrorMessage       string  `json:"error_message"`
	EstimatedTime      string  `json:"estimated_time"`
	Extract            bool    `json:"extract"`
	FileID             int     `json:"file_id"`
	FinishedAt         *Time   `json:"finished_at"`
	ID                 int     `json:"id"`
	IsPrivate          bool    `json:"is_private"`
	MagnetURI          string  `json:"magneturi"`
	Name               string  `json:"name"`
	PeersConnected     int     `json:"peers_connected"`
	PeersGettingFromUs int     `json:"peers_getting_from_us"`
	PeersSendingToUs   int     `json:"peers_sending_to_us"`
	PercentDone        int     `json:"percent_done"`
	SaveParentID       int     `json:"save_parent_id"`
	SecondsSeeding     int     `json:"seconds_seeding"`
	Size               int     `json:"size"`
	Source             string  `json:"source"`
	Status             string  `json:"status"`
	StatusMessage      string  `json:"status_message"`
	SubscriptionID     int     `json:"subscription_id"`
	TorrentLink        string  `json:"torrent_link"`
	TrackerMessage     string  `json:"tracker_message"`
	Trackers           string  `json:"tracker"`
	Type               string  `json:"type"`
	UploadSpeed        int     `json:"up_speed"`
	Uploaded           int     `json:"uploaded"`
}

// Info represents user's account information.
type Info struct {
	AccountActive           bool   `json:"account_active"`
	AvatarURL               string `json:"avatar_url"`
	DaysUntilFilesDeletion  int    `json:"days_until_files_deletion"`
	DefaultSubtitleLanguage string `json:"default_subtitle_language"`
	Disk                    struct {
		Avail int `json:"avail"`
		Size  int `json:"size"`
		Used  int `json:"used"`
	} `json:"disk"`
	HasVoucher                int      `json:"has_voucher"`
	Mail                      string   `json:"mail"`
	PassiveAccount            bool     `json:"passive_account"`
	PlanExpirationDate        string   `json:"plan_expiration_date"`
	Settings                  Settings `json:"settings"`
	SimultaneousDownloadLimit int      `json:"simultaneous_download_limit"`
	SubtitleLanguages         []string `json:"subtitle_languages"`
	UserID                    int      `json:"user_id"`
	Username                  string   `json:"username"`
}

// Settings represents user's personal settings.
type Settings struct {
	CallbackURL             string      `json:"callback_url"`
	DefaultDownloadFolder   int         `json:"default_download_folder"`
	DefaultSubtitleLanguage string      `json:"default_subtitle_language"`
	DownloadFolderUnset     bool        `json:"download_folder_unset"`
	IsInvisible             bool        `json:"is_invisible"`
	Nextepisode             bool        `json:"nextepisode"`
	PrivateDownloadHostIP   interface{} `json:"private_download_host_ip"`
	PushoverToken           string      `json:"pushover_token"`
	Routing                 string      `json:"routing"`
	Sorting                 string      `json:"sorting"`
	SSLEnabled              bool        `json:"ssl_enabled"`
	StartFrom               bool        `json:"start_from"`
	SubtitleLanguages       []string    `json:"subtitle_languages"`
}

// Friend represents Put.io user's friend.
type Friend struct {
	ID        int    `json:"id"`
	Name      string `json:"name"`
	AvatarURL string `json:"avatar_url"`
}

// Zip represents Put.io zip file.
type Zip struct {
	ID        int   `json:"id"`
	CreatedAt *Time `json:"created_at"`

	Size   int    `json:"size"`
	Status string `json:"status"`
	URL    string `json:"url"`

	// FIXME: missing_files field is missin
	MissingFiles string
}

// Subtitle represents a subtitle.
type Subtitle struct {
	Key      string
	Language string
	Name     string
	Source   string
}

// Event represents a Put.io event. It could be a transfer or a shared file.
type Event struct {
	ID           int    `json:"id"`
	FileID       int    `json:"file_id"`
	Source       string `json:"source"`
	Type         string `json:"type"`
	TransferName string `json:"transfer_name"`
	TransferSize int    `json:"transfer_size"`
	CreatedAt    *Time  `json:"created_at"`
}

type share struct {
	FileID   int    `json:"file_id"`
	Filename string `json:"file_name"`
	// Number of friends the file is shared with
	SharedWith int `json:"shared_with"`
}

// errorResponse represents a common error message that Put.io v2 API sends on
// error.
type errorResponse struct {
	ErrorMessage string `json:"error_message"`
	ErrorType    string `json:"error_type"`
	ErrorURI     string `json:"error_uri"`
	Status       string `json:"status"`
	StatusCode   int    `json:"status_code"`
}

func (e errorResponse) Error() string {
	return fmt.Sprintf("StatusCode: %v ErrorType: %v ErrorMsg: %v", e.StatusCode, e.ErrorType, e.ErrorMessage)
}
<commit_msg>zip: unexport MissingFiles field of Zip<commit_after>package putio

import "fmt"

// File represents a Put.io file.
type File struct {
	ID                int    `json:"id"`
	Filename          string `json:"name"`
	Filesize          int64  `json:"size"`
	ContentType       string `json:"content_type"`
	CreatedAt         *Time  `json:"created_at"`
	FirstAccessedAt   *Time  `json:"first_accessed_at"`
	ParentID          int    `json:"parent_id"`
	Screenshot        string `json:"screenshot"`
	OpensubtitlesHash string `json:"opensubtitles_hash"`
	IsMP4Available    bool   `json:"is_mp4_available"`
	Icon              string `json:"icon"`
	CRC32             string `json:"crc32"`
	IsShared          bool   `json:"is_shared"`
}

// Upload represents a Put.io upload. If the uploaded file is a torrent file,
// Transfer field will represent the status of the transfer.
type Upload struct {
	File     *File     `json:"file"`
	Transfer *Transfer `json:"transfer"`
}

// Search represents a search response.
type Search struct {
	Files []File `json:"files"`
	Next  string `json:"next"`
}

// Transfer represents a Put.io transfer state.
type Transfer struct {
	Availability       string  `json:"availability"`
	CallbackURL        string  `json:"callback_url"`
	CreatedAt          *Time   `json:"created_at"`
	CreatedTorrent     bool    `json:"created_torrent"`
	ClientIP           string  `json:"client_ip"`
	CurrentRatio       float32 `json:"current_ratio"`
	DownloadSpeed      int     `json:"down_speed"`
	Downloaded         int     `json:"downloaded"`
	DownloadID         int     `json:"download_id"`
	ErrorMessage       string  `json:"error_message"`
	EstimatedTime      string  `json:"estimated_time"`
	Extract            bool    `json:"extract"`
	FileID             int     `json:"file_id"`
	FinishedAt         *Time   `json:"finished_at"`
	ID                 int     `json:"id"`
	IsPrivate          bool    `json:"is_private"`
	MagnetURI          string  `json:"magneturi"`
	Name               string  `json:"name"`
	PeersConnected     int     `json:"peers_connected"`
	PeersGettingFromUs int     `json:"peers_getting_from_us"`
	PeersSendingToUs   int     `json:"peers_sending_to_us"`
	PercentDone        int     `json:"percent_done"`
	SaveParentID       int     `json:"save_parent_id"`
	SecondsSeeding     int     `json:"seconds_seeding"`
	Size               int     `json:"size"`
	Source             string  `json:"source"`
	Status             string  `json:"status"`
	StatusMessage      string  `json:"status_message"`
	SubscriptionID     int     `json:"subscription_id"`
	TorrentLink        string  `json:"torrent_link"`
	TrackerMessage     string  `json:"tracker_message"`
	Trackers           string  `json:"tracker"`
	Type               string  `json:"type"`
	UploadSpeed        int     `json:"up_speed"`
	Uploaded           int     `json:"uploaded"`
}

// Info represents user's account information.
type Info struct {
	AccountActive           bool   `json:"account_active"`
	AvatarURL               string `json:"avatar_url"`
	DaysUntilFilesDeletion  int    `json:"days_until_files_deletion"`
	DefaultSubtitleLanguage string `json:"default_subtitle_language"`
	Disk                    struct {
		Avail int `json:"avail"`
		Size  int `json:"size"`
		Used  int `json:"used"`
	} `json:"disk"`
	HasVoucher                int      `json:"has_voucher"`
	Mail                      string   `json:"mail"`
	PassiveAccount            bool     `json:"passive_account"`
	PlanExpirationDate        string   `json:"plan_expiration_date"`
	Settings                  Settings `json:"settings"`
	SimultaneousDownloadLimit int      `json:"simultaneous_download_limit"`
	SubtitleLanguages         []string `json:"subtitle_languages"`
	UserID                    int      `json:"user_id"`
	Username                  string   `json:"username"`
}

// Settings represents user's personal settings.
type Settings struct {
	CallbackURL             string      `json:"callback_url"`
	DefaultDownloadFolder   int         `json:"default_download_folder"`
	DefaultSubtitleLanguage string      `json:"default_subtitle_language"`
	DownloadFolderUnset     bool        `json:"download_folder_unset"`
	IsInvisible             bool        `json:"is_invisible"`
	Nextepisode             bool        `json:"nextepisode"`
	PrivateDownloadHostIP   interface{} `json:"private_download_host_ip"`
	PushoverToken           string      `json:"pushover_token"`
	Routing                 string      `json:"routing"`
	Sorting                 string      `json:"sorting"`
	SSLEnabled              bool        `json:"ssl_enabled"`
	StartFrom               bool        `json:"start_from"`
	SubtitleLanguages       []string    `json:"subtitle_languages"`
}

// Friend represents Put.io user's friend.
type Friend struct {
	ID        int    `json:"id"`
	Name      string `json:"name"`
	AvatarURL string `json:"avatar_url"`
}

// Zip represents Put.io zip file.
type Zip struct {
	ID        int   `json:"id"`
	CreatedAt *Time `json:"created_at"`

	Size   int    `json:"size"`
	Status string `json:"status"`
	URL    string `json:"url"`

	// FIXME: missing_files field is missin
	missingFiles string
}

// Subtitle represents a subtitle.
type Subtitle struct {
	Key      string
	Language string
	Name     string
	Source   string
}

// Event represents a Put.io event. It could be a transfer or a shared file.
type Event struct {
	ID           int    `json:"id"`
	FileID       int    `json:"file_id"`
	Source       string `json:"source"`
	Type         string `json:"type"`
	TransferName string `json:"transfer_name"`
	TransferSize int    `json:"transfer_size"`
	CreatedAt    *Time  `json:"created_at"`
}

type share struct {
	FileID   int    `json:"file_id"`
	Filename string `json:"file_name"`
	// Number of friends the file is shared with
	SharedWith int `json:"shared_with"`
}

// errorResponse represents a common error message that Put.io v2 API sends on
// error.
type errorResponse struct {
	ErrorMessage string `json:"error_message"`
	ErrorType    string `json:"error_type"`
	ErrorURI     string `json:"error_uri"`
	Status       string `json:"status"`
	StatusCode   int    `json:"status_code"`
}

func (e errorResponse) Error() string {
	return fmt.Sprintf("StatusCode: %v ErrorType: %v ErrorMsg: %v", e.StatusCode, e.ErrorType, e.ErrorMessage)
}
<|endoftext|>


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

// Package boomer provides commands to run load tests and display results.
package boomer

import (
	"crypto/tls"
	"github.com/valyala/fasthttp"
	"net/url"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/sschepens/pb"
)

type result struct {
	err           error
	statusCode    int
	duration      time.Duration
	contentLength int
}

type Boomer struct {
	// Request is the request to be made.
	Request *fasthttp.Request

	RequestBody string

	// N is the total number of requests to make.
	N int

	// C is the concurrency level, the number of concurrent workers to run.
	C int

	// Timeout in seconds.
	Timeout int

	// Qps is the rate limit.
	Qps int

	// AllowInsecure is an option to allow insecure TLS/SSL certificates.
	AllowInsecure bool

	// DisableCompression is an option to disable compression in response
	DisableCompression bool

	// DisableKeepAlives is an option to prevents re-use of TCP connections between different HTTP requests
	DisableKeepAlives bool

	// Output represents the output type. If "csv" is provided, the
	// output will be dumped as a csv stream.
	Output string

	// ProxyAddr is the address of HTTP proxy server in the format on "host:port".
	// Optional.
	ProxyAddr *url.URL

	// ReadAll determines whether the body of the response needs
	// to be fully consumed.
	ReadAll bool

	bar     *pb.ProgressBar
	results chan *result
	stop    chan struct{}
}

func (b *Boomer) startProgress() {
	if b.Output != "" {
		return
	}
	b.bar = pb.New(b.N)
	b.bar.Format("Bom !")
	b.bar.BarStart = "Pl"
	b.bar.BarEnd = "!"
	b.bar.Empty = " "
	b.bar.Current = "a"
	b.bar.CurrentN = "a"
	b.bar.Start()
}

func (b *Boomer) finalizeProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Finish()
}

func (b *Boomer) incProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Increment()
}

// Run makes all the requests, prints the summary. It blocks until
// all work is done.
func (b *Boomer) Run() {
	b.results = make(chan *result, b.N)
	b.stop = make(chan struct{})
	b.startProgress()

	start := time.Now()
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		<-c
		close(b.stop)
	}()

	b.runWorkers()
	b.finalizeProgress()
	newReport(b.N, b.results, b.Output, time.Now().Sub(start)).finalize()
	close(b.results)
}

func (b *Boomer) runWorker(wg *sync.WaitGroup, ch chan *fasthttp.Request) {
	client := &fasthttp.Client{
		TLSConfig: &tls.Config{
			InsecureSkipVerify: b.AllowInsecure,
		},
		MaxConnsPerHost: 65000,
	}
	resp := fasthttp.AcquireResponse()
	for req := range ch {
		s := time.Now()

		var code int
		var size int

		err := client.Do(req, resp)
		if err == nil {
			size = resp.Header.ContentLength()
			code = resp.Header.StatusCode()
		}

		b.incProgress()
		b.results <- &result{
			statusCode:    code,
			duration:      time.Now().Sub(s),
			err:           err,
			contentLength: size,
		}
	}
	wg.Done()
}

func (b *Boomer) runWorkers() {
	var wg sync.WaitGroup
	wg.Add(b.C)

	var throttle <-chan time.Time
	if b.Qps > 0 {
		throttle = time.Tick(time.Duration(1e6/(b.Qps)) * time.Microsecond)
	}

	jobsch := make(chan *fasthttp.Request, b.C)
	for i := 0; i < b.C; i++ {
		go b.runWorker(&wg, jobsch)
	}

Loop:
	for i := 0; i < b.N; i++ {
		if b.Qps > 0 {
			<-throttle
		}
		select {
		case <-b.stop:
			break Loop
		case jobsch <- cloneRequest(b.Request):
			continue
		}
	}
	close(jobsch)
	wg.Wait()
}

// cloneRequest returns a clone of the provided *http.Request.
// The clone is a shallow copy of the struct and its Header map.
func cloneRequest(r *fasthttp.Request) *fasthttp.Request {
	req := fasthttp.AcquireRequest()
	r.CopyTo(req)
	return req
}
<commit_msg>improve behaviour on stop<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package boomer provides commands to run load tests and display results.
package boomer

import (
	"crypto/tls"
	"github.com/valyala/fasthttp"
	"net/url"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/sschepens/pb"
)

type result struct {
	err           error
	statusCode    int
	duration      time.Duration
	contentLength int
}

type Boomer struct {
	// Request is the request to be made.
	Request *fasthttp.Request

	RequestBody string

	// N is the total number of requests to make.
	N int

	// C is the concurrency level, the number of concurrent workers to run.
	C int

	// Timeout in seconds.
	Timeout int

	// Qps is the rate limit.
	Qps int

	// AllowInsecure is an option to allow insecure TLS/SSL certificates.
	AllowInsecure bool

	// DisableCompression is an option to disable compression in response
	DisableCompression bool

	// DisableKeepAlives is an option to prevents re-use of TCP connections between different HTTP requests
	DisableKeepAlives bool

	// Output represents the output type. If "csv" is provided, the
	// output will be dumped as a csv stream.
	Output string

	// ProxyAddr is the address of HTTP proxy server in the format on "host:port".
	// Optional.
	ProxyAddr *url.URL

	// ReadAll determines whether the body of the response needs
	// to be fully consumed.
	ReadAll bool

	bar     *pb.ProgressBar
	results chan *result
	stop    chan struct{}
}

func (b *Boomer) startProgress() {
	if b.Output != "" {
		return
	}
	b.bar = pb.New(b.N)
	b.bar.Format("Bom !")
	b.bar.BarStart = "Pl"
	b.bar.BarEnd = "!"
	b.bar.Empty = " "
	b.bar.Current = "a"
	b.bar.CurrentN = "a"
	b.bar.Start()
}

func (b *Boomer) finalizeProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Finish()
}

func (b *Boomer) incProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Increment()
}

// Run makes all the requests, prints the summary. It blocks until
// all work is done.
func (b *Boomer) Run() {
	b.results = make(chan *result, b.N)
	b.stop = make(chan struct{})
	b.startProgress()

	start := time.Now()
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		<-c
		b.finalizeProgress()
		close(b.stop)
	}()

	b.runWorkers()
	b.finalizeProgress()
	newReport(b.N, b.results, b.Output, time.Now().Sub(start)).finalize()
	close(b.results)
}

func (b *Boomer) runWorker(wg *sync.WaitGroup, ch chan *fasthttp.Request) {
	client := &fasthttp.Client{
		TLSConfig: &tls.Config{
			InsecureSkipVerify: b.AllowInsecure,
		},
		MaxConnsPerHost: 65000,
	}
	resp := fasthttp.AcquireResponse()
	for req := range ch {
		s := time.Now()

		var code int
		var size int

		err := client.Do(req, resp)
		if err == nil {
			size = resp.Header.ContentLength()
			code = resp.Header.StatusCode()
		}

		b.incProgress()
		b.results <- &result{
			statusCode:    code,
			duration:      time.Now().Sub(s),
			err:           err,
			contentLength: size,
		}
	}
	wg.Done()
}

func (b *Boomer) runWorkers() {
	var wg sync.WaitGroup
	wg.Add(b.C)

	var throttle <-chan time.Time
	if b.Qps > 0 {
		throttle = time.Tick(time.Duration(1e6/(b.Qps)) * time.Microsecond)
	}

	jobsch := make(chan *fasthttp.Request, b.C)
	for i := 0; i < b.C; i++ {
		go b.runWorker(&wg, jobsch)
	}

Loop:
	for i := 0; i < b.N; i++ {
		if b.Qps > 0 {
			<-throttle
		}
		select {
		case <-b.stop:
			break Loop
		case jobsch <- cloneRequest(b.Request):
			continue
		}
	}
	close(jobsch)
	wg.Wait()
}

// cloneRequest returns a clone of the provided *http.Request.
// The clone is a shallow copy of the struct and its Header map.
func cloneRequest(r *fasthttp.Request) *fasthttp.Request {
	req := fasthttp.AcquireRequest()
	r.CopyTo(req)
	return req
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.

package gosnowflake

import (
	"bytes"
	"database/sql"
	"database/sql/driver"
	"fmt"
)

const (
	fixedType byte = iota
	realType
	textType
	dateType
	variantType
	timestampLtzType
	timestampNtzType
	timestampTzType
	objectType
	arrayType
	binaryType
	timeType
	booleanType
)

var (
	// DataTypeFixed is a FIXED datatype.
	DataTypeFixed = []byte{fixedType}
	// DataTypeReal is a REAL datatype.
	DataTypeReal = []byte{realType}
	// DataTypeText is a TEXT datatype.
	DataTypeText = []byte{textType}
	// DataTypeDate is a Date datatype.
	DataTypeDate = []byte{dateType}
	// DataTypeVariant is a TEXT datatype.
	DataTypeVariant = []byte{variantType}
	// DataTypeTimestampLtz is a TIMESTAMP_LTZ datatype.
	DataTypeTimestampLtz = []byte{timestampLtzType}
	// DataTypeTimestampNtz is a TIMESTAMP_NTZ datatype.
	DataTypeTimestampNtz = []byte{timestampNtzType}
	// DataTypeTimestampTz is a TIMESTAMP_TZ datatype.
	DataTypeTimestampTz = []byte{timestampTzType}
	// DataTypeObject is a OBJECT datatype.
	DataTypeObject = []byte{objectType}
	// DataTypeArray is a ARRAY datatype.
	DataTypeArray = []byte{arrayType}
	// DataTypeBinary is a BINARY datatype.
	DataTypeBinary = []byte{binaryType}
	// DataTypeTime is a TIME datatype.
	DataTypeTime = []byte{timeType}
	// DataTypeBoolean is a BOOLEAN datatype.
	DataTypeBoolean = []byte{booleanType}
)

// dataTypeMode returns the subsequent data type in a string representation.
func dataTypeMode(v driver.Value) (tsmode string, err error) {
	if bd, ok := v.([]byte); ok {
		switch {
		case bytes.Equal(bd, DataTypeDate):
			tsmode = "DATE"
		case bytes.Equal(bd, DataTypeTime):
			tsmode = "TIME"
		case bytes.Equal(bd, DataTypeTimestampLtz):
			tsmode = "TIMESTAMP_LTZ"
		case bytes.Equal(bd, DataTypeTimestampNtz):
			tsmode = "TIMESTAMP_NTZ"
		case bytes.Equal(bd, DataTypeTimestampTz):
			tsmode = "TIMESTAMP_TZ"
		case bytes.Equal(bd, DataTypeBinary):
			tsmode = "BINARY"
		default:
			return "", fmt.Errorf(errMsgInvalidByteArray, v)
		}
	} else {
		return "", fmt.Errorf(errMsgInvalidByteArray, v)
	}
	return tsmode, nil
}

// SnowflakeParameter includes the columns output from SHOW PARAMETER command.
type SnowflakeParameter struct {
	Key              string
	Value            string
	Default          string
	Level            string
	Description      string
	SetByUser        string
	SetInJob         string
	SetOn            string
	SetByThreadID    string
	SetByThreadName  string
	SetByClass       string
	ParameterComment string
	Type             string
	Unknown          string // Reserve for added parameter
}

func populateSnowflakeParameter(colname string, p *SnowflakeParameter) interface{} {
	switch colname {
	case "key":
		return &p.Key
	case "value":
		return &p.Value
	case "default":
		return &p.Default
	case "level":
		return &p.Level
	case "description":
		return &p.Description
	case "set_by_user":
		return &p.SetByUser
	case "set_in_job":
		return &p.SetInJob
	case "set_on":
		return &p.SetOn
	case "set_by_thread_id":
		return &p.SetByThreadID
	case "set_by_thread_name":
		return &p.SetByThreadName
	case "set_by_class":
		return &p.SetByClass
	case "parameter_comment":
		return &p.ParameterComment
	case "type":
		return &p.Type
	default:
		debugPanicf("unknown type: %v", colname)
		return &p.Unknown
	}
}

// ScanSnowflakeParameter binds SnowflakeParameter variable with an array of column buffer.
func ScanSnowflakeParameter(rows *sql.Rows) (*SnowflakeParameter, error) {
	var err error
	var columns []string
	columns, err = rows.Columns()
	if err != nil {
		return nil, err
	}
	colNum := len(columns)
	p := SnowflakeParameter{}
	cols := make([]interface{}, colNum)
	for i := 0; i < colNum; i++ {
		cols[i] = populateSnowflakeParameter(columns[i], &p)
	}
	err = rows.Scan(cols...)
	return &p, err
}
<commit_msg>Fix Snowflakeparameter column struct<commit_after>// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.

package gosnowflake

import (
	"bytes"
	"database/sql"
	"database/sql/driver"
	"fmt"
)

const (
	fixedType byte = iota
	realType
	textType
	dateType
	variantType
	timestampLtzType
	timestampNtzType
	timestampTzType
	objectType
	arrayType
	binaryType
	timeType
	booleanType
)

var (
	// DataTypeFixed is a FIXED datatype.
	DataTypeFixed = []byte{fixedType}
	// DataTypeReal is a REAL datatype.
	DataTypeReal = []byte{realType}
	// DataTypeText is a TEXT datatype.
	DataTypeText = []byte{textType}
	// DataTypeDate is a Date datatype.
	DataTypeDate = []byte{dateType}
	// DataTypeVariant is a TEXT datatype.
	DataTypeVariant = []byte{variantType}
	// DataTypeTimestampLtz is a TIMESTAMP_LTZ datatype.
	DataTypeTimestampLtz = []byte{timestampLtzType}
	// DataTypeTimestampNtz is a TIMESTAMP_NTZ datatype.
	DataTypeTimestampNtz = []byte{timestampNtzType}
	// DataTypeTimestampTz is a TIMESTAMP_TZ datatype.
	DataTypeTimestampTz = []byte{timestampTzType}
	// DataTypeObject is a OBJECT datatype.
	DataTypeObject = []byte{objectType}
	// DataTypeArray is a ARRAY datatype.
	DataTypeArray = []byte{arrayType}
	// DataTypeBinary is a BINARY datatype.
	DataTypeBinary = []byte{binaryType}
	// DataTypeTime is a TIME datatype.
	DataTypeTime = []byte{timeType}
	// DataTypeBoolean is a BOOLEAN datatype.
	DataTypeBoolean = []byte{booleanType}
)

// dataTypeMode returns the subsequent data type in a string representation.
func dataTypeMode(v driver.Value) (tsmode string, err error) {
	if bd, ok := v.([]byte); ok {
		switch {
		case bytes.Equal(bd, DataTypeDate):
			tsmode = "DATE"
		case bytes.Equal(bd, DataTypeTime):
			tsmode = "TIME"
		case bytes.Equal(bd, DataTypeTimestampLtz):
			tsmode = "TIMESTAMP_LTZ"
		case bytes.Equal(bd, DataTypeTimestampNtz):
			tsmode = "TIMESTAMP_NTZ"
		case bytes.Equal(bd, DataTypeTimestampTz):
			tsmode = "TIMESTAMP_TZ"
		case bytes.Equal(bd, DataTypeBinary):
			tsmode = "BINARY"
		default:
			return "", fmt.Errorf(errMsgInvalidByteArray, v)
		}
	} else {
		return "", fmt.Errorf(errMsgInvalidByteArray, v)
	}
	return tsmode, nil
}

// SnowflakeParameter includes the columns output from SHOW PARAMETER command.
type SnowflakeParameter struct {
	Key                       string
	Value                     string
	Default                   string
	Level                     string
	Description               string
	SetByUser                 string
	SetInJob                  string
	SetOn                     string
	SetByThreadID             string
	SetByThreadName           string
	SetByClass                string
	ParameterComment          string
	Type                      string
	IsExpired                 string
	ExpiresAt                 string
	SetByControllingParameter string
	Unknown                   string // Reserve for added parameter
}

func populateSnowflakeParameter(colname string, p *SnowflakeParameter) interface{} {
	switch colname {
	case "key":
		return &p.Key
	case "value":
		return &p.Value
	case "default":
		return &p.Default
	case "level":
		return &p.Level
	case "description":
		return &p.Description
	case "set_by_user":
		return &p.SetByUser
	case "set_in_job":
		return &p.SetInJob
	case "set_on":
		return &p.SetOn
	case "set_by_thread_id":
		return &p.SetByThreadID
	case "set_by_thread_name":
		return &p.SetByThreadName
	case "set_by_class":
		return &p.SetByClass
	case "parameter_comment":
		return &p.ParameterComment
	case "type":
		return &p.Type
	case "is_expired":
		return &p.IsExpired
	case "expires_at":
		return &p.ExpiresAt
	case "set_by_controlling_parameter":
		return &p.SetByControllingParameter
	default:
		debugPanicf("unknown type: %v", colname)
		return &p.Unknown
	}
}

// ScanSnowflakeParameter binds SnowflakeParameter variable with an array of column buffer.
func ScanSnowflakeParameter(rows *sql.Rows) (*SnowflakeParameter, error) {
	var err error
	var columns []string
	columns, err = rows.Columns()
	if err != nil {
		return nil, err
	}
	colNum := len(columns)
	p := SnowflakeParameter{}
	cols := make([]interface{}, colNum)
	for i := 0; i < colNum; i++ {
		cols[i] = populateSnowflakeParameter(columns[i], &p)
	}
	err = rows.Scan(cols...)
	return &p, err
}
<|endoftext|>


<commit_before>package brands

import (
	"errors"
	"fmt"

	"github.com/Financial-Times/neo-model-utils-go/mapper"
	log "github.com/Sirupsen/logrus"
	"github.com/jmcvetta/neoism"
)

// Driver interface
type Driver interface {
	Read(id string) (brand Brand, found bool, err error)
	CheckConnectivity() error
}

// CypherDriver struct
type CypherDriver struct {
	db  *neoism.Database
	env string
}

//NewCypherDriver instantiate driver
func NewCypherDriver(db *neoism.Database, env string) CypherDriver {
	return CypherDriver{db, env}
}

// CheckConnectivity tests neo4j by running a simple cypher query
func (pcw CypherDriver) CheckConnectivity() error {
	results := []struct {
		ID int
	}{}
	query := &neoism.CypherQuery{
		Statement: "MATCH (x) RETURN ID(x) LIMIT 1",
		Result:    &results,
	}
	err := pcw.db.Cypher(query)
	log.Debugf("CheckConnectivity results:%+v  err: %+v", results, err)
	return err
}

func (pcw CypherDriver) Read(uuid string) (brand Brand, found bool, err error) {
	results := []struct {
		Brand
	}{}

	query := &neoism.CypherQuery{
		Statement: `
                        MATCH (b:Brand{uuid:{uuid}})
                        OPTIONAL MATCH (b)-[:HAS_PARENT]->(p:Thing)
                        OPTIONAL MATCH (b)<-[:HAS_PARENT]-(c:Thing)
                        RETURN  b.uuid as id, labels(b) as types, b.prefLabel as prefLabel,
                                b.description as description, b.descriptionXML as descriptionXML,
                                b.strapline as strapline, b.imageUrl as _imageUrl,
                                { id: p.uuid, types: labels(p), prefLabel: p.prefLabel } AS parentBrand,
                                collect (
                                { id: c.uuid, types: labels(c), prefLabel: c.prefLabel }
                                ) AS childBrands
                `,
		Parameters: neoism.Props{"uuid": uuid},
		Result:     &results,
	}
	err = pcw.db.Cypher(query)
	if err != nil {
		log.Errorf("Error looking up uuid %s with query %s from neoism: %+v\n", uuid, query.Statement, err)
		return Brand{}, false, fmt.Errorf("Error accessing Brands datastore for uuid: %s", uuid)
	}
	log.Debugf("CypherResult Read Brand for uuid: %s was: %+v", uuid, results)
	if (len(results)) == 0 {
		return Brand{}, false, nil
	} else if len(results) != 1 {
		errMsg := fmt.Sprintf("Multiple brands found with the same uuid:%s !", uuid)
		log.Error(errMsg)
		return Brand{}, true, errors.New(errMsg)
	}
	publicAPITransformation(&results[0].Brand, pcw.env)
	log.Debugf("Returning %v", results[0].Brand)
	return results[0].Brand, true, nil
}

func publicAPITransformation(brand *Brand, env string) {
	if brand.Parent.ID != "" {
		brand.Parent.APIURL = mapper.APIURL(brand.Parent.ID, brand.Parent.Types, env)
		brand.Parent.ID = mapper.IDURL(brand.Parent.ID)
		brand.Parent.Types = mapper.TypeURIs(brand.Parent.Types)
	} else {
		brand.Parent = nil
	}
	if brand.Children[0].ID != "" {
		for idx := range brand.Children {
			brand.Children[idx].APIURL = mapper.APIURL(brand.Children[idx].ID, brand.Children[idx].Types, env)
			brand.Children[idx].ID = mapper.IDURL(brand.Children[idx].ID)
			brand.Children[idx].Types = mapper.TypeURIs(brand.Children[idx].Types)
		}
	} else {
		var empty = []*Thing{}
		brand.Children = empty
	}
	brand.APIURL = mapper.APIURL(brand.ID, brand.Types, env)
	brand.Types = mapper.TypeURIs(brand.Types)
	brand.ID = mapper.IDURL(brand.ID)
}
<commit_msg>Post review updates, pt1<commit_after>package brands

import (
	"errors"
	"fmt"

	"github.com/Financial-Times/neo-model-utils-go/mapper"
	log "github.com/Sirupsen/logrus"
	"github.com/jmcvetta/neoism"
)

// Driver interface
type Driver interface {
	Read(id string) (brand Brand, found bool, err error)
	CheckConnectivity() error
}

// CypherDriver struct
type CypherDriver struct {
	db  *neoism.Database
	env string
}

//NewCypherDriver instantiate driver
func NewCypherDriver(db *neoism.Database, env string) CypherDriver {
	return CypherDriver{db, env}
}

// CheckConnectivity tests neo4j by running a simple cypher query
func (driver CypherDriver) CheckConnectivity() error {
	results := []struct {
		ID int
	}{}
	query := &neoism.CypherQuery{
		Statement: "MATCH (x) RETURN ID(x) LIMIT 1",
		Result:    &results,
	}
	err := driver.db.Cypher(query)
	log.Debugf("CheckConnectivity results:%+v  err: %+v", results, err)
	return err
}

func (driver CypherDriver) Read(uuid string) (brand Brand, found bool, err error) {
	results := []struct {
		Brand
	}{}

	query := &neoism.CypherQuery{
		Statement: `
                        MATCH (b:Brand{uuid:{uuid}})
                        OPTIONAL MATCH (b)-[:HAS_PARENT]->(p:Thing)
                        OPTIONAL MATCH (b)<-[:HAS_PARENT]-(c:Thing)
                        RETURN  b.uuid as id, labels(b) as types, b.prefLabel as prefLabel,
                                b.description as description, b.descriptionXML as descriptionXML,
                                b.strapline as strapline, b.imageUrl as _imageUrl,
                                { id: p.uuid, types: labels(p), prefLabel: p.prefLabel } AS parentBrand,
                                collect (
                                { id: c.uuid, types: labels(c), prefLabel: c.prefLabel }
                                ) AS childBrands
                `,
		Parameters: neoism.Props{"uuid": uuid},
		Result:     &results,
	}
	err = driver.db.Cypher(query)
	if err != nil {
		log.Errorf("Error looking up uuid %s with query %s from neoism: %+v\n", uuid, query.Statement, err)
		return Brand{}, false, fmt.Errorf("Error accessing Brands datastore for uuid: %s", uuid)
	}
	log.Debugf("CypherResult Read Brand for uuid: %s was: %+v", uuid, results)
	if (len(results)) == 0 {
		return Brand{}, false, nil
	} else if len(results) != 1 {
		errMsg := fmt.Sprintf("Multiple brands found with the same uuid:%s !", uuid)
		log.Error(errMsg)
		return Brand{}, true, errors.New(errMsg)
	}
	publicAPITransformation(&results[0].Brand, pcw.env)
	log.Debugf("Returning %v", results[0].Brand)
	return results[0].Brand, true, nil
}

func publicAPITransformation(brand *Brand, env string) {
	if brand.Parent.ID != "" {
		brand.Parent.APIURL = mapper.APIURL(brand.Parent.ID, brand.Parent.Types, env)
		brand.Parent.ID = mapper.IDURL(brand.Parent.ID)
		brand.Parent.Types = mapper.TypeURIs(brand.Parent.Types)
	} else {
		brand.Parent = nil
	}
	if brand.Children[0].ID != "" {
		for idx := range brand.Children {
			brand.Children[idx].APIURL = mapper.APIURL(brand.Children[idx].ID, brand.Children[idx].Types, env)
			brand.Children[idx].ID = mapper.IDURL(brand.Children[idx].ID)
			brand.Children[idx].Types = mapper.TypeURIs(brand.Children[idx].Types)
		}
	} else {
		var empty = []*Thing{}
		brand.Children = empty
	}
	brand.APIURL = mapper.APIURL(brand.ID, brand.Types, env)
	brand.Types = mapper.TypeURIs(brand.Types)
	brand.ID = mapper.IDURL(brand.ID)
}
<|endoftext|>


<commit_before>package db

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/cheggaaa/pb/v3"
	"github.com/go-redis/redis/v8"
	"github.com/inconshreveable/log15"
	"github.com/spf13/viper"
	"github.com/vulsio/go-cpe-dictionary/config"
	"github.com/vulsio/go-cpe-dictionary/models"
	"golang.org/x/xerrors"
)

/**
# Redis Data Structure
- Sets
  ┌─────────────────────────────┬──────────────────────┬──────────────────────────────────┐
  │       KEY                   │  MEMBER              │             PURPOSE              │
  └─────────────────────────────┴──────────────────────┴──────────────────────────────────┘
  ┌─────────────────────────────┬──────────────────────┬──────────────────────────────────┐
  │ CPE#VendorProducts          │ ${vendor}#${product} │ Get ALL Vendor Products          │
  ├─────────────────────────────┼──────────────────────┼──────────────────────────────────┤
  │ CPE#VP#${vendor}#${product} │ CPEURI               │ Get CPEURI by vendor and product │
  ├─────────────────────────────┼──────────────────────┼──────────────────────────────────┤
  │ CPE#DeprecatedCPEs          │ CPEURI               │ Get DeprecatedCPEs               │
  └─────────────────────────────┴──────────────────────┴──────────────────────────────────┘

- Hash
  ┌───┬────────────────┬───────────────┬────────┬──────────────────────────────────────────────────┐
  │NO │    KEY         │   FIELD       │  VALUE │                     PURPOSE                      │
  └───┴────────────────┴───────────────┴────────┴──────────────────────────────────────────────────┘
  ┌───┬────────────────┬───────────────┬────────┬──────────────────────────────────────────────────┐
  │ 1 │ CPE#DEP        │    NVD/JVN    │  JSON  │ TO DELETE OUTDATED AND UNNEEDED FIELD AND MEMBER │
  ├───┼────────────────┼───────────────┼────────┼──────────────────────────────────────────────────┤
  │ 2 │ CPE#FETCHMETA  │   Revision    │ string │ GET Go-Cpe-Dictionary Binary Revision            │
  ├───┼────────────────┼───────────────┼────────┼──────────────────────────────────────────────────┤
  │ 3 │ CPE#FETCHMETA  │ SchemaVersion │  uint  │ GET Go-Cpe-Dictionary Schema Version             │
  └───┴────────────────┴───────────────┴────────┴──────────────────────────────────────────────────┘
**/

const (
	dialectRedis      = "redis"
	vpKeyFormat       = "CPE#VP#%s#%s"
	vpListKey         = "CPE#VendorProducts"
	vpSeparator       = "##"
	deprecatedCPEsKey = "CPE#DeprecatedCPEs"
	depKey            = "CPE#DEP"
	fetchMetaKey      = "CPE#FETCHMETA"
)

// RedisDriver is Driver for Redis
type RedisDriver struct {
	name string
	conn *redis.Client
}

// Name return db name
func (r *RedisDriver) Name() string {
	return r.name
}

// OpenDB opens Database
func (r *RedisDriver) OpenDB(_, dbPath string, _ bool, option Option) (bool, error) {
	return false, r.connectRedis(dbPath, option)
}

func (r *RedisDriver) connectRedis(dbPath string, option Option) error {
	opt, err := redis.ParseURL(dbPath)
	if err != nil {
		return xerrors.Errorf("Failed to parse url. err: %w", err)
	}
	if 0 < option.RedisTimeout.Seconds() {
		opt.ReadTimeout = option.RedisTimeout
	}
	r.conn = redis.NewClient(opt)
	return r.conn.Ping(context.Background()).Err()
}

// CloseDB close Database
func (r *RedisDriver) CloseDB() error {
	if err := r.conn.Close(); err != nil {
		return xerrors.Errorf("Failed to close DB. Type: %s, err: %w", r.name, err)
	}
	return nil
}

// MigrateDB migrates Database
func (r *RedisDriver) MigrateDB() error {
	return nil
}

// IsGoCPEDictModelV1 determines if the DB was created at the time of go-cpe-dictionary Model v1
func (r *RedisDriver) IsGoCPEDictModelV1() (bool, error) {
	ctx := context.Background()

	exists, err := r.conn.Exists(ctx, fetchMetaKey).Result()
	if err != nil {
		return false, xerrors.Errorf("Failed to Exists. err: %w", err)
	}
	if exists == 0 {
		keys, _, err := r.conn.Scan(ctx, 0, "CPE#*", 1).Result()
		if err != nil {
			return false, xerrors.Errorf("Failed to Scan. err: %w", err)
		}
		if len(keys) == 0 {
			return false, nil
		}
		return true, nil
	}

	return false, nil
}

// GetFetchMeta get FetchMeta from Database
func (r *RedisDriver) GetFetchMeta() (*models.FetchMeta, error) {
	ctx := context.Background()

	exists, err := r.conn.Exists(ctx, fetchMetaKey).Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to Exists. err: %w", err)
	}
	if exists == 0 {
		return &models.FetchMeta{GoCPEDictRevision: config.Revision, SchemaVersion: models.LatestSchemaVersion}, nil
	}

	revision, err := r.conn.HGet(ctx, fetchMetaKey, "Revision").Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to HGet Revision. err: %w", err)
	}

	verstr, err := r.conn.HGet(ctx, fetchMetaKey, "SchemaVersion").Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to HGet SchemaVersion. err: %w", err)
	}
	version, err := strconv.ParseUint(verstr, 10, 8)
	if err != nil {
		return nil, xerrors.Errorf("Failed to ParseUint. err: %w", err)
	}

	return &models.FetchMeta{GoCPEDictRevision: revision, SchemaVersion: uint(version)}, nil
}

// UpsertFetchMeta upsert FetchMeta to Database
func (r *RedisDriver) UpsertFetchMeta(fetchMeta *models.FetchMeta) error {
	return r.conn.HSet(context.Background(), fetchMetaKey, map[string]interface{}{"Revision": fetchMeta.GoCPEDictRevision, "SchemaVersion": fetchMeta.SchemaVersion}).Err()
}

// GetVendorProducts : GetVendorProducts
func (r *RedisDriver) GetVendorProducts() (vendorProducts []models.VendorProduct, err error) {
	ctx := context.Background()
	result, err := r.conn.SMembers(ctx, vpListKey).Result()
	if err != nil {
		return nil, err
	}
	for _, vp := range result {
		vpParts := strings.Split(vp, vpSeparator)
		if len(vpParts) != 2 {
			continue
		}
		vendorProducts = append(vendorProducts, models.VendorProduct{
			Vendor:  vpParts[0],
			Product: vpParts[1],
		})
	}
	return vendorProducts, nil
}

// GetCpesByVendorProduct : GetCpesByVendorProduct
func (r *RedisDriver) GetCpesByVendorProduct(vendor, product string) ([]string, []string, error) {
	if vendor == "" || product == "" {
		return nil, nil, nil
	}
	result, err := r.conn.SMembers(context.Background(), fmt.Sprintf(vpKeyFormat, vendor, product)).Result()
	if err != nil {
		return nil, nil, xerrors.Errorf("Failed to SMembers CPE. err: %w", err)
	}

	cpeURIs, deprecated := []string{}, []string{}
	for _, cpeURI := range result {
		ok, err := r.IsDeprecated(cpeURI)
		if err != nil {
			return nil, nil, xerrors.Errorf("Failed to get deprecated CPE. err :%s", err)
		}
		if ok {
			deprecated = append(deprecated, cpeURI)
		} else {
			cpeURIs = append(cpeURIs, cpeURI)
		}
	}
	return cpeURIs, deprecated, nil
}

// InsertCpes Select Cve information from DB.
func (r *RedisDriver) InsertCpes(fetchType models.FetchType, cpes []models.CategorizedCpe) (err error) {
	ctx := context.Background()
	batchSize := viper.GetInt("batch-size")
	if batchSize < 1 {
		return xerrors.Errorf("Failed to set batch-size. err: batch-size option is not set properly")
	}

	// newDeps, oldDeps: {"VP": {"${part}#${vendor}": {"CPEURI": {}}}, "VendorProducts": {"${part}#${vendor}": {}}, "DeprecatedCPEs": {"CPEURI": {}}}
	newDeps := map[string]map[string]map[string]struct{}{
		"VP":             {},
		"VendorProducts": {},
		"DeprecatedCPEs": {},
	}
	oldDepsStr, err := r.conn.HGet(ctx, depKey, string(fetchType)).Result()
	if err != nil {
		if !errors.Is(err, redis.Nil) {
			return xerrors.Errorf("Failed to Get key: %s. err: %w", depKey, err)
		}
		oldDepsStr = `{
			"VP": {},
			"VendorProducts": {},
			"DeprecatedCPEs": {}
		}`
	}
	var oldDeps map[string]map[string]map[string]struct{}
	if err := json.Unmarshal([]byte(oldDepsStr), &oldDeps); err != nil {
		return xerrors.Errorf("Failed to unmarshal JSON. err: %w", err)
	}

	bar := pb.StartNew(len(cpes))
	for idx := range chunkSlice(len(cpes), batchSize) {
		pipe := r.conn.Pipeline()
		for _, c := range cpes[idx.From:idx.To] {
			bar.Increment()
			vendorProductStr := fmt.Sprintf("%s%s%s", c.Vendor, vpSeparator, c.Product)
			if err := pipe.SAdd(ctx, vpListKey, vendorProductStr).Err(); err != nil {
				return xerrors.Errorf("Failed to SAdd vendorProduct. err: %w", err)
			}
			if err := pipe.SAdd(ctx, fmt.Sprintf(vpKeyFormat, c.Vendor, c.Product), c.CpeURI).Err(); err != nil {
				return xerrors.Errorf("Failed to SAdd CpeURI. err: %w", err)
			}
			if _, ok := newDeps["VP"][vendorProductStr]; !ok {
				newDeps["VP"][vendorProductStr] = map[string]struct{}{}
			}
			newDeps["VP"][vendorProductStr][c.CpeURI] = struct{}{}
			if _, ok := oldDeps["VP"][vendorProductStr]; ok {
				delete(oldDeps["VP"][vendorProductStr], c.CpeURI)
				if len(oldDeps["VP"][vendorProductStr]) == 0 {
					delete(oldDeps["VP"], vendorProductStr)
				}
			}
			newDeps["VendorProducts"][vendorProductStr] = map[string]struct{}{}
			delete(oldDeps["VendorProducts"], vendorProductStr)

			if c.Deprecated {
				if err := pipe.SAdd(ctx, deprecatedCPEsKey, c.CpeURI).Err(); err != nil {
					return xerrors.Errorf("Failed to set to deprecated CPE. err: %w", err)
				}
				newDeps["DeprecatedCPEs"][c.CpeURI] = map[string]struct{}{}
				delete(oldDeps["DeprecatedCPEs"], c.CpeURI)
			}
		}
		if _, err = pipe.Exec(ctx); err != nil {
			return xerrors.Errorf("Failed to exec pipeline. err: %w", err)
		}
	}
	bar.Finish()
	log15.Info(fmt.Sprintf("Refreshed %d CPEs.", len(cpes)))

	pipe := r.conn.Pipeline()
	for vendorProductStr, cpeURIs := range oldDeps["VP"] {
		for cpeURI := range cpeURIs {
			ss := strings.Split(vendorProductStr, "#")
			if err := pipe.SRem(ctx, fmt.Sprintf(vpKeyFormat, ss[0], ss[1]), cpeURI).Err(); err != nil {
				return xerrors.Errorf("Failed to SRem. err: %w", err)
			}
		}
	}
	for vendorProductStr := range oldDeps["VendorProducts"] {
		if err := pipe.SRem(ctx, vpListKey, vendorProductStr).Err(); err != nil {
			return xerrors.Errorf("Failed to SRem. err: %w", err)
		}
	}
	for cpeURI := range oldDeps["DeprecatedCPEs"] {
		if err := pipe.SRem(ctx, deprecatedCPEsKey, cpeURI).Err(); err != nil {
			return xerrors.Errorf("Failed to SRem. err: %w", err)
		}
	}

	newDepsJSON, err := json.Marshal(newDeps)
	if err != nil {
		return xerrors.Errorf("Failed to Marshal JSON. err: %w", err)
	}
	if err := pipe.HSet(ctx, depKey, string(fetchType), string(newDepsJSON)).Err(); err != nil {
		return xerrors.Errorf("Failed to Set depkey. err: %w", err)
	}
	if _, err = pipe.Exec(ctx); err != nil {
		return xerrors.Errorf("Failed to exec pipeline. err: %w", err)
	}

	return nil
}

// IsDeprecated : IsDeprecated
func (r *RedisDriver) IsDeprecated(cpeURI string) (bool, error) {
	result, err := r.conn.SIsMember(context.Background(), deprecatedCPEsKey, cpeURI).Result()
	if err != nil {
		if err == redis.Nil {
			return false, nil
		}
		return false, xerrors.Errorf("Failed to SIsMember. err :%s", err)
	}
	return result, nil
}
<commit_msg>fix(redis): update Redis key relationship successfully (#48)<commit_after>package db

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/cheggaaa/pb/v3"
	"github.com/go-redis/redis/v8"
	"github.com/inconshreveable/log15"
	"github.com/spf13/viper"
	"github.com/vulsio/go-cpe-dictionary/config"
	"github.com/vulsio/go-cpe-dictionary/models"
	"golang.org/x/xerrors"
)

/**
# Redis Data Structure
- Sets
  ┌─────────────────────────────┬───────────────────────┬──────────────────────────────────┐
  │       KEY                   │  MEMBER               │             PURPOSE              │
  └─────────────────────────────┴───────────────────────┴──────────────────────────────────┘
  ┌─────────────────────────────┬───────────────────────┬──────────────────────────────────┐
  │ CPE#VendorProducts          │ ${vendor}##${product} │ Get ALL Vendor Products          │
  ├─────────────────────────────┼───────────────────────┼──────────────────────────────────┤
  │ CPE#VP#${vendor}##${product} │ CPEURI               │ Get CPEURI by vendor and product │
  ├─────────────────────────────┼───────────────────────┼──────────────────────────────────┤
  │ CPE#DeprecatedCPEs          │ CPEURI                │ Get DeprecatedCPEs               │
  └─────────────────────────────┴───────────────────────┴──────────────────────────────────┘

- Hash
  ┌───┬────────────────┬───────────────┬────────┬──────────────────────────────────────────────────┐
  │NO │    KEY         │   FIELD       │  VALUE │                     PURPOSE                      │
  └───┴────────────────┴───────────────┴────────┴──────────────────────────────────────────────────┘
  ┌───┬────────────────┬───────────────┬────────┬──────────────────────────────────────────────────┐
  │ 1 │ CPE#DEP        │    NVD/JVN    │  JSON  │ TO DELETE OUTDATED AND UNNEEDED FIELD AND MEMBER │
  ├───┼────────────────┼───────────────┼────────┼──────────────────────────────────────────────────┤
  │ 2 │ CPE#FETCHMETA  │   Revision    │ string │ GET Go-Cpe-Dictionary Binary Revision            │
  ├───┼────────────────┼───────────────┼────────┼──────────────────────────────────────────────────┤
  │ 3 │ CPE#FETCHMETA  │ SchemaVersion │  uint  │ GET Go-Cpe-Dictionary Schema Version             │
  └───┴────────────────┴───────────────┴────────┴──────────────────────────────────────────────────┘
**/

const (
	dialectRedis      = "redis"
	vpKeyFormat       = "CPE#VP#%s##%s"
	vpListKey         = "CPE#VendorProducts"
	vpSeparator       = "##"
	deprecatedCPEsKey = "CPE#DeprecatedCPEs"
	depKey            = "CPE#DEP"
	fetchMetaKey      = "CPE#FETCHMETA"
)

// RedisDriver is Driver for Redis
type RedisDriver struct {
	name string
	conn *redis.Client
}

// Name return db name
func (r *RedisDriver) Name() string {
	return r.name
}

// OpenDB opens Database
func (r *RedisDriver) OpenDB(_, dbPath string, _ bool, option Option) (bool, error) {
	return false, r.connectRedis(dbPath, option)
}

func (r *RedisDriver) connectRedis(dbPath string, option Option) error {
	opt, err := redis.ParseURL(dbPath)
	if err != nil {
		return xerrors.Errorf("Failed to parse url. err: %w", err)
	}
	if 0 < option.RedisTimeout.Seconds() {
		opt.ReadTimeout = option.RedisTimeout
	}
	r.conn = redis.NewClient(opt)
	return r.conn.Ping(context.Background()).Err()
}

// CloseDB close Database
func (r *RedisDriver) CloseDB() error {
	if err := r.conn.Close(); err != nil {
		return xerrors.Errorf("Failed to close DB. Type: %s, err: %w", r.name, err)
	}
	return nil
}

// MigrateDB migrates Database
func (r *RedisDriver) MigrateDB() error {
	return nil
}

// IsGoCPEDictModelV1 determines if the DB was created at the time of go-cpe-dictionary Model v1
func (r *RedisDriver) IsGoCPEDictModelV1() (bool, error) {
	ctx := context.Background()

	exists, err := r.conn.Exists(ctx, fetchMetaKey).Result()
	if err != nil {
		return false, xerrors.Errorf("Failed to Exists. err: %w", err)
	}
	if exists == 0 {
		keys, _, err := r.conn.Scan(ctx, 0, "CPE#*", 1).Result()
		if err != nil {
			return false, xerrors.Errorf("Failed to Scan. err: %w", err)
		}
		if len(keys) == 0 {
			return false, nil
		}
		return true, nil
	}

	return false, nil
}

// GetFetchMeta get FetchMeta from Database
func (r *RedisDriver) GetFetchMeta() (*models.FetchMeta, error) {
	ctx := context.Background()

	exists, err := r.conn.Exists(ctx, fetchMetaKey).Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to Exists. err: %w", err)
	}
	if exists == 0 {
		return &models.FetchMeta{GoCPEDictRevision: config.Revision, SchemaVersion: models.LatestSchemaVersion}, nil
	}

	revision, err := r.conn.HGet(ctx, fetchMetaKey, "Revision").Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to HGet Revision. err: %w", err)
	}

	verstr, err := r.conn.HGet(ctx, fetchMetaKey, "SchemaVersion").Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to HGet SchemaVersion. err: %w", err)
	}
	version, err := strconv.ParseUint(verstr, 10, 8)
	if err != nil {
		return nil, xerrors.Errorf("Failed to ParseUint. err: %w", err)
	}

	return &models.FetchMeta{GoCPEDictRevision: revision, SchemaVersion: uint(version)}, nil
}

// UpsertFetchMeta upsert FetchMeta to Database
func (r *RedisDriver) UpsertFetchMeta(fetchMeta *models.FetchMeta) error {
	return r.conn.HSet(context.Background(), fetchMetaKey, map[string]interface{}{"Revision": fetchMeta.GoCPEDictRevision, "SchemaVersion": fetchMeta.SchemaVersion}).Err()
}

// GetVendorProducts : GetVendorProducts
func (r *RedisDriver) GetVendorProducts() (vendorProducts []models.VendorProduct, err error) {
	ctx := context.Background()
	result, err := r.conn.SMembers(ctx, vpListKey).Result()
	if err != nil {
		return nil, err
	}
	for _, vp := range result {
		vpParts := strings.Split(vp, vpSeparator)
		if len(vpParts) != 2 {
			continue
		}
		vendorProducts = append(vendorProducts, models.VendorProduct{
			Vendor:  vpParts[0],
			Product: vpParts[1],
		})
	}
	return vendorProducts, nil
}

// GetCpesByVendorProduct : GetCpesByVendorProduct
func (r *RedisDriver) GetCpesByVendorProduct(vendor, product string) ([]string, []string, error) {
	if vendor == "" || product == "" {
		return nil, nil, nil
	}
	result, err := r.conn.SMembers(context.Background(), fmt.Sprintf(vpKeyFormat, vendor, product)).Result()
	if err != nil {
		return nil, nil, xerrors.Errorf("Failed to SMembers CPE. err: %w", err)
	}

	cpeURIs, deprecated := []string{}, []string{}
	for _, cpeURI := range result {
		ok, err := r.IsDeprecated(cpeURI)
		if err != nil {
			return nil, nil, xerrors.Errorf("Failed to get deprecated CPE. err :%s", err)
		}
		if ok {
			deprecated = append(deprecated, cpeURI)
		} else {
			cpeURIs = append(cpeURIs, cpeURI)
		}
	}
	return cpeURIs, deprecated, nil
}

// InsertCpes Select Cve information from DB.
func (r *RedisDriver) InsertCpes(fetchType models.FetchType, cpes []models.CategorizedCpe) (err error) {
	ctx := context.Background()
	batchSize := viper.GetInt("batch-size")
	if batchSize < 1 {
		return xerrors.Errorf("Failed to set batch-size. err: batch-size option is not set properly")
	}

	// newDeps, oldDeps: {"VP": {"${part}#${vendor}": {"CPEURI": {}}}, "VendorProducts": {"${part}#${vendor}": {}}, "DeprecatedCPEs": {"CPEURI": {}}}
	newDeps := map[string]map[string]map[string]struct{}{
		"VP":             {},
		"VendorProducts": {},
		"DeprecatedCPEs": {},
	}
	oldDepsStr, err := r.conn.HGet(ctx, depKey, string(fetchType)).Result()
	if err != nil {
		if !errors.Is(err, redis.Nil) {
			return xerrors.Errorf("Failed to Get key: %s. err: %w", depKey, err)
		}
		oldDepsStr = `{
			"VP": {},
			"VendorProducts": {},
			"DeprecatedCPEs": {}
		}`
	}
	var oldDeps map[string]map[string]map[string]struct{}
	if err := json.Unmarshal([]byte(oldDepsStr), &oldDeps); err != nil {
		return xerrors.Errorf("Failed to unmarshal JSON. err: %w", err)
	}

	bar := pb.StartNew(len(cpes))
	for idx := range chunkSlice(len(cpes), batchSize) {
		pipe := r.conn.Pipeline()
		for _, c := range cpes[idx.From:idx.To] {
			bar.Increment()
			vendorProductStr := fmt.Sprintf("%s%s%s", c.Vendor, vpSeparator, c.Product)
			_ = pipe.SAdd(ctx, vpListKey, vendorProductStr)
			_ = pipe.SAdd(ctx, fmt.Sprintf(vpKeyFormat, c.Vendor, c.Product), c.CpeURI)
			if _, ok := newDeps["VP"][vendorProductStr]; !ok {
				newDeps["VP"][vendorProductStr] = map[string]struct{}{}
			}
			newDeps["VP"][vendorProductStr][c.CpeURI] = struct{}{}
			if _, ok := oldDeps["VP"][vendorProductStr]; ok {
				delete(oldDeps["VP"][vendorProductStr], c.CpeURI)
				if len(oldDeps["VP"][vendorProductStr]) == 0 {
					delete(oldDeps["VP"], vendorProductStr)
				}
			}
			newDeps["VendorProducts"][vendorProductStr] = map[string]struct{}{}
			delete(oldDeps["VendorProducts"], vendorProductStr)

			if c.Deprecated {
				_ = pipe.SAdd(ctx, deprecatedCPEsKey, c.CpeURI)
				newDeps["DeprecatedCPEs"][c.CpeURI] = map[string]struct{}{}
				delete(oldDeps["DeprecatedCPEs"], c.CpeURI)
			}
		}
		if _, err = pipe.Exec(ctx); err != nil {
			return xerrors.Errorf("Failed to exec pipeline. err: %w", err)
		}
	}
	bar.Finish()
	log15.Info(fmt.Sprintf("Refreshed %d CPEs.", len(cpes)))

	pipe := r.conn.Pipeline()
	for vendorProductStr, cpeURIs := range oldDeps["VP"] {
		for cpeURI := range cpeURIs {
			ss := strings.Split(vendorProductStr, vpSeparator)
			_ = pipe.SRem(ctx, fmt.Sprintf(vpKeyFormat, ss[0], ss[1]), cpeURI)
		}
	}
	for vendorProductStr := range oldDeps["VendorProducts"] {
		_ = pipe.SRem(ctx, vpListKey, vendorProductStr)
	}
	for cpeURI := range oldDeps["DeprecatedCPEs"] {
		_ = pipe.SRem(ctx, deprecatedCPEsKey, cpeURI)
	}

	newDepsJSON, err := json.Marshal(newDeps)
	if err != nil {
		return xerrors.Errorf("Failed to Marshal JSON. err: %w", err)
	}
	_ = pipe.HSet(ctx, depKey, string(fetchType), string(newDepsJSON))
	if _, err = pipe.Exec(ctx); err != nil {
		return xerrors.Errorf("Failed to exec pipeline. err: %w", err)
	}

	return nil
}

// IsDeprecated : IsDeprecated
func (r *RedisDriver) IsDeprecated(cpeURI string) (bool, error) {
	result, err := r.conn.SIsMember(context.Background(), deprecatedCPEsKey, cpeURI).Result()
	if err != nil {
		if err == redis.Nil {
			return false, nil
		}
		return false, xerrors.Errorf("Failed to SIsMember. err :%s", err)
	}
	return result, nil
}
<|endoftext|>


<commit_before>package middleware

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/DeedleFake/Go-PhysicsFS/physfs"
	"github.com/carbonsrv/carbon/modules/glue"
	"github.com/carbonsrv/carbon/modules/helpers"
	"github.com/carbonsrv/carbon/modules/scheduler"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/pmylund/go-cache"
	"github.com/vifino/golua/lua"
	"github.com/vifino/luar"
	"net/http"
	"time"
)

// Cache
var kvstore *cache.Cache
var cbc *cache.Cache
var cfe *cache.Cache
var LDumper *lua.State

func cacheDump(file string) (string, error, bool) {
	data_tmp, found := cbc.Get(file)
	if found == false {
		data, err := fileRead(file)
		if err != nil {
			return "", err, false
		}
		res, err := bcdump(data)
		if err != nil {
			return "", err, true
		}
		cbc.Set(file, res, cache.DefaultExpiration)
		return res, nil, false
	} else {
		//debug("Using Bytecode-cache for " + file)
		return data_tmp.(string), nil, false
	}
}
func bcdump(data string) (string, error) {
	if LDumper.LoadString(data) != 0 {
		return "", errors.New(LDumper.ToString(-1))
	}
	defer LDumper.Pop(1)
	return LDumper.FDump(), nil
}

// FS
var filesystem http.FileSystem

func fileExists(file string) bool {
	data_tmp, found := cfe.Get(file)
	if found == false {
		exists := physfs.Exists(file)
		cfe.Set(file, exists, cache.DefaultExpiration)
		return exists
	} else {
		return data_tmp.(bool)
	}
}

func fileRead(file string) (string, error) {
	f, err := filesystem.Open(file)
	defer f.Close()
	if err != nil {
		return "", err
	}
	fi, err := f.Stat()
	if err != nil {
		return "", err
	}
	r := bufio.NewReader(f)
	buf := make([]byte, fi.Size())
	_, err = r.Read(buf)
	if err != nil {
		return "", err
	}
	return string(buf), err
}

// Preloader/Starter
var jobs int
var Preloaded chan *lua.State

func Preloader() {
	Preloaded = make(chan *lua.State, jobs)
	for {
		//fmt.Println("preloading")
		L := luar.Init()
		Bind(L)
		err := L.DoString(glue.MainGlue())
		if err != nil {
			panic(err)
		}
		err = L.DoString(glue.RouteGlue())
		if err != nil {
			panic(err)
		}
		Preloaded <- L
	}
}
func GetInstance() *lua.State {
	//fmt.Println("grabbing instance")
	L := <-Preloaded
	//fmt.Println("Done")
	return L
}

// Init
func Init(j int, cfe_new *cache.Cache, kvstore_new *cache.Cache) {
	cfe = cfe_new
	kvstore = kvstore_new
	jobs = j
	filesystem = physfs.FileSystem()
	cbc = cache.New(5*time.Minute, 30*time.Second) // Initialize cache with 5 minute lifetime and purge every 30 seconds
	LDumper = luar.Init()
}

// PHP-like lua scripts
func Lua() func(*gin.Context) {
	//LDumper := luar.Init()
	return func(context *gin.Context) {
		file := context.Request.URL.Path
		if fileExists(file) {
			//fmt.Println("start")
			L := GetInstance()
			//fmt.Println("after start")
			defer scheduler.Add(func() {
				L.Close()
			})
			//fmt.Println("after after start")
			BindContext(L, context)
			//fmt.Println("before cache")
			code, err, lerr := cacheDump(file)
			//fmt.Println("after cache")
			if err != nil {
				if lerr == false {
					context.Next()
					return
				} else {
					helpers.HTMLString(context, http.StatusInternalServerError, `
					<title>Syntax Error in `+context.Request.URL.Path+`</title>

Syntax Error in file `+context.Request.URL.Path+`:

						`+string(err.Error())+`

					`)
					context.Abort()
					return
				}
			}
			//fmt.Println("before loadbuffer")
			L.LoadBuffer(code, len(code), file) // This shouldn't error, was checked earlier.
			if L.Pcall(0, 0, 0) != 0 {          // != 0 means error in execution
				helpers.HTMLString(context, http.StatusInternalServerError, `
				<title>Runtime Error in `+context.Request.URL.Path+`</title>

Runtime Error in file `+context.Request.URL.Path+`:

					`+L.ToString(-1)+`

				`)
				context.Abort()
				return
			}
			/*L.DoString("return CONTENT_TO_RETURN")
			v := luar.CopyTableToMap(L, nil, -1)
			m := v.(map[string]interface{})
			i := int(m["code"].(float64))
			if err != nil {
				i = http.StatusOK
			}*/
			//helpers.HTMLString(context, i, m["content"].(string))
		} else {
			context.Next()
		}
	}
}

// Route creation by lua
func DLR_NS(bcode string, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) {
	/*code, err := bcdump(code)
	if err != nil {
		return func(*gin.Context) {}, err
	}*/
	return func(context *gin.Context) {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		defer scheduler.Add(func() {
			L.Close()
		})
		BindContext(L, context)
		//fmt.Println("before loadbuffer")
		/*if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Syntax Error in `+context.Request.URL.Path+`</title>

Syntax Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}*/
		L.LoadBuffer(bcode, len(bcode), "route")
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Runtime Error on `+context.Request.URL.Path+`</title>

Runtime Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}
	}, nil
}
func DLR_RUS(bcode string, instances int, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) { // Same as above, but reuses states. Much faster. Higher memory use though, because more states.
	insts := instances
	if instances < 0 {
		insts = 2
		if jobs/2 > 1 {
			insts = jobs
		}
	}
	schan := make(chan *lua.State, insts)
	for i := 0; i < jobs/2; i++ {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			return func(context *gin.Context) {}, errors.New(L.ToString(-1))
		}
		L.PushValue(-1)
		schan <- L
	}
	return func(context *gin.Context) {
		L := <-schan
		BindContext(L, context)
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Runtime Error on `+context.Request.URL.Path+`</title>

Runtime Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}
		L.PushValue(-1)
		schan <- L
	}, nil
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func DLRWS_RUS(bcode string, instances int, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) { // Same as above, but for websockets. Not working because?!
	insts := instances
	if instances < 0 {
		insts = 2
		if jobs/2 > 1 {
			insts = jobs
		}
	}
	schan := make(chan *lua.State, insts)
	for i := 0; i < jobs/2; i++ {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			return func(context *gin.Context) {}, errors.New(L.ToString(-1))
		}
		L.PushValue(-1)
		schan <- L
	}
	return func(context *gin.Context) {
		L := <-schan
		BindContext(L, context)
		conn, err := upgrader.Upgrade(context.Writer, context.Request, nil)
		if err != nil {
			fmt.Println("Websocket error: " + err.Error()) // silent error.
		}
		luar.Register(L, "ws", luar.Map{
			"BinaryMessage": websocket.BinaryMessage,
			"TextMessage":   websocket.TextMessage,
			//"read":          conn.ReadMessage,
			//"send":          conn.SendMessage,
			"read": (func() (int, string, error) {
				messageType, p, err := conn.ReadMessage()
				if err != nil {
					return -1, "", err
				}
				return messageType, string(p), nil
			}),
			"send": (func(t int, cnt string) error {
				return conn.WriteMessage(t, []byte(cnt))
			}),
		})
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			fmt.Println("Websocket error: " + L.ToString(-1))
			context.Abort()
			return
		}
		L.PushValue(-1)
		schan <- L
	}, nil
}
<commit_msg>Applying hacks.<commit_after>package middleware

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/DeedleFake/Go-PhysicsFS/physfs"
	"github.com/carbonsrv/carbon/modules/glue"
	"github.com/carbonsrv/carbon/modules/helpers"
	"github.com/carbonsrv/carbon/modules/scheduler"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/pmylund/go-cache"
	"github.com/vifino/golua/lua"
	"github.com/vifino/luar"
	"net/http"
	"time"
)

// Cache
var kvstore *cache.Cache
var cbc *cache.Cache
var cfe *cache.Cache
var LDumper *lua.State

func cacheDump(file string) (string, error, bool) {
	data_tmp, found := cbc.Get(file)
	if found == false {
		data, err := fileRead(file)
		if err != nil {
			return "", err, false
		}
		res, err := bcdump(data)
		if err != nil {
			return "", err, true
		}
		cbc.Set(file, res, cache.DefaultExpiration)
		return res, nil, false
	} else {
		//debug("Using Bytecode-cache for " + file)
		return data_tmp.(string), nil, false
	}
}
func bcdump(data string) (string, error) {
	if LDumper.LoadString(data) != 0 {
		return "", errors.New(LDumper.ToString(-1))
	}
	defer LDumper.Pop(1)
	return LDumper.FDump(), nil
}

// FS
var filesystem http.FileSystem

func fileExists(file string) bool {
	data_tmp, found := cfe.Get(file)
	if found == false {
		exists := physfs.Exists(file)
		cfe.Set(file, exists, cache.DefaultExpiration)
		return exists
	} else {
		return data_tmp.(bool)
	}
}

func fileRead(file string) (string, error) {
	f, err := filesystem.Open(file)
	defer f.Close()
	if err != nil {
		return "", err
	}
	fi, err := f.Stat()
	if err != nil {
		return "", err
	}
	r := bufio.NewReader(f)
	buf := make([]byte, fi.Size())
	_, err = r.Read(buf)
	if err != nil {
		return "", err
	}
	return string(buf), err
}

// Preloader/Starter
var jobs int
var Preloaded chan *lua.State

func Preloader() {
	Preloaded = make(chan *lua.State, jobs)
	for {
		//fmt.Println("preloading")
		L := luar.Init()
		Bind(L)
		err := L.DoString(glue.MainGlue())
		if err != nil {
			panic(err)
		}
		err = L.DoString(glue.RouteGlue())
		if err != nil {
			panic(err)
		}
		Preloaded <- L
	}
}
func GetInstance() *lua.State {
	//fmt.Println("grabbing instance")
	L := <-Preloaded
	//fmt.Println("Done")
	return L
}

// Init
func Init(j int, cfe_new *cache.Cache, kvstore_new *cache.Cache) {
	cfe = cfe_new
	kvstore = kvstore_new
	jobs = j
	filesystem = physfs.FileSystem()
	cbc = cache.New(5*time.Minute, 30*time.Second) // Initialize cache with 5 minute lifetime and purge every 30 seconds
	LDumper = luar.Init()
}

// PHP-like lua scripts
func Lua() func(*gin.Context) {
	//LDumper := luar.Init()
	return func(context *gin.Context) {
		file := context.Request.URL.Path
		if fileExists(file) {
			//fmt.Println("start")
			L := GetInstance()
			//fmt.Println("after start")
			defer scheduler.Add(func() {
				L.Close()
			})
			//fmt.Println("after after start")
			BindContext(L, context)
			//fmt.Println("before cache")
			code, err, lerr := cacheDump(file)
			//fmt.Println("after cache")
			if err != nil {
				if lerr == false {
					context.Next()
					return
				} else {
					helpers.HTMLString(context, http.StatusInternalServerError, `
					<title>Syntax Error in `+context.Request.URL.Path+`</title>

Syntax Error in file `+context.Request.URL.Path+`:

						`+string(err.Error())+`

					`)
					context.Abort()
					return
				}
			}
			//fmt.Println("before loadbuffer")
			L.LoadBuffer(code, len(code), file) // This shouldn't error, was checked earlier.
			if L.Pcall(0, 0, 0) != 0 {          // != 0 means error in execution
				helpers.HTMLString(context, http.StatusInternalServerError, `
				<title>Runtime Error in `+context.Request.URL.Path+`</title>

Runtime Error in file `+context.Request.URL.Path+`:

					`+L.ToString(-1)+`

				`)
				context.Abort()
				return
			}
			/*L.DoString("return CONTENT_TO_RETURN")
			v := luar.CopyTableToMap(L, nil, -1)
			m := v.(map[string]interface{})
			i := int(m["code"].(float64))
			if err != nil {
				i = http.StatusOK
			}*/
			//helpers.HTMLString(context, i, m["content"].(string))
		} else {
			context.Next()
		}
	}
}

// Route creation by lua
func DLR_NS(bcode string, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) {
	/*code, err := bcdump(code)
	if err != nil {
		return func(*gin.Context) {}, err
	}*/
	return func(context *gin.Context) {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		defer scheduler.Add(func() {
			L.Close()
		})
		BindContext(L, context)
		//fmt.Println("before loadbuffer")
		/*if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Syntax Error in `+context.Request.URL.Path+`</title>

Syntax Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}*/
		L.LoadBuffer(bcode, len(bcode), "route")
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Runtime Error on `+context.Request.URL.Path+`</title>

Runtime Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}
	}, nil
}
func DLR_RUS(bcode string, instances int, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) { // Same as above, but reuses states. Much faster. Higher memory use though, because more states.
	insts := instances
	if instances < 0 {
		insts = 2
		if jobs/2 > 1 {
			insts = jobs
		}
	}
	schan := make(chan *lua.State, insts)
	for i := 0; i < jobs/2; i++ {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			return func(context *gin.Context) {}, errors.New(L.ToString(-1))
		}
		L.PushValue(-1)
		schan <- L
	}
	return func(context *gin.Context) {
		L := <-schan
		BindContext(L, context)
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Runtime Error on `+context.Request.URL.Path+`</title>

Runtime Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}
		L.PushValue(-1)
		schan <- L
	}, nil
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func DLRWS_RUS(bcode string, instances int, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) { // Same as above, but for websockets. Not working because?!
	insts := instances
	if instances < 0 {
		insts = 2
		if jobs/2 > 1 {
			insts = jobs
		}
	}
	schan := make(chan *lua.State, insts)
	for i := 0; i < jobs/2; i++ {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			return func(context *gin.Context) {}, errors.New(L.ToString(-1))
		}
		L.PushValue(-1)
		schan <- L
	}
	return func(context *gin.Context) {
		L := <-schan
		BindContext(L, context)
		conn, err := upgrader.Upgrade(context.Writer, context.Request, nil)
		if err != nil {
			fmt.Println("Websocket error: " + err.Error()) // silent error.
		}
		luar.Register(L, "ws", luar.Map{
			"BinaryMessage": websocket.BinaryMessage,
			"TextMessage":   websocket.TextMessage,
			//"read":          conn.ReadMessage,
			//"send":          conn.SendMessage,
			"read": (func() (int, string, error) {
				messageType, p, err := conn.ReadMessage()
				if err != nil {
					return -1, "", err
				}
				return messageType, string(p), nil
			}),
			"send": (func(t int, cnt string) error {
				return conn.WriteMessage(t, []byte(cnt))
			}),
		})
		L.LoadBuffer(bcode, len(bcode), "route")
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			fmt.Println("Websocket error: " + L.ToString(-1))
			context.Abort()
			return
		}
		L.PushValue(-1)
		schan <- L
	}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"github.com/ear7h/m0ney/data"
	"github.com/ear7h/m0ney/log"
	"net/http"
	"strings"
	"time"
)

const (
	QUOTES_URL       = "https://api.robinhood.com/quotes/?symbols="
	NASDAQ_HOURS_URL = "https://api.robinhood.com/markets/XNAS/hours/"
)

var SYMBOLS []string = []string{"AAL", "AAPL", "ADBE", "ADI", "ADP", "ADSK", "AKAM", "ALXN", "AMD", "AMAT", "AMGN", "AMZN", "ATVI", "AVGO", "BIDU", "BIIB", "BMRN", "CA", "CELG", "CERN", "CHKP", "CHTR", "CTRP", "CTAS", "CSCO", "CTXS", "CMCSA", "COST", "CSX", "CTSH", "DISCA", "DISCK", "DISH", "DLTR", "EA", "EBAY", "ESRX", "EXPE", "FAST", "FB", "FISV", "FOX", "FOXA", "GILD", "GOOG", "GOOGL", "HAS", "HSIC", "HOLX", "ILMN", "INCY", "INTC", "INTU", "ISRG", "JBHT", "JD", "KLAC", "KHC", "LBTYK", "LILA", "LBTYA", "QCOM", "QVCA", "MELI", "MAR", "MAT", "MDLZ", "MNST", "MSFT", "MU", "MXIM", "MYL", "NCLH", "NFLX", "NTES", "NVDA", "PAYX", "PCLN", "PYPL", "QCOM", "REGN", "ROST", "SHPG", "SIRI", "SWKS", "SBUX", "SYMC", "TSCO", "TXN", "TMUS", "ULTA", "VIAB", "VOD", "VRTX", "WBA", "WDC", "XRAY", "IDXX", "LILAK", "LRCX", "MCHP", "ORLY", "PCAR", "STX", "TSLA", "VRSK", "WYNN", "XLNX"}

func insertPrices() {
	addr := QUOTES_URL + strings.Join(SYMBOLS, ",")
	log.Enter(log.OK, addr)
	resp, err := http.Get(addr)
	if err != nil {
		log.Enter(log.ERROR, err)
		return
	}
	defer resp.Body.Close()

	resData, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Enter(log.ERROR, err)
		return
	}

	var dat map[string][]data.RhQuote

	err = json.Unmarshal(resData, &dat)
	if err != nil {
		log.Enter(log.ERROR, err)
		return
	}

	for _, v := range dat["results"] {
		err := data.InsertRhQuote(v)
		if err != nil {
			log.Enter(log.ERROR, err)
			return
		}
	}
}

func getMarketHours() (time.Time, time.Time) {

	url := NASDAQ_HOURS_URL + time.Now().Format("2006-01-02") + "/"
	//tag for goto statement
L:
	//get today's market info
	resp, err := http.Get(url)
	if err != nil {
		log.Enter(log.ERROR, err)
		return time.Time{}, time.Time{}
	}
	defer resp.Body.Close()

	byt, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Enter(log.ERROR, err)
		return time.Time{}, time.Time{}
	}

	var m map[string]interface{}

	err = json.Unmarshal(byt, &m)
	if err != nil {
		log.Enter(log.ERROR, err)
		return time.Time{}, time.Time{}
	}

	//if the market is open on fetched day and the close time is in the future
	//then parse time strings
	if m["is_open"].(bool) {
		openTime, err := time.Parse(time.RFC3339, m["opens_at"].(string))
		closeTime, err := time.Parse(time.RFC3339, m["closes_at"].(string))
		if err != nil {
			log.Enter(log.ERROR, err)
			return time.Time{}, time.Time{}
		}

		if closeTime.After(time.Now()) {
			return openTime, closeTime
		}

	}
	//else get the next open hours
	url = m["next_open_hours"].(string)
	//and fetch again
	goto L
}

func insertDataSets(d time.Duration) {

	_, err := data.DB.Exec("INSERT INTO sets " +
		"(`symbol`, `start`, `end`, `scale`, `table`) " +
		"SELECT `symbol`, min(`updated_at`), max(`updated_at`), ?, 'moment' FROM `moment` " +
		"WHERE DATE(`updated_at`) = DATE(NOW()) GROUP BY `symbol`, DATE(`updated_at`) ;", d)

	if err != nil {
		log.Enter(log.ERROR, err)
		return
	}

}

func dayLoop(start, end time.Time) {
	fmt.Println()

	//add data set after completion of day loop
	defer func() {
		insertDataSets(time.Second)
	}()

	//if the market opens in the future
	//then wait
	if time.Now().Before(start) {
		time.Sleep(start.Sub(time.Now()))
	}

	fmt.Println("retriever starting")

	//fetch and insert prices until the close time is in the past
	for time.Now().Before(end) {
		insertPrices()
		time.Sleep(time.Second)
	}

	fmt.Println("retriver finished")

}

//checks config to see if the retriever should run another loop
func shouldRunToday() (retrieveAnother bool) {
	byt, err := ioutil.ReadFile("config.json")
	if err != nil {
		log.Enter(log.ERROR, err)
		return false
	}

	v := map[string]bool{}

	err = json.Unmarshal(byt, &v)
	if err != nil {
		log.Enter(log.ERROR, err)
		return false
	}

	retrieveAnother = v["retrieveAnother"]
	log.Enter(log.OK, "retrieveAnother = ", retrieveAnother)

	return

}

//program entry point
func momentRetriever() error {

	//program loop
	for true {
		s, e := getMarketHours()

		for (time.Time{} == s) || (time.Time{} == e) {
			log.Enter(log.WARNING, "could not market hrs trying again in 10 seconds")
			time.Sleep(10 * time.Second)
			s, e = getMarketHours()
		}

		log.Enter(log.OK, "market hrs: ", s, " - ", e)

		if !shouldRunToday() {
			time.Sleep(time.Until(e))
			continue
		}

		dayLoop(s, e)
	}

	return nil
}
<commit_msg>fixed log issue relating to new use of constants removed shitty use of defer statement<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"github.com/ear7h/m0ney/data"
	"github.com/ear7h/m0ney/log"
	"net/http"
	"strings"
	"time"
)

const (
	QUOTES_URL       = "https://api.robinhood.com/quotes/?symbols="
	NASDAQ_HOURS_URL = "https://api.robinhood.com/markets/XNAS/hours/"
)

var SYMBOLS []string = []string{"AAL", "AAPL", "ADBE", "ADI", "ADP", "ADSK", "AKAM", "ALXN", "AMD", "AMAT", "AMGN", "AMZN", "ATVI", "AVGO", "BIDU", "BIIB", "BMRN", "CA", "CELG", "CERN", "CHKP", "CHTR", "CTRP", "CTAS", "CSCO", "CTXS", "CMCSA", "COST", "CSX", "CTSH", "DISCA", "DISCK", "DISH", "DLTR", "EA", "EBAY", "ESRX", "EXPE", "FAST", "FB", "FISV", "FOX", "FOXA", "GILD", "GOOG", "GOOGL", "HAS", "HSIC", "HOLX", "ILMN", "INCY", "INTC", "INTU", "ISRG", "JBHT", "JD", "KLAC", "KHC", "LBTYK", "LILA", "LBTYA", "QCOM", "QVCA", "MELI", "MAR", "MAT", "MDLZ", "MNST", "MSFT", "MU", "MXIM", "MYL", "NCLH", "NFLX", "NTES", "NVDA", "PAYX", "PCLN", "PYPL", "QCOM", "REGN", "ROST", "SHPG", "SIRI", "SWKS", "SBUX", "SYMC", "TSCO", "TXN", "TMUS", "ULTA", "VIAB", "VOD", "VRTX", "WBA", "WDC", "XRAY", "IDXX", "LILAK", "LRCX", "MCHP", "ORLY", "PCAR", "STX", "TSLA", "VRSK", "WYNN", "XLNX"}

func insertPrices() {
	addr := QUOTES_URL + strings.Join(SYMBOLS, ",")
	log.Enter(log.OK, addr)
	resp, err := http.Get(addr)
	if err != nil {
		log.Enter(log.ERROR, err)
		return
	}
	defer resp.Body.Close()

	resData, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Enter(log.ERROR, err)
		return
	}

	var dat map[string][]data.RhQuote

	err = json.Unmarshal(resData, &dat)
	if err != nil {
		log.Enter(log.ERROR, err)
		return
	}

	for _, v := range dat["results"] {
		err := data.InsertRhQuote(v)
		if err != nil {
			log.Enter(log.ERROR, err)
			return
		}
	}
}

func getMarketHours() (time.Time, time.Time) {

	url := NASDAQ_HOURS_URL + time.Now().Format("2006-01-02") + "/"
	//tag for goto statement
L:
	//get today's market info
	resp, err := http.Get(url)
	if err != nil {
		log.Enter(log.ERROR, err)
		return time.Time{}, time.Time{}
	}
	defer resp.Body.Close()

	byt, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Enter(log.ERROR, err)
		return time.Time{}, time.Time{}
	}

	var m map[string]interface{}

	err = json.Unmarshal(byt, &m)
	if err != nil {
		log.Enter(log.ERROR, err)
		return time.Time{}, time.Time{}
	}

	//if the market is open on fetched day and the close time is in the future
	//then parse time strings
	if m["is_open"].(bool) {
		openTime, err := time.Parse(time.RFC3339, m["opens_at"].(string))
		closeTime, err := time.Parse(time.RFC3339, m["closes_at"].(string))
		if err != nil {
			log.Enter(log.ERROR, err)
			return time.Time{}, time.Time{}
		}

		if closeTime.After(time.Now()) {
			return openTime, closeTime
		}

	}
	//else get the next open hours
	url = m["next_open_hours"].(string)
	//and fetch again
	goto L
}

func insertDataSets(d time.Duration) {

	_, err := data.DB.Exec("INSERT INTO sets " +
		"(`symbol`, `start`, `end`, `scale`, `table`) " +
		"SELECT `symbol`, min(`updated_at`), max(`updated_at`), ?, 'moment' FROM `moment` " +
		"WHERE DATE(`updated_at`) = DATE(NOW()) GROUP BY `symbol`, DATE(`updated_at`) ;", d)

	if err != nil {
		log.Enter(log.ERROR, err)
		return
	}

}

func dayLoop(start, end time.Time) {
	//add data set after completion of day loop
	defer insertDataSets(time.Second)

	//if the market opens in the future
	//then wait
	if time.Now().Before(start) {
		time.Sleep(start.Sub(time.Now()))
	}

	fmt.Println("retriever starting")

	//fetch and insert prices until the close time is in the past
	for time.Now().Before(end) {
		insertPrices()
		time.Sleep(time.Second)
	}

	fmt.Println("retriver finished")

}

//checks config to see if the retriever should run another loop
func shouldRunToday() (retrieveAnother bool) {
	byt, err := ioutil.ReadFile("config.json")
	if err != nil {
		log.Enter(log.ERROR, err)
		return false
	}

	v := map[string]bool{}

	err = json.Unmarshal(byt, &v)
	if err != nil {
		log.Enter(log.ERROR, err)
		return false
	}

	retrieveAnother = v["retrieveAnother"]
	log.Enter(log.OK, "retrieveAnother = ", retrieveAnother)

	return

}

//program entry point
func momentRetriever() error {

	//program loop
	for true {
		s, e := getMarketHours()

		for (time.Time{} == s) || (time.Time{} == e) {
			log.Enter(log.WARNING, "could not market hrs trying again in 10 seconds")
			time.Sleep(10 * time.Second)
			s, e = getMarketHours()
		}

		log.Enter(log.OK, "market hrs: ", s, " - ", e)

		if !shouldRunToday() {
			time.Sleep(time.Until(e))
			continue
		}

		dayLoop(s, e)
	}

	return nil
}
<|endoftext|>


<commit_before>package blocks

import (
	"github.com/bitly/go-nsq"
	"github.com/bitly/go-simplejson"
	"log"
)

type readWriteHandler struct {
	OutChans map[string]chan *simplejson.Json
}

func (self readWriteHandler) HandleMessage(message *nsq.Message) error {
	out, err := simplejson.NewJson(message.Body)
	if err != nil {
		log.Fatal(err.Error())
	}
	broadcast(self.OutChans, out)
	return nil
}

func FromNSQ(b *Block) {

	type fromNSQRule struct {
		ReadTopic   string
		LookupdAddr string
		MaxInFlight int
		ReadChannel string
	}

	rule := &fromNSQRule{}

	unmarshal(<-b.Routes["set_rule"], &rule)

	reader, err := nsq.NewReader(rule.ReadTopic, rule.ReadChannel)
	if err != nil {
		log.Fatal(err.Error())
	}
	reader.SetMaxInFlight(rule.MaxInFlight)

	h := readWriteHandler{b.OutChans}
	reader.AddHandler(h)
	err = reader.ConnectToLookupd(rule.LookupdAddr)
	if err != nil {
		log.Fatal(err.Error())
	}
	for {
		select {
		case <-reader.ExitChan:
			break
		case msg := <-b.AddChan:
			updateOutChans(msg, b)
		case <-b.QuitChan:
			quit(b)
			return
		}
	}
}
<commit_msg>converted to bmsg<commit_after>package blocks

import (
	"github.com/bitly/go-nsq"
	"encoding/json"
	"log"
)

type readWriteHandler struct {
	OutChans map[string]chan BMsg
}

func (self readWriteHandler) HandleMessage(message *nsq.Message) error {
	var msg BMsg
	err := json.Unmarshal(message.Body, &msg)
	if err != nil {
		log.Println(err.Error())
	}

	broadcast(self.OutChans, msg)
	return nil
}

func FromNSQ(b *Block) {

	type fromNSQRule struct {
		ReadTopic   string
		LookupdAddr string
		MaxInFlight int
		ReadChannel string
	}

	rule := &fromNSQRule{}

	unmarshal(<-b.Routes["set_rule"], &rule)

	reader, err := nsq.NewReader(rule.ReadTopic, rule.ReadChannel)
	if err != nil {
		log.Fatal(err.Error())
	}
	reader.SetMaxInFlight(rule.MaxInFlight)

	h := readWriteHandler{b.OutChans}
	reader.AddHandler(h)
	err = reader.ConnectToLookupd(rule.LookupdAddr)
	if err != nil {
		log.Fatal(err.Error())
	}
	for {
		select {
		case <-reader.ExitChan:
			break
		case msg := <-b.AddChan:
			updateOutChans(msg, b)
		case <-b.QuitChan:
			reader.ExitChan <- 1
			quit(b)
			return
		}
	}
}
<|endoftext|>