<commit_before>// Copyright 2015 clair 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 pgsql

import "strconv"

const (
	lockVulnerabilityAffects = `LOCK Vulnerability_Affects_FeatureVersion IN SHARE ROW EXCLUSIVE MODE`
	disableHashJoin          = `SET LOCAL enable_hashjoin = off`
	disableMergeJoin         = `SET LOCAL enable_mergejoin = off`

	// keyvalue.go
	updateKeyValue = `UPDATE KeyValue SET value = $1 WHERE key = $2`
	insertKeyValue = `INSERT INTO KeyValue(key, value) VALUES($1, $2)`
	searchKeyValue = `SELECT value FROM KeyValue WHERE key = $1`

	// namespace.go
	soiNamespace = `
		WITH new_namespace AS (
			INSERT INTO Namespace(name)
			SELECT CAST($1 AS VARCHAR)
			WHERE NOT EXISTS (SELECT name FROM Namespace WHERE name = $1)
			RETURNING id
		)
		SELECT id FROM Namespace WHERE name = $1
		UNION
		SELECT id FROM new_namespace`

	searchNamespace = `SELECT id FROM Namespace WHERE name = $1`
	listNamespace   = `SELECT id, name FROM Namespace`

	// feature.go
	soiFeature = `
		WITH new_feature AS (
			INSERT INTO Feature(name, namespace_id)
			SELECT CAST($1 AS VARCHAR), CAST($2 AS INTEGER)
			WHERE NOT EXISTS (SELECT id FROM Feature WHERE name = $1 AND namespace_id = $2)
			RETURNING id
		)
		SELECT id FROM Feature WHERE name = $1 AND namespace_id = $2
		UNION
		SELECT id FROM new_feature`

	searchFeatureVersion = `
		SELECT id FROM FeatureVersion WHERE feature_id = $1 AND version = $2`

	soiFeatureVersion = `
		WITH new_featureversion AS (
			INSERT INTO FeatureVersion(feature_id, version)
			SELECT CAST($1 AS INTEGER), CAST($2 AS VARCHAR)
			WHERE NOT EXISTS (SELECT id FROM FeatureVersion WHERE feature_id = $1 AND version = $2)
			RETURNING id
		)
		SELECT 'exi', id FROM FeatureVersion WHERE feature_id = $1 AND version = $2
		UNION
		SELECT 'new', id FROM new_featureversion`

	searchVulnerabilityFixedInFeature = `
		SELECT id, vulnerability_id, version FROM Vulnerability_FixedIn_Feature
    WHERE feature_id = $1`

	insertVulnerabilityAffectsFeatureVersion = `
		INSERT INTO Vulnerability_Affects_FeatureVersion(vulnerability_id,
    featureversion_id, fixedin_id) VALUES($1, $2, $3)`

	// layer.go
	searchLayer = `
		SELECT l.id, l.name, l.engineversion, p.id, p.name, n.id, n.name
		FROM Layer l
			LEFT JOIN Layer p ON l.parent_id = p.id
			LEFT JOIN Namespace n ON l.namespace_id = n.id
		WHERE l.name = $1;`

	searchLayerFeatureVersion = `
		WITH RECURSIVE layer_tree(id, name, parent_id, depth, path, cycle) AS(
			SELECT l.id, l.name, l.parent_id, 1, ARRAY[l.id], false
			FROM Layer l
			WHERE l.id = $1
		UNION ALL
			SELECT l.id, l.name, l.parent_id, lt.depth + 1, path || l.id, l.id = ANY(path)
			FROM Layer l, layer_tree lt
			WHERE l.id = lt.parent_id
		)
		SELECT ldf.featureversion_id, ldf.modification, fn.id, fn.name, f.id, f.name, fv.id, fv.version, ltree.id, ltree.name
		FROM Layer_diff_FeatureVersion ldf
		JOIN (
			SELECT row_number() over (ORDER BY depth DESC), id, name FROM layer_tree
		) AS ltree (ordering, id, name) ON ldf.layer_id = ltree.id, FeatureVersion fv, Feature f, Namespace fn
		WHERE ldf.featureversion_id = fv.id AND fv.feature_id = f.id AND f.namespace_id = fn.id
		ORDER BY ltree.ordering`

	searchFeatureVersionVulnerability = `
			SELECT vafv.featureversion_id, v.id, v.name, v.description, v.link, v.severity, v.metadata,
				vn.name, vfif.version
			FROM Vulnerability_Affects_FeatureVersion vafv, Vulnerability v,
					 Namespace vn, Vulnerability_FixedIn_Feature vfif
			WHERE vafv.featureversion_id = ANY($1::integer[])
						AND vfif.vulnerability_id = v.id
						AND vafv.fixedin_id = vfif.id
						AND v.namespace_id = vn.id
						AND v.deleted_at IS NULL`

	insertLayer = `
		INSERT INTO Layer(name, engineversion, parent_id, namespace_id, created_at)
    VALUES($1, $2, $3, $4, CURRENT_TIMESTAMP)
    RETURNING id`

	updateLayer = `UPDATE LAYER SET engineversion = $2, namespace_id = $3 WHERE id = $1`

	removeLayerDiffFeatureVersion = `
		DELETE FROM Layer_diff_FeatureVersion
		WHERE layer_id = $1`

	insertLayerDiffFeatureVersion = `
		INSERT INTO Layer_diff_FeatureVersion(layer_id, featureversion_id, modification)
			SELECT $1, fv.id, $2
			FROM FeatureVersion fv
			WHERE fv.id = ANY($3::integer[])`

	removeLayer = `DELETE FROM Layer WHERE name = $1`

	// lock.go
	insertLock        = `INSERT INTO Lock(name, owner, until) VALUES($1, $2, $3)`
	searchLock        = `SELECT owner, until FROM Lock WHERE name = $1`
	updateLock        = `UPDATE Lock SET until = $3 WHERE name = $1 AND owner = $2`
	removeLock        = `DELETE FROM Lock WHERE name = $1 AND owner = $2`
	removeLockExpired = `DELETE FROM LOCK WHERE until < CURRENT_TIMESTAMP`

	// vulnerability.go
	searchVulnerabilityBase = `
	  SELECT v.id, v.name, n.id, n.name, v.description, v.link, v.severity, v.metadata
	  FROM Vulnerability v JOIN Namespace n ON v.namespace_id = n.id`
	searchVulnerabilityForUpdate          = ` FOR UPDATE OF v`
	searchVulnerabilityByNamespaceAndName = ` WHERE n.name = $1 AND v.name = $2 AND v.deleted_at IS NULL`
	searchVulnerabilityByID               = ` WHERE v.id = $1`
	searchVulnerabilityByNamespace        = ` WHERE n.name = $1 AND v.deleted_at IS NULL
		  				  AND v.id >= $2
						  ORDER BY v.id
						  LIMIT $3`

	searchVulnerabilityFixedIn = `
		SELECT vfif.version, f.id, f.Name
		FROM Vulnerability_FixedIn_Feature vfif JOIN Feature f ON vfif.feature_id = f.id
		WHERE vfif.vulnerability_id = $1`

	insertVulnerability = `
		INSERT INTO Vulnerability(namespace_id, name, description, link, severity, metadata, created_at)
		VALUES($1, $2, $3, $4, $5, $6, CURRENT_TIMESTAMP)
		RETURNING id`

	insertVulnerabilityFixedInFeature = `
		INSERT INTO Vulnerability_FixedIn_Feature(vulnerability_id, feature_id, version)
		VALUES($1, $2, $3)
		RETURNING id`

	searchFeatureVersionByFeature = `SELECT id, version FROM FeatureVersion WHERE feature_id = $1`

	removeVulnerability = `
		UPDATE Vulnerability
    SET deleted_at = CURRENT_TIMESTAMP
    WHERE namespace_id = (SELECT id FROM Namespace WHERE name = $1)
          AND name = $2
          AND deleted_at IS NULL
    RETURNING id`

	// notification.go
	insertNotification = `
		INSERT INTO Vulnerability_Notification(name, created_at, old_vulnerability_id, new_vulnerability_id)
    VALUES($1, CURRENT_TIMESTAMP, $2, $3)`

	updatedNotificationNotified = `
		UPDATE Vulnerability_Notification
		SET notified_at = CURRENT_TIMESTAMP
		WHERE name = $1`

	removeNotification = `
		UPDATE Vulnerability_Notification
	  SET deleted_at = CURRENT_TIMESTAMP
	  WHERE name = $1`

	searchNotificationAvailable = `
		SELECT id, name, created_at, notified_at, deleted_at
		FROM Vulnerability_Notification
		WHERE (notified_at IS NULL OR notified_at < $1)
					AND deleted_at IS NULL
					AND name NOT IN (SELECT name FROM Lock)
		ORDER BY Random()
		LIMIT 1`

	searchNotification = `
		SELECT id, name, created_at, notified_at, deleted_at, old_vulnerability_id, new_vulnerability_id
		FROM Vulnerability_Notification
		WHERE name = $1`

	searchNotificationLayerIntroducingVulnerability = `
		SELECT l.ID, l.name
		FROM Vulnerability v, Vulnerability_Affects_FeatureVersion vafv, FeatureVersion fv, Layer_diff_FeatureVersion ldfv, Layer l
		WHERE v.id = $1
					AND v.id = vafv.vulnerability_id
					AND vafv.featureversion_id = fv.id
					AND fv.id = ldfv.featureversion_id
					AND ldfv.modification = 'add'
					AND ldfv.layer_id = l.id
					AND l.id >= $2
		ORDER BY l.ID
		LIMIT $3`

	// complex_test.go
	searchComplexTestFeatureVersionAffects = `
		SELECT v.name
    FROM FeatureVersion fv
      LEFT JOIN Vulnerability_Affects_FeatureVersion vaf ON fv.id = vaf.featureversion_id
      JOIN Vulnerability v ON vaf.vulnerability_id = v.id
    WHERE featureversion_id = $1`
)

// buildInputArray constructs a PostgreSQL input array from the specified integers.
// Useful to use the `= ANY($1::integer[])` syntax that let us use a IN clause while using
// a single placeholder.
func buildInputArray(ints []int) string {
	str := "{"
	for i := 0; i < len(ints)-1; i++ {
		str = str + strconv.Itoa(ints[i]) + ","
	}
	str = str + strconv.Itoa(ints[len(ints)-1]) + "}"
	return str
}
<commit_msg>pgsql: remove unnecessary join used in GetNotification (#179)<commit_after>// Copyright 2015 clair 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 pgsql

import "strconv"

const (
	lockVulnerabilityAffects = `LOCK Vulnerability_Affects_FeatureVersion IN SHARE ROW EXCLUSIVE MODE`
	disableHashJoin          = `SET LOCAL enable_hashjoin = off`
	disableMergeJoin         = `SET LOCAL enable_mergejoin = off`

	// keyvalue.go
	updateKeyValue = `UPDATE KeyValue SET value = $1 WHERE key = $2`
	insertKeyValue = `INSERT INTO KeyValue(key, value) VALUES($1, $2)`
	searchKeyValue = `SELECT value FROM KeyValue WHERE key = $1`

	// namespace.go
	soiNamespace = `
		WITH new_namespace AS (
			INSERT INTO Namespace(name)
			SELECT CAST($1 AS VARCHAR)
			WHERE NOT EXISTS (SELECT name FROM Namespace WHERE name = $1)
			RETURNING id
		)
		SELECT id FROM Namespace WHERE name = $1
		UNION
		SELECT id FROM new_namespace`

	searchNamespace = `SELECT id FROM Namespace WHERE name = $1`
	listNamespace   = `SELECT id, name FROM Namespace`

	// feature.go
	soiFeature = `
		WITH new_feature AS (
			INSERT INTO Feature(name, namespace_id)
			SELECT CAST($1 AS VARCHAR), CAST($2 AS INTEGER)
			WHERE NOT EXISTS (SELECT id FROM Feature WHERE name = $1 AND namespace_id = $2)
			RETURNING id
		)
		SELECT id FROM Feature WHERE name = $1 AND namespace_id = $2
		UNION
		SELECT id FROM new_feature`

	searchFeatureVersion = `
		SELECT id FROM FeatureVersion WHERE feature_id = $1 AND version = $2`

	soiFeatureVersion = `
		WITH new_featureversion AS (
			INSERT INTO FeatureVersion(feature_id, version)
			SELECT CAST($1 AS INTEGER), CAST($2 AS VARCHAR)
			WHERE NOT EXISTS (SELECT id FROM FeatureVersion WHERE feature_id = $1 AND version = $2)
			RETURNING id
		)
		SELECT 'exi', id FROM FeatureVersion WHERE feature_id = $1 AND version = $2
		UNION
		SELECT 'new', id FROM new_featureversion`

	searchVulnerabilityFixedInFeature = `
		SELECT id, vulnerability_id, version FROM Vulnerability_FixedIn_Feature
    WHERE feature_id = $1`

	insertVulnerabilityAffectsFeatureVersion = `
		INSERT INTO Vulnerability_Affects_FeatureVersion(vulnerability_id,
    featureversion_id, fixedin_id) VALUES($1, $2, $3)`

	// layer.go
	searchLayer = `
		SELECT l.id, l.name, l.engineversion, p.id, p.name, n.id, n.name
		FROM Layer l
			LEFT JOIN Layer p ON l.parent_id = p.id
			LEFT JOIN Namespace n ON l.namespace_id = n.id
		WHERE l.name = $1;`

	searchLayerFeatureVersion = `
		WITH RECURSIVE layer_tree(id, name, parent_id, depth, path, cycle) AS(
			SELECT l.id, l.name, l.parent_id, 1, ARRAY[l.id], false
			FROM Layer l
			WHERE l.id = $1
		UNION ALL
			SELECT l.id, l.name, l.parent_id, lt.depth + 1, path || l.id, l.id = ANY(path)
			FROM Layer l, layer_tree lt
			WHERE l.id = lt.parent_id
		)
		SELECT ldf.featureversion_id, ldf.modification, fn.id, fn.name, f.id, f.name, fv.id, fv.version, ltree.id, ltree.name
		FROM Layer_diff_FeatureVersion ldf
		JOIN (
			SELECT row_number() over (ORDER BY depth DESC), id, name FROM layer_tree
		) AS ltree (ordering, id, name) ON ldf.layer_id = ltree.id, FeatureVersion fv, Feature f, Namespace fn
		WHERE ldf.featureversion_id = fv.id AND fv.feature_id = f.id AND f.namespace_id = fn.id
		ORDER BY ltree.ordering`

	searchFeatureVersionVulnerability = `
			SELECT vafv.featureversion_id, v.id, v.name, v.description, v.link, v.severity, v.metadata,
				vn.name, vfif.version
			FROM Vulnerability_Affects_FeatureVersion vafv, Vulnerability v,
					 Namespace vn, Vulnerability_FixedIn_Feature vfif
			WHERE vafv.featureversion_id = ANY($1::integer[])
						AND vfif.vulnerability_id = v.id
						AND vafv.fixedin_id = vfif.id
						AND v.namespace_id = vn.id
						AND v.deleted_at IS NULL`

	insertLayer = `
		INSERT INTO Layer(name, engineversion, parent_id, namespace_id, created_at)
    VALUES($1, $2, $3, $4, CURRENT_TIMESTAMP)
    RETURNING id`

	updateLayer = `UPDATE LAYER SET engineversion = $2, namespace_id = $3 WHERE id = $1`

	removeLayerDiffFeatureVersion = `
		DELETE FROM Layer_diff_FeatureVersion
		WHERE layer_id = $1`

	insertLayerDiffFeatureVersion = `
		INSERT INTO Layer_diff_FeatureVersion(layer_id, featureversion_id, modification)
			SELECT $1, fv.id, $2
			FROM FeatureVersion fv
			WHERE fv.id = ANY($3::integer[])`

	removeLayer = `DELETE FROM Layer WHERE name = $1`

	// lock.go
	insertLock        = `INSERT INTO Lock(name, owner, until) VALUES($1, $2, $3)`
	searchLock        = `SELECT owner, until FROM Lock WHERE name = $1`
	updateLock        = `UPDATE Lock SET until = $3 WHERE name = $1 AND owner = $2`
	removeLock        = `DELETE FROM Lock WHERE name = $1 AND owner = $2`
	removeLockExpired = `DELETE FROM LOCK WHERE until < CURRENT_TIMESTAMP`

	// vulnerability.go
	searchVulnerabilityBase = `
	  SELECT v.id, v.name, n.id, n.name, v.description, v.link, v.severity, v.metadata
	  FROM Vulnerability v JOIN Namespace n ON v.namespace_id = n.id`
	searchVulnerabilityForUpdate          = ` FOR UPDATE OF v`
	searchVulnerabilityByNamespaceAndName = ` WHERE n.name = $1 AND v.name = $2 AND v.deleted_at IS NULL`
	searchVulnerabilityByID               = ` WHERE v.id = $1`
	searchVulnerabilityByNamespace        = ` WHERE n.name = $1 AND v.deleted_at IS NULL
		  				  AND v.id >= $2
						  ORDER BY v.id
						  LIMIT $3`

	searchVulnerabilityFixedIn = `
		SELECT vfif.version, f.id, f.Name
		FROM Vulnerability_FixedIn_Feature vfif JOIN Feature f ON vfif.feature_id = f.id
		WHERE vfif.vulnerability_id = $1`

	insertVulnerability = `
		INSERT INTO Vulnerability(namespace_id, name, description, link, severity, metadata, created_at)
		VALUES($1, $2, $3, $4, $5, $6, CURRENT_TIMESTAMP)
		RETURNING id`

	insertVulnerabilityFixedInFeature = `
		INSERT INTO Vulnerability_FixedIn_Feature(vulnerability_id, feature_id, version)
		VALUES($1, $2, $3)
		RETURNING id`

	searchFeatureVersionByFeature = `SELECT id, version FROM FeatureVersion WHERE feature_id = $1`

	removeVulnerability = `
		UPDATE Vulnerability
    SET deleted_at = CURRENT_TIMESTAMP
    WHERE namespace_id = (SELECT id FROM Namespace WHERE name = $1)
          AND name = $2
          AND deleted_at IS NULL
    RETURNING id`

	// notification.go
	insertNotification = `
		INSERT INTO Vulnerability_Notification(name, created_at, old_vulnerability_id, new_vulnerability_id)
    VALUES($1, CURRENT_TIMESTAMP, $2, $3)`

	updatedNotificationNotified = `
		UPDATE Vulnerability_Notification
		SET notified_at = CURRENT_TIMESTAMP
		WHERE name = $1`

	removeNotification = `
		UPDATE Vulnerability_Notification
	  SET deleted_at = CURRENT_TIMESTAMP
	  WHERE name = $1`

	searchNotificationAvailable = `
		SELECT id, name, created_at, notified_at, deleted_at
		FROM Vulnerability_Notification
		WHERE (notified_at IS NULL OR notified_at < $1)
					AND deleted_at IS NULL
					AND name NOT IN (SELECT name FROM Lock)
		ORDER BY Random()
		LIMIT 1`

	searchNotification = `
		SELECT id, name, created_at, notified_at, deleted_at, old_vulnerability_id, new_vulnerability_id
		FROM Vulnerability_Notification
		WHERE name = $1`

	searchNotificationLayerIntroducingVulnerability = `
		SELECT l.ID, l.name
		FROM Vulnerability_Affects_FeatureVersion vafv, FeatureVersion fv, Layer_diff_FeatureVersion ldfv, Layer l
		WHERE l.id >= $2
					AND vafv.vulnerability_id = $1
					AND vafv.featureversion_id = fv.id
					AND ldfv.featureversion_id = fv.id
					AND ldfv.modification = 'add'
					AND ldfv.layer_id = l.id
		ORDER BY l.ID
		LIMIT $3`

	// complex_test.go
	searchComplexTestFeatureVersionAffects = `
		SELECT v.name
    FROM FeatureVersion fv
      LEFT JOIN Vulnerability_Affects_FeatureVersion vaf ON fv.id = vaf.featureversion_id
      JOIN Vulnerability v ON vaf.vulnerability_id = v.id
    WHERE featureversion_id = $1`
)

// buildInputArray constructs a PostgreSQL input array from the specified integers.
// Useful to use the `= ANY($1::integer[])` syntax that let us use a IN clause while using
// a single placeholder.
func buildInputArray(ints []int) string {
	str := "{"
	for i := 0; i < len(ints)-1; i++ {
		str = str + strconv.Itoa(ints[i]) + ","
	}
	str = str + strconv.Itoa(ints[len(ints)-1]) + "}"
	return str
}
<|endoftext|>


<commit_before>package gapi

import (
	"bytes"
	"encoding/json"
	"fmt"
)

// DatasourcePermission has information such as a datasource, user, team, role and permission.
type DatasourcePermission struct {
	ID           int64  `json:"id"`
	DatasourceID int64  `json:"datasourceId"`
	UserID       int64  `json:"userId"`
	UserEmail    string `json:"userEmail"`
	TeamID       int64  `json:"teamId"`

	// Permission levels are
	// 1 = Query
	Permission     int64  `json:"permission"`
	PermissionName string `json:"permissionName"`
}

type DatasourcePermissionsResponse struct {
	DatasourceID int64 `json:"datasourceId"`
	Enabled      bool  `json:"enabled"`
	Permissions  []*DatasourcePermission
}

type DatasourcePermissionAddPayload struct {
	UserID     int64 `json:"userId"`
	TeamID     int64 `json:"teamId"`
	Permission int64 `json:"permission"`
}

// DatasourcePermissions fetches and returns the permissions for the datasource whose ID it's passed.
func (c *Client) DatasourcePermissions(id int64) (*DatasourcePermissionsResponse, error) {
	path := fmt.Sprintf("/api/datasources/%d/permissions", id)
	var out *DatasourcePermissionsResponse
	err := c.request("GET", path, nil, nil, &out)
	if err != nil {
		return out, fmt.Errorf("error getting permissions at %s: %w", path, err)
	}

	return out, nil
}

// AddDatasourcePermission adds the given permission item
func (c *Client) AddDatasourcePermission(id int64, item *DatasourcePermissionAddPayload) error {
	path := fmt.Sprintf("/api/datasources/%d/permissions", id)
	data, err := json.Marshal(item)
	if err != nil {
		return fmt.Errorf("marshal err: %w", err)
	}

	if err = c.request("POST", path, nil, bytes.NewBuffer(data), nil); err != nil {
		return fmt.Errorf("error adding permissions at %s: %w", path, err)
	}

	return nil
}

// RemoveDatasourcePermission removes the permission with the given id
func (c *Client) RemoveDatasourcePermission(id, permissionID int64) error {
	path := fmt.Sprintf("/api/datasources/%d/permissions/%d", id, permissionID)
	if err := c.request("DELETE", path, nil, nil, nil); err != nil {
		return fmt.Errorf("error deleting permissions at %s: %w", path, err)
	}

	return nil
}
<commit_msg>Add enable/disable permissions<commit_after>package gapi

import (
	"bytes"
	"encoding/json"
	"fmt"
)

// DatasourcePermission has information such as a datasource, user, team, role and permission.
type DatasourcePermission struct {
	ID           int64  `json:"id"`
	DatasourceID int64  `json:"datasourceId"`
	UserID       int64  `json:"userId"`
	UserEmail    string `json:"userEmail"`
	TeamID       int64  `json:"teamId"`

	// Permission levels are
	// 1 = Query
	Permission     int64  `json:"permission"`
	PermissionName string `json:"permissionName"`
}

type DatasourcePermissionsResponse struct {
	DatasourceID int64 `json:"datasourceId"`
	Enabled      bool  `json:"enabled"`
	Permissions  []*DatasourcePermission
}

type DatasourcePermissionAddPayload struct {
	UserID     int64 `json:"userId"`
	TeamID     int64 `json:"teamId"`
	Permission int64 `json:"permission"`
}

// EnableDatasourcePermissions enables the datasource permissions (this is a datasource setting)
func (c *Client) EnableDatasourcePermissions(id int64) error {
	path := fmt.Sprintf("/api/datasources/%d/enable-permissions", id)
	if err := c.request("POST", path, nil, nil, nil); err != nil {
		return fmt.Errorf("error enabling permissions at %s: %w", path, err)
	}
	return nil
}

// DisableDatasourcePermissions disables the datasource permissions (this is a datasource setting)
func (c *Client) DisableDatasourcePermissions(id int64) error {
	path := fmt.Sprintf("/api/datasources/%d/disable-permissions", id)
	if err := c.request("POST", path, nil, nil, nil); err != nil {
		return fmt.Errorf("error disabling permissions at %s: %w", path, err)
	}
	return nil
}

// DatasourcePermissions fetches and returns the permissions for the datasource whose ID it's passed.
func (c *Client) DatasourcePermissions(id int64) (*DatasourcePermissionsResponse, error) {
	path := fmt.Sprintf("/api/datasources/%d/permissions", id)
	var out *DatasourcePermissionsResponse
	err := c.request("GET", path, nil, nil, &out)
	if err != nil {
		return out, fmt.Errorf("error getting permissions at %s: %w", path, err)
	}

	return out, nil
}

// AddDatasourcePermission adds the given permission item
func (c *Client) AddDatasourcePermission(id int64, item *DatasourcePermissionAddPayload) error {
	path := fmt.Sprintf("/api/datasources/%d/permissions", id)
	data, err := json.Marshal(item)
	if err != nil {
		return fmt.Errorf("marshal err: %w", err)
	}

	if err = c.request("POST", path, nil, bytes.NewBuffer(data), nil); err != nil {
		return fmt.Errorf("error adding permissions at %s: %w", path, err)
	}

	return nil
}

// RemoveDatasourcePermission removes the permission with the given id
func (c *Client) RemoveDatasourcePermission(id, permissionID int64) error {
	path := fmt.Sprintf("/api/datasources/%d/permissions/%d", id, permissionID)
	if err := c.request("DELETE", path, nil, nil, nil); err != nil {
		return fmt.Errorf("error deleting permissions at %s: %w", path, err)
	}

	return nil
}
<|endoftext|>


<commit_before>package db_test

import (
	"fmt"
	"time"

	"github.com/cloudfoundry-incubator/bbs/db/etcd/internal/test_helpers"
	"github.com/cloudfoundry/storeadapter/storerunner/etcdstorerunner"
	etcdclient "github.com/coreos/go-etcd/etcd"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-golang/clock/fakeclock"
	"github.com/pivotal-golang/lager/lagertest"

	"testing"
)

var etcdPort int
var etcdUrl string
var etcdRunner *etcdstorerunner.ETCDClusterRunner
var etcdClient *etcdclient.Client

var logger *lagertest.TestLogger
var clock *fakeclock.FakeClock
var testHelper *test_helpers.TestHelper

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

var _ = BeforeSuite(func() {
	logger = lagertest.NewTestLogger("test")

	clock = fakeclock.NewFakeClock(time.Unix(0, 1138))

	etcdPort = 4001 + GinkgoParallelNode()
	etcdUrl = fmt.Sprintf("http://127.0.0.1:%d", etcdPort)
	etcdRunner = etcdstorerunner.NewETCDClusterRunner(etcdPort, 1, nil)

	etcdRunner.Start()
})

var _ = AfterSuite(func() {
	etcdRunner.Stop()
})

var _ = BeforeEach(func() {
	etcdRunner.Reset()
	etcdClient = etcdRunner.Client()
	testHelper = test_helpers.NewTestHelper(etcdClient)
})
<commit_msg>Set strong consistency for tests also<commit_after>package db_test

import (
	"fmt"
	"time"

	"github.com/cloudfoundry-incubator/bbs/db/etcd/internal/test_helpers"
	"github.com/cloudfoundry/storeadapter/storerunner/etcdstorerunner"
	etcdclient "github.com/coreos/go-etcd/etcd"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-golang/clock/fakeclock"
	"github.com/pivotal-golang/lager/lagertest"

	"testing"
)

var etcdPort int
var etcdUrl string
var etcdRunner *etcdstorerunner.ETCDClusterRunner
var etcdClient *etcdclient.Client

var logger *lagertest.TestLogger
var clock *fakeclock.FakeClock
var testHelper *test_helpers.TestHelper

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

var _ = BeforeSuite(func() {
	logger = lagertest.NewTestLogger("test")

	clock = fakeclock.NewFakeClock(time.Unix(0, 1138))

	etcdPort = 4001 + GinkgoParallelNode()
	etcdUrl = fmt.Sprintf("http://127.0.0.1:%d", etcdPort)
	etcdRunner = etcdstorerunner.NewETCDClusterRunner(etcdPort, 1, nil)

	etcdRunner.Start()
})

var _ = AfterSuite(func() {
	etcdRunner.Stop()
})

var _ = BeforeEach(func() {
	etcdRunner.Reset()
	etcdClient = etcdRunner.Client()
	etcdClient.SetConsistency(etcdclient.STRONG_CONSISTENCY)
	testHelper = test_helpers.NewTestHelper(etcdClient)
})

<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 cassandra

import (
	"bytes"
	"github.com/go-errors/errors"
	"github.com/ligato/cn-infra/db/sql"
	r "reflect"
	"strings"
	"github.com/ligato/cn-infra/utils/structs"
	"fmt"
)

// PutExpToString converts expression to string & slice of bindings
func PutExpToString(whereCondition sql.Expression, entity interface{}) (sqlStr string, bindings []interface{},
	err error) {

	whereCondtionStr := &toStringVisitor{entity: entity}
	whereCondition.Accept(whereCondtionStr)

	statement, _, err := updateSetExpToString(r.Indirect(r.ValueOf(entity)).Type().Name(), /*TODO extract method / make customizable*/
		entity                                                                             /*, TODO TTL*/)
	if err != nil {
		return "", nil, err
	}
	bindings = structFieldPtrs(entity)
	whereBinding := whereCondtionStr.Binding()
	if whereBinding != nil {
		bindings = append(bindings, whereBinding...)
	}

	return strings.Trim(statement+" WHERE"+whereCondtionStr.String(), " "), bindings, nil
}

// SelectExpToString converts expression to string & slice of bindings
func SelectExpToString(fromWhere sql.Expression) (sqlStr string, bindings []interface{},
	err error) {

	findEntity := &findEntityVisitor{}
	fromWhere.Accept(findEntity)

	fromWhereStr := &toStringVisitor{entity: findEntity.entity}
	fromWhere.Accept(fromWhereStr)

	fieldsStr := selectFields(findEntity.entity)
	if err != nil {
		return "", nil, err
	}
	fromWhereBindings := fromWhereStr.Binding()

	return "SELECT " + fieldsStr + " " + fromWhereStr.String(), fromWhereBindings, nil
}

// ExpToString converts expression to string & slice of bindings
func ExpToString(exp sql.Expression) (sql string, bindings []interface{}, err error) {
	findEntity := &findEntityVisitor{}
	exp.Accept(findEntity)

	stringer := &toStringVisitor{entity: findEntity.entity}
	exp.Accept(stringer)

	return stringer.String(), stringer.Binding(), stringer.lastError
}

type toStringVisitor struct {
	entity    interface{}
	generated bytes.Buffer
	binding   []interface{}
	lastError error
}

// String converts generated byte Buffer to string
func (visitor *toStringVisitor) String() string {
	return visitor.generated.String()
}

// Binding is a getter...
func (visitor *toStringVisitor) Binding() []interface{} {
	return visitor.binding
}

// VisitPrefixedExp generates part of SQL expression
func (visitor *toStringVisitor) VisitPrefixedExp(exp *sql.PrefixedExp) {
	visitor.generated.WriteString(" ")
	visitor.generated.WriteString(exp.Prefix)
	if exp.Prefix == "FROM" {
		visitor.generated.WriteString(" ")
		visitor.generated.WriteString(entityName(visitor.entity))
	}
	if exp.AfterPrefix != nil {
		exp.AfterPrefix.Accept(visitor)
	}
	visitor.generated.WriteString(exp.Suffix)

	if exp.Prefix != "FROM" && exp.Binding != nil && len(exp.Binding) > 0 {
		if visitor.binding != nil {
			visitor.binding = append(visitor.binding, exp.Binding)
		} else {
			visitor.binding = exp.Binding
		}
	}
}
func entityName(entity interface{}) string {
	return r.Indirect(r.ValueOf(entity)).Type().Name()
}

// VisitPrefixedExp generates part of SQL expression
func (visitor *toStringVisitor) VisitFieldExpression(exp *sql.FieldExpression) {
	if visitor.entity == nil {
		visitor.lastError = errors.New("not found entity")
	} else {
		field, found := structs.FindField(exp.PointerToAField, visitor.entity)
		if !found {
			visitor.lastError = errors.New("not found field in entity")
			return
		}
		fieldName, found := fieldName(field)
		if !found {
			visitor.lastError = errors.New("not exported field in entity")
			return
		}
		visitor.generated.WriteString(" ")
		visitor.generated.WriteString(fieldName)

		if exp.AfterField != nil {
			exp.AfterField.Accept(visitor)
		}
	}
}

// fieldName checks the cql tag in StructField and parses the field name
func fieldName(field *r.StructField) (name string, exported bool) {
	cql := field.Tag.Get("cql")
	if len(cql) > 0 {
		if cql == "-" {
			return cql, false
		}
		return cql, true
	}
	return field.Name, true
}

// selectFields generates comma separated field names string
func selectFields(val interface{} /*, opts Options*/) (statement string) {
	fields := structs.ListExportedFields(val)
	ret := bytes.Buffer{}
	first := true
	for _, field := range fields {
		if first {
			first = false
		} else {
			ret.WriteString(", ")
		}
		fieldName, exported := fieldName(field)
		if exported {
			ret.WriteString(fieldName)
		}
	}

	return ret.String()
}

// SliceOfFields generates slice of translated (cql tag) field names
func sliceOfFields(val interface{} /*, opts Options*/) (fieldNames []string) {
	fields := structs.ListExportedFields(val)
	fieldNames = []string{}
	for _, field := range fields {
		fieldName, exported := fieldName(field)
		if exported {
			fieldNames = append(fieldNames, fieldName)
		}
	}

	return fieldNames
}

// SliceOfFieldsWithVals generates slice of translated (cql tag) field names with field values
func SliceOfFieldsWithVals(val interface{} /*, opts Options*/) (fieldNames []string, vals []interface{}) {
	fields, vals := structs.ListExportedFieldsWithVals(val)

	fieldNames = []string{}
	for _, field := range fields {
		fieldName, exported := fieldName(field)
		if exported {
			fieldNames = append(fieldNames, fieldName)
		}
	}

	return fieldNames, vals
}

// updateSetExpToString generates UPDATE + SET part of SQL statement
// for fields of an entity
func updateSetExpToString(cfName string, val interface{} /*, opts Options*/) (
	statement string, fields []string, err error) {

	fields = sliceOfFields(val)

	statement = updateStatement(cfName, fields)
	return statement, fields, nil
}

// UPDATE keyspace.Movies SET col1 = val1, col2 = val2
func updateStatement(cfName string, fields []string /*, opts Options*/) (statement string) {
	buf := new(bytes.Buffer)
	buf.WriteString(fmt.Sprintf("UPDATE %s ", cfName))

	/*
		// Apply options
		if opts.TTL != 0 {
			buf.WriteString("USING TTL ")
			buf.WriteString(strconv.FormatFloat(opts.TTL.Seconds(), 'f', 0, 64))
			buf.WriteRune(' ')
		}*/

	buf.WriteString("SET ")
	first := true
	for _, fieldName := range fields {
		if !first {
			buf.WriteString(", ")
		} else {
			first = false
		}
		buf.WriteString(fieldName)
		buf.WriteString(` = ?`)
	}

	return buf.String()
}

type findEntityVisitor struct {
	entity interface{}
}

// VisitPrefixedExp checks for "FROM" expression to find out the entity
func (visitor *findEntityVisitor) VisitPrefixedExp(exp *sql.PrefixedExp) {
	if exp.Prefix == "FROM" {
		if len(exp.Binding) == 1 && r.Indirect(r.ValueOf(exp.Binding[0])).Kind() == r.Struct {
			visitor.entity = exp.Binding[0]
		}
	} else if exp.AfterPrefix != nil {
		exp.AfterPrefix.Accept(visitor)
	}
}

// VisitFieldExpression just propagates to AfterFieldExpression
func (visitor *findEntityVisitor) VisitFieldExpression(exp *sql.FieldExpression) {
	if exp.AfterField != nil {
		exp.AfterField.Accept(visitor)
	}
}
<commit_msg>ODPM-419 fix import errors<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 cassandra

import (
	"bytes"
	"errors"
	"github.com/ligato/cn-infra/db/sql"
	r "reflect"
	"strings"
	"github.com/ligato/cn-infra/utils/structs"
	"fmt"
)

// PutExpToString converts expression to string & slice of bindings
func PutExpToString(whereCondition sql.Expression, entity interface{}) (sqlStr string, bindings []interface{},
	err error) {

	whereCondtionStr := &toStringVisitor{entity: entity}
	whereCondition.Accept(whereCondtionStr)

	statement, _, err := updateSetExpToString(r.Indirect(r.ValueOf(entity)).Type().Name(), /*TODO extract method / make customizable*/
		entity                                                                             /*, TODO TTL*/)
	if err != nil {
		return "", nil, err
	}
	bindings = structFieldPtrs(entity)
	whereBinding := whereCondtionStr.Binding()
	if whereBinding != nil {
		bindings = append(bindings, whereBinding...)
	}

	return strings.Trim(statement+" WHERE"+whereCondtionStr.String(), " "), bindings, nil
}

// SelectExpToString converts expression to string & slice of bindings
func SelectExpToString(fromWhere sql.Expression) (sqlStr string, bindings []interface{},
	err error) {

	findEntity := &findEntityVisitor{}
	fromWhere.Accept(findEntity)

	fromWhereStr := &toStringVisitor{entity: findEntity.entity}
	fromWhere.Accept(fromWhereStr)

	fieldsStr := selectFields(findEntity.entity)
	if err != nil {
		return "", nil, err
	}
	fromWhereBindings := fromWhereStr.Binding()

	return "SELECT " + fieldsStr + " " + fromWhereStr.String(), fromWhereBindings, nil
}

// ExpToString converts expression to string & slice of bindings
func ExpToString(exp sql.Expression) (sql string, bindings []interface{}, err error) {
	findEntity := &findEntityVisitor{}
	exp.Accept(findEntity)

	stringer := &toStringVisitor{entity: findEntity.entity}
	exp.Accept(stringer)

	return stringer.String(), stringer.Binding(), stringer.lastError
}

type toStringVisitor struct {
	entity    interface{}
	generated bytes.Buffer
	binding   []interface{}
	lastError error
}

// String converts generated byte Buffer to string
func (visitor *toStringVisitor) String() string {
	return visitor.generated.String()
}

// Binding is a getter...
func (visitor *toStringVisitor) Binding() []interface{} {
	return visitor.binding
}

// VisitPrefixedExp generates part of SQL expression
func (visitor *toStringVisitor) VisitPrefixedExp(exp *sql.PrefixedExp) {
	visitor.generated.WriteString(" ")
	visitor.generated.WriteString(exp.Prefix)
	if exp.Prefix == "FROM" {
		visitor.generated.WriteString(" ")
		visitor.generated.WriteString(entityName(visitor.entity))
	}
	if exp.AfterPrefix != nil {
		exp.AfterPrefix.Accept(visitor)
	}
	visitor.generated.WriteString(exp.Suffix)

	if exp.Prefix != "FROM" && exp.Binding != nil && len(exp.Binding) > 0 {
		if visitor.binding != nil {
			visitor.binding = append(visitor.binding, exp.Binding)
		} else {
			visitor.binding = exp.Binding
		}
	}
}
func entityName(entity interface{}) string {
	return r.Indirect(r.ValueOf(entity)).Type().Name()
}

// VisitPrefixedExp generates part of SQL expression
func (visitor *toStringVisitor) VisitFieldExpression(exp *sql.FieldExpression) {
	if visitor.entity == nil {
		visitor.lastError = errors.New("not found entity")
	} else {
		field, found := structs.FindField(exp.PointerToAField, visitor.entity)
		if !found {
			visitor.lastError = errors.New("not found field in entity")
			return
		}
		fieldName, found := fieldName(field)
		if !found {
			visitor.lastError = errors.New("not exported field in entity")
			return
		}
		visitor.generated.WriteString(" ")
		visitor.generated.WriteString(fieldName)

		if exp.AfterField != nil {
			exp.AfterField.Accept(visitor)
		}
	}
}

// fieldName checks the cql tag in StructField and parses the field name
func fieldName(field *r.StructField) (name string, exported bool) {
	cql := field.Tag.Get("cql")
	if len(cql) > 0 {
		if cql == "-" {
			return cql, false
		}
		return cql, true
	}
	return field.Name, true
}

// selectFields generates comma separated field names string
func selectFields(val interface{} /*, opts Options*/) (statement string) {
	fields := structs.ListExportedFields(val)
	ret := bytes.Buffer{}
	first := true
	for _, field := range fields {
		if first {
			first = false
		} else {
			ret.WriteString(", ")
		}
		fieldName, exported := fieldName(field)
		if exported {
			ret.WriteString(fieldName)
		}
	}

	return ret.String()
}

// SliceOfFields generates slice of translated (cql tag) field names
func sliceOfFields(val interface{} /*, opts Options*/) (fieldNames []string) {
	fields := structs.ListExportedFields(val)
	fieldNames = []string{}
	for _, field := range fields {
		fieldName, exported := fieldName(field)
		if exported {
			fieldNames = append(fieldNames, fieldName)
		}
	}

	return fieldNames
}

// SliceOfFieldsWithVals generates slice of translated (cql tag) field names with field values
func SliceOfFieldsWithVals(val interface{} /*, opts Options*/) (fieldNames []string, vals []interface{}) {
	fields, vals := structs.ListExportedFieldsWithVals(val)

	fieldNames = []string{}
	for _, field := range fields {
		fieldName, exported := fieldName(field)
		if exported {
			fieldNames = append(fieldNames, fieldName)
		}
	}

	return fieldNames, vals
}

// updateSetExpToString generates UPDATE + SET part of SQL statement
// for fields of an entity
func updateSetExpToString(cfName string, val interface{} /*, opts Options*/) (
	statement string, fields []string, err error) {

	fields = sliceOfFields(val)

	statement = updateStatement(cfName, fields)
	return statement, fields, nil
}

// UPDATE keyspace.Movies SET col1 = val1, col2 = val2
func updateStatement(cfName string, fields []string /*, opts Options*/) (statement string) {
	buf := new(bytes.Buffer)
	buf.WriteString(fmt.Sprintf("UPDATE %s ", cfName))

	/*
		// Apply options
		if opts.TTL != 0 {
			buf.WriteString("USING TTL ")
			buf.WriteString(strconv.FormatFloat(opts.TTL.Seconds(), 'f', 0, 64))
			buf.WriteRune(' ')
		}*/

	buf.WriteString("SET ")
	first := true
	for _, fieldName := range fields {
		if !first {
			buf.WriteString(", ")
		} else {
			first = false
		}
		buf.WriteString(fieldName)
		buf.WriteString(` = ?`)
	}

	return buf.String()
}

type findEntityVisitor struct {
	entity interface{}
}

// VisitPrefixedExp checks for "FROM" expression to find out the entity
func (visitor *findEntityVisitor) VisitPrefixedExp(exp *sql.PrefixedExp) {
	if exp.Prefix == "FROM" {
		if len(exp.Binding) == 1 && r.Indirect(r.ValueOf(exp.Binding[0])).Kind() == r.Struct {
			visitor.entity = exp.Binding[0]
		}
	} else if exp.AfterPrefix != nil {
		exp.AfterPrefix.Accept(visitor)
	}
}

// VisitFieldExpression just propagates to AfterFieldExpression
func (visitor *findEntityVisitor) VisitFieldExpression(exp *sql.FieldExpression) {
	if exp.AfterField != nil {
		exp.AfterField.Accept(visitor)
	}
}
<|endoftext|>


<commit_before>package dotnetframework_test

import (
	"bytes"
	"dotnetcore/dotnetframework"
	"io/ioutil"
	"os"
	"path/filepath"

	libbuildpack "github.com/cloudfoundry/libbuildpack"
	"github.com/cloudfoundry/libbuildpack/ansicleaner"
	gomock "github.com/golang/mock/gomock"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

//go:generate mockgen -source=dotnetframework.go --destination=mocks_dotnetframework_test.go --package=dotnetframework_test

var _ = Describe("Dotnetframework", func() {
	var (
		err          error
		depDir       string
		buildDir     string
		subject      *dotnetframework.DotnetFramework
		mockCtrl     *gomock.Controller
		mockManifest *MockManifest
		buffer       *bytes.Buffer
		logger       *libbuildpack.Logger
	)

	BeforeEach(func() {
		depDir, err = ioutil.TempDir("", "dotnetcore-buildpack.deps.")
		buildDir, err = ioutil.TempDir("", "dotnetcore-buildpack.build.")
		Expect(err).To(BeNil())

		mockCtrl = gomock.NewController(GinkgoT())
		mockManifest = NewMockManifest(mockCtrl)

		buffer = new(bytes.Buffer)
		logger = libbuildpack.NewLogger(ansicleaner.New(buffer))

		subject = dotnetframework.New(depDir, buildDir, mockManifest, logger)
	})

	AfterEach(func() {
		mockCtrl.Finish()
		Expect(os.RemoveAll(depDir)).To(Succeed())
		Expect(os.RemoveAll(buildDir)).To(Succeed())
	})

	FDescribe("Install", func() {
		Context("Versions installed == [1.2.3, 4.5.6]", func() {
			BeforeEach(func() {
				Expect(os.MkdirAll(filepath.Join(depDir, "shared", "Microsoft.NETCore.App", "1.2.3"), 0755)).To(Succeed())
				Expect(os.MkdirAll(filepath.Join(depDir, "shared", "Microsoft.NETCore.App", "4.5.6"), 0755)).To(Succeed())
			})
			Context("when required version is discovered via .runtimeconfig.json", func() {
				Context("Versions required == [4.5.6]", func() {
					BeforeEach(func() {
						Expect(ioutil.WriteFile(filepath.Join(buildDir, "foo.runtimeconfig.json"),
							[]byte(`{ "runtimeOptions": { "framework": { "name": "Microsoft.NETCore.App", "version": "4.5.6" } } }`), 0644)).To(Succeed())
					})

					It("does not install the framework again", func() {
						Expect(subject.Install()).To(Succeed())
					})
				})

				Context("Versions required == [7.8.9]", func() {
					BeforeEach(func() {
						Expect(ioutil.WriteFile(filepath.Join(buildDir, "foo.runtimeconfig.json"),
							[]byte(`{ "runtimeOptions": { "framework": { "name": "Microsoft.NETCore.App", "version": "7.8.9" } } }`), 0644)).To(Succeed())
					})

					It("installs the additional framework", func() {
						mockManifest.EXPECT().InstallDependency(libbuildpack.Dependency{Name: "dotnet-framework", Version: "7.8.9"}, filepath.Join(depDir, "dotnet"))
						Expect(subject.Install()).To(Succeed())
					})
				})
			})

			Context("when required versions are discovered via restored packages", func() {
				Context("Versions required == [4.5.6]", func() {
					BeforeEach(func() {
						Expect(os.MkdirAll(filepath.Join(depDir, ".nuget", "packages", "microsoft.netcore.app", "4.5.6"), 0755)).To(Succeed())
					})

					It("does not install the framework again", func() {
						Expect(subject.Install()).To(Succeed())
					})
				})

				Context("Versions required == [7.8.9]", func() {
					BeforeEach(func() {
						Expect(os.MkdirAll(filepath.Join(depDir, ".nuget", "packages", "microsoft.netcore.app", "7.8.9"), 0755)).To(Succeed())
					})

					It("installs the additional framework", func() {
						mockManifest.EXPECT().InstallDependency(libbuildpack.Dependency{Name: "dotnet-framework", Version: "7.8.9"}, filepath.Join(depDir, "dotnet"))
						Expect(subject.Install()).To(Succeed())
					})
				})
			})
		})
	})
})
<commit_msg>Unfocus test<commit_after>package dotnetframework_test

import (
	"bytes"
	"dotnetcore/dotnetframework"
	"io/ioutil"
	"os"
	"path/filepath"

	libbuildpack "github.com/cloudfoundry/libbuildpack"
	"github.com/cloudfoundry/libbuildpack/ansicleaner"
	gomock "github.com/golang/mock/gomock"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

//go:generate mockgen -source=dotnetframework.go --destination=mocks_dotnetframework_test.go --package=dotnetframework_test

var _ = Describe("Dotnetframework", func() {
	var (
		err          error
		depDir       string
		buildDir     string
		subject      *dotnetframework.DotnetFramework
		mockCtrl     *gomock.Controller
		mockManifest *MockManifest
		buffer       *bytes.Buffer
		logger       *libbuildpack.Logger
	)

	BeforeEach(func() {
		depDir, err = ioutil.TempDir("", "dotnetcore-buildpack.deps.")
		buildDir, err = ioutil.TempDir("", "dotnetcore-buildpack.build.")
		Expect(err).To(BeNil())

		mockCtrl = gomock.NewController(GinkgoT())
		mockManifest = NewMockManifest(mockCtrl)

		buffer = new(bytes.Buffer)
		logger = libbuildpack.NewLogger(ansicleaner.New(buffer))

		subject = dotnetframework.New(depDir, buildDir, mockManifest, logger)
	})

	AfterEach(func() {
		mockCtrl.Finish()
		Expect(os.RemoveAll(depDir)).To(Succeed())
		Expect(os.RemoveAll(buildDir)).To(Succeed())
	})

	Describe("Install", func() {
		Context("Versions installed == [1.2.3, 4.5.6]", func() {
			BeforeEach(func() {
				Expect(os.MkdirAll(filepath.Join(depDir, "shared", "Microsoft.NETCore.App", "1.2.3"), 0755)).To(Succeed())
				Expect(os.MkdirAll(filepath.Join(depDir, "shared", "Microsoft.NETCore.App", "4.5.6"), 0755)).To(Succeed())
			})
			Context("when required version is discovered via .runtimeconfig.json", func() {
				Context("Versions required == [4.5.6]", func() {
					BeforeEach(func() {
						Expect(ioutil.WriteFile(filepath.Join(buildDir, "foo.runtimeconfig.json"),
							[]byte(`{ "runtimeOptions": { "framework": { "name": "Microsoft.NETCore.App", "version": "4.5.6" } } }`), 0644)).To(Succeed())
					})

					It("does not install the framework again", func() {
						Expect(subject.Install()).To(Succeed())
					})
				})

				Context("Versions required == [7.8.9]", func() {
					BeforeEach(func() {
						Expect(ioutil.WriteFile(filepath.Join(buildDir, "foo.runtimeconfig.json"),
							[]byte(`{ "runtimeOptions": { "framework": { "name": "Microsoft.NETCore.App", "version": "7.8.9" } } }`), 0644)).To(Succeed())
					})

					It("installs the additional framework", func() {
						mockManifest.EXPECT().InstallDependency(libbuildpack.Dependency{Name: "dotnet-framework", Version: "7.8.9"}, filepath.Join(depDir, "dotnet"))
						Expect(subject.Install()).To(Succeed())
					})
				})
			})

			Context("when required versions are discovered via restored packages", func() {
				Context("Versions required == [4.5.6]", func() {
					BeforeEach(func() {
						Expect(os.MkdirAll(filepath.Join(depDir, ".nuget", "packages", "microsoft.netcore.app", "4.5.6"), 0755)).To(Succeed())
					})

					It("does not install the framework again", func() {
						Expect(subject.Install()).To(Succeed())
					})
				})

				Context("Versions required == [7.8.9]", func() {
					BeforeEach(func() {
						Expect(os.MkdirAll(filepath.Join(depDir, ".nuget", "packages", "microsoft.netcore.app", "7.8.9"), 0755)).To(Succeed())
					})

					It("installs the additional framework", func() {
						mockManifest.EXPECT().InstallDependency(libbuildpack.Dependency{Name: "dotnet-framework", Version: "7.8.9"}, filepath.Join(depDir, "dotnet"))
						Expect(subject.Install()).To(Succeed())
					})
				})
			})
		})
	})
})
<|endoftext|>


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

package main

import (
	"os"

	"github.com/blevesearch/bleve"

	log "github.com/couchbaselabs/clog"
)

func (pindex *PIndex) Run(mgr PIndexManager) {
	close := true
	cleanup := true

	var err error = nil

	if pindex.IndexType == "bleve" {
		close, cleanup, err = RunBleveStream(mgr, pindex, pindex.Stream,
			pindex.Impl.(bleve.Index))
		if err != nil {
			log.Printf("error: RunBleveStream, close: %t, cleanup: %t, err: %v",
				close, cleanup, err)
		} else {
			log.Printf("done: RunBleveStream, close: %t, cleanup: %t",
				close, cleanup)
		}
	} else {
		log.Printf("error: PIndex.Run() saw unknown IndexType: %s", pindex.IndexType)
	}

	// NOTE: We expect the PIndexImpl to handle any inflight, concurrent
	// queries, access and Close() correctly with its own locking.
	if close {
		pindex.Impl.Close()
	}

	if cleanup {
		os.RemoveAll(pindex.Path)
	}
}

func RunBleveStream(mgr PIndexManager, pindex *PIndex, stream Stream,
	bindex bleve.Index) (bool, bool, error) {
	for req := range stream {
		var err error

		// TODO: maybe need a more batchy API?  Perhaps, yet another
		// goroutine that clumps up up updates into bigger batches?

		switch req.Op {
		case STREAM_OP_NOOP:
			// Do nothing, so stream source can use NOOP like a ping.
			log.Printf("bleve stream noop, key: %s", string(req.Key))

		case STREAM_OP_UPDATE:
			log.Printf("bleve stream udpate, key: %s", string(req.Key))
			err = bindex.Index(string(req.Key), req.Val)

		case STREAM_OP_DELETE:
			log.Printf("bleve stream delete, key: %s", string(req.Key))
			err = bindex.Delete(string(req.Key))

		case STREAM_OP_FLUSH:
			// TODO: Need to delete all records here.  So, why not
			// implement this the same as rollback to zero?

		case STREAM_OP_ROLLBACK:
			// TODO: Implement partial rollback one day.
			// Implementation sketch: we expect bleve to one day to
			// provide an additional Snapshot() and Rollback() API,
			// where Snapshot() returns some opaque and persistable
			// snapshot ID ("SID"), which cbft can occasionally record
			// into the bleve's Get/SetInternal() "side" storage.  A
			// stream rollback operation then needs to loop through
			// appropriate candidate SID's until a Rollback(SID)
			// succeeds.  Else, we eventually devolve down to
			// restarting/rebuilding everything from scratch or zero.
			//
			// For now, always rollback to zero, in which we close the
			// pindex and have the janitor rebuild from scratch.
			pindex.Impl.Close()
			os.RemoveAll(pindex.Path)

			// First, respond to the stream source (example: the feed)
			// so that it can unblock.
			if req.DoneCh != nil {
				close(req.DoneCh)
			}

			// Because, here the manager/janitor will synchronously
			// ask the feed to close and we don't want a deadlock.
			mgr.ClosePIndex(pindex)
			mgr.Kick("stream-rollback")

			return false, false, nil

		case STREAM_OP_GET_META:
			v, err := bindex.GetInternal(req.Key)
			if req.Misc != nil {
				c, ok := req.Misc.(chan []byte)
				if ok && c != nil && err == nil {
					c <- v
				}
				close(c)
			}

		case STREAM_OP_SET_META:
			err = bindex.SetInternal(req.Key, req.Val)
		}

		if err != nil {
			log.Printf("error: bleve stream, op: %s, req: %#v, err: %v",
				StreamOpNames[req.Op], req, err)
		}

		if req.DoneCh != nil {
			if err != nil {
				req.DoneCh <- err
			}
			close(req.DoneCh)
		}
	}

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

package main

import (
	"os"

	"github.com/blevesearch/bleve"

	log "github.com/couchbaselabs/clog"
)

func (pindex *PIndex) Run(mgr PIndexManager) {
	close := true
	cleanup := true

	var err error = nil

	if pindex.IndexType == "bleve" {
		close, cleanup, err = RunBleveStream(mgr, pindex, pindex.Stream,
			pindex.Impl.(bleve.Index))
		if err != nil {
			log.Printf("error: RunBleveStream, close: %t, cleanup: %t, err: %v",
				close, cleanup, err)
		} else {
			log.Printf("done: RunBleveStream, close: %t, cleanup: %t",
				close, cleanup)
		}
	} else {
		log.Printf("error: PIndex.Run() saw unknown IndexType: %s", pindex.IndexType)
	}

	// NOTE: We expect the PIndexImpl to handle any inflight, concurrent
	// queries, access and Close() correctly with its own locking.
	if close {
		pindex.Impl.Close()
	}

	if cleanup {
		os.RemoveAll(pindex.Path)
	}
}

func RunBleveStream(mgr PIndexManager, pindex *PIndex, stream Stream,
	bindex bleve.Index) (bool, bool, error) {
	for req := range stream {
		var err error

		// TODO: maybe need a more batchy API?  Perhaps, yet another
		// goroutine that clumps up up updates into bigger batches?

		switch req.Op {
		case STREAM_OP_NOOP:
			// Do nothing, so stream source can use NOOP like a ping.
			log.Printf("bleve stream noop, partition: %s, key: %s",
				req.Partition, string(req.Key))

		case STREAM_OP_UPDATE:
			log.Printf("bleve stream udpate, partition: %s, key: %s",
				req.Partition, string(req.Key))

			err = bindex.Index(string(req.Key), req.Val)

		case STREAM_OP_DELETE:
			log.Printf("bleve stream delete, partition: %s, key: %s",
				req.Partition, string(req.Key))

			err = bindex.Delete(string(req.Key))

		case STREAM_OP_FLUSH:
			// TODO: Need to delete all records here.  So, why not
			// implement this the same as rollback to zero?

		case STREAM_OP_ROLLBACK:
			log.Printf("bleve stream rollback, partition: %s",
				req.Partition)

			// TODO: Implement partial rollback one day.
			// Implementation sketch: we expect bleve to one day to
			// provide an additional Snapshot() and Rollback() API,
			// where Snapshot() returns some opaque and persistable
			// snapshot ID ("SID"), which cbft can occasionally record
			// into the bleve's Get/SetInternal() "side" storage.  A
			// stream rollback operation then needs to loop through
			// appropriate candidate SID's until a Rollback(SID)
			// succeeds.  Else, we eventually devolve down to
			// restarting/rebuilding everything from scratch or zero.
			//
			// For now, always rollback to zero, in which we close the
			// pindex and have the janitor rebuild from scratch.
			pindex.Impl.Close()
			os.RemoveAll(pindex.Path)

			// First, respond to the stream source (example: the feed)
			// so that it can unblock.
			if req.DoneCh != nil {
				close(req.DoneCh)
			}

			// Because, here the manager/janitor will synchronously
			// ask the feed to close and we don't want a deadlock.
			mgr.ClosePIndex(pindex)
			mgr.Kick("stream-rollback")

			return false, false, nil

		case STREAM_OP_GET_META:
			log.Printf("bleve stream get-meta, partition: %s, key: %s",
				req.Partition, string(req.Key))

			v, err := bindex.GetInternal(req.Key)
			if req.Misc != nil {
				c, ok := req.Misc.(chan []byte)
				if ok && c != nil && err == nil {
					c <- v
				}
				close(c)
			}

		case STREAM_OP_SET_META:
			log.Printf("bleve stream set-meta, partition: %s, key: %s",
				req.Partition, string(req.Key))

			err = bindex.SetInternal(req.Key, req.Val)
		}

		if err != nil {
			log.Printf("error: bleve stream, op: %s, req: %#v, err: %v",
				StreamOpNames[req.Op], req, err)
		}

		if req.DoneCh != nil {
			if err != nil {
				req.DoneCh <- err
			}
			close(req.DoneCh)
		}
	}

	return true, true, nil
}
<|endoftext|>


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

// The block writer example program is a write-only workload intended to insert
// a large amount of data into cockroach quickly. This example is intended to
// trigger range splits and rebalances.
package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/url"
	"os"
	"sync/atomic"
	"time"

	_ "github.com/lib/pq"

	"github.com/satori/go.uuid"
)

const (
	insertBlockStmt = `INSERT INTO blocks (block_id, writer_id, block_num, raw_bytes) VALUES ($1, $2, $3, $4)`
)

// concurrency = number of concurrent insertion processes.
var concurrency = flag.Int("concurrency", 3, "Number of concurrent writers inserting blocks")

var tolerateErrors = flag.Bool("tolerate-errors", false, "Keep running on error")

// outputInterval = interval at which information is output to console.
var outputInterval = flag.Duration("output-interval", 1*time.Second, "Interval of output")

// Minimum and maximum size of inserted blocks.
var minBlockSizeBytes = flag.Int("min-block-bytes", 256, "Minimum amount of raw data written with each insertion")
var maxBlockSizeBytes = flag.Int("max-block-bytes", 1024, "Maximum amount of raw data written with each insertion")

// numBlocks keeps a global count of successfully written blocks.
var numBlocks uint64

// A blockWriter writes blocks of random data into cockroach in an infinite
// loop.
type blockWriter struct {
	id         string
	blockCount uint64
	db         *sql.DB
	rand       *rand.Rand
}

func newBlockWriter(db *sql.DB) blockWriter {
	source := rand.NewSource(int64(time.Now().UnixNano()))
	return blockWriter{
		db:   db,
		id:   uuid.NewV4().String(),
		rand: rand.New(source),
	}
}

// run is an infinite loop in which the blockWriter continuously attempts to
// write blocks of random data into a table in cockroach DB.
func (bw blockWriter) run(errCh chan<- error) {
	for {
		blockID := bw.rand.Int63()
		blockData := bw.randomBlock()
		bw.blockCount++
		if _, err := bw.db.Exec(insertBlockStmt, blockID, bw.id, bw.blockCount, blockData); err != nil {
			errCh <- fmt.Errorf("error running blockwriter %s: %s", bw.id, err)
		} else {
			atomic.AddUint64(&numBlocks, 1)
		}
	}
}

// randomBlock generates a slice of randomized bytes. Random data is preferred
// to prevent compression in storage.
func (bw blockWriter) randomBlock() []byte {
	blockSize := bw.rand.Intn(*maxBlockSizeBytes-*minBlockSizeBytes) + *minBlockSizeBytes
	blockData := make([]byte, blockSize)
	for i := range blockData {
		blockData[i] = byte(bw.rand.Int() & 0xff)
	}
	return blockData
}

// setupDatabase performs initial setup for the example, creating a database and
// with a single table. If the desired table already exists on the cluster, the
// existing table will be dropped.
func setupDatabase(dbURL string) (*sql.DB, error) {
	parsedURL, err := url.Parse(dbURL)
	if err != nil {
		return nil, err
	}
	parsedURL.Path = "datablocks"

	// Open connection to server and create a database.
	db, err := sql.Open("postgres", parsedURL.String())
	if err != nil {
		return nil, err
	}
	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS datablocks"); err != nil {
		return nil, err
	}

	// Allow a maximum of concurrency+1 connections to the database.
	db.SetMaxOpenConns(*concurrency + 1)

	// Create the initial table for storing blocks.
	if _, err := db.Exec(`
	CREATE TABLE IF NOT EXISTS blocks (
	  block_id BIGINT NOT NULL,
	  writer_id STRING NOT NULL,
	  block_num BIGINT NOT NULL,
	  raw_bytes BYTES NOT NULL,
	  PRIMARY KEY (block_id, writer_id, block_num)
	)`); err != nil {
		return nil, err
	}

	return db, nil
}

var usage = func() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s <db URL>\n\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	flag.Usage = usage
	flag.Parse()

	if flag.NArg() != 1 {
		usage()
		os.Exit(2)
	}

	dbURL := flag.Arg(0)

	if *concurrency < 1 {
		log.Fatalf("Value of 'concurrency' flag (%d) must be greater than or equal to 1", *concurrency)
	}

	if max, min := *maxBlockSizeBytes, *minBlockSizeBytes; max < min {
		log.Fatalf("Value of 'max-block-bytes' (%d) must be greater than or equal to value of 'min-block-bytes' (%d)", max, min)
	}

	var db *sql.DB
	{
		var err error
		for err == nil || *tolerateErrors {
			db, err = setupDatabase(dbURL)
			if err == nil {
				break
			}
			if !*tolerateErrors {
				log.Fatal(err)
			}
		}
	}

	lastNow := time.Now()
	var lastNumDumps uint64
	writers := make([]blockWriter, *concurrency)

	errCh := make(chan error)
	for i := range writers {
		writers[i] = newBlockWriter(db)
		go writers[i].run(errCh)
	}

	var numErr int
	for range time.Tick(*outputInterval) {
		now := time.Now()
		elapsed := time.Since(lastNow)
		dumps := atomic.LoadUint64(&numBlocks)
		log.Printf("%d dumps were executed at %.1f/second (%d total errors)", (dumps - lastNumDumps), float64(dumps-lastNumDumps)/elapsed.Seconds(), numErr)
		for {
			select {
			case err := <-errCh:
				numErr++
				if !*tolerateErrors {
					log.Fatal(err)
				} else {
					log.Print(err)
				}
				continue
			default:
			}
			break
		}
		lastNumDumps = dumps
		lastNow = now
	}
}
<commit_msg>block_writer: simplify status output<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Matt Tracy

// The block writer example program is a write-only workload intended to insert
// a large amount of data into cockroach quickly. This example is intended to
// trigger range splits and rebalances.
package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/url"
	"os"
	"sync/atomic"
	"time"

	_ "github.com/lib/pq"

	"github.com/satori/go.uuid"
)

const (
	insertBlockStmt = `INSERT INTO blocks (block_id, writer_id, block_num, raw_bytes) VALUES ($1, $2, $3, $4)`
)

// concurrency = number of concurrent insertion processes.
var concurrency = flag.Int("concurrency", 3, "Number of concurrent writers inserting blocks")

var tolerateErrors = flag.Bool("tolerate-errors", false, "Keep running on error")

// outputInterval = interval at which information is output to console.
var outputInterval = flag.Duration("output-interval", 1*time.Second, "Interval of output")

// Minimum and maximum size of inserted blocks.
var minBlockSizeBytes = flag.Int("min-block-bytes", 256, "Minimum amount of raw data written with each insertion")
var maxBlockSizeBytes = flag.Int("max-block-bytes", 1024, "Maximum amount of raw data written with each insertion")

// numBlocks keeps a global count of successfully written blocks.
var numBlocks uint64

// A blockWriter writes blocks of random data into cockroach in an infinite
// loop.
type blockWriter struct {
	id         string
	blockCount uint64
	db         *sql.DB
	rand       *rand.Rand
}

func newBlockWriter(db *sql.DB) blockWriter {
	source := rand.NewSource(int64(time.Now().UnixNano()))
	return blockWriter{
		db:   db,
		id:   uuid.NewV4().String(),
		rand: rand.New(source),
	}
}

// run is an infinite loop in which the blockWriter continuously attempts to
// write blocks of random data into a table in cockroach DB.
func (bw blockWriter) run(errCh chan<- error) {
	for {
		blockID := bw.rand.Int63()
		blockData := bw.randomBlock()
		bw.blockCount++
		if _, err := bw.db.Exec(insertBlockStmt, blockID, bw.id, bw.blockCount, blockData); err != nil {
			errCh <- fmt.Errorf("error running blockwriter %s: %s", bw.id, err)
		} else {
			atomic.AddUint64(&numBlocks, 1)
		}
	}
}

// randomBlock generates a slice of randomized bytes. Random data is preferred
// to prevent compression in storage.
func (bw blockWriter) randomBlock() []byte {
	blockSize := bw.rand.Intn(*maxBlockSizeBytes-*minBlockSizeBytes) + *minBlockSizeBytes
	blockData := make([]byte, blockSize)
	for i := range blockData {
		blockData[i] = byte(bw.rand.Int() & 0xff)
	}
	return blockData
}

// setupDatabase performs initial setup for the example, creating a database and
// with a single table. If the desired table already exists on the cluster, the
// existing table will be dropped.
func setupDatabase(dbURL string) (*sql.DB, error) {
	parsedURL, err := url.Parse(dbURL)
	if err != nil {
		return nil, err
	}
	parsedURL.Path = "datablocks"

	// Open connection to server and create a database.
	db, err := sql.Open("postgres", parsedURL.String())
	if err != nil {
		return nil, err
	}
	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS datablocks"); err != nil {
		return nil, err
	}

	// Allow a maximum of concurrency+1 connections to the database.
	db.SetMaxOpenConns(*concurrency + 1)

	// Create the initial table for storing blocks.
	if _, err := db.Exec(`
	CREATE TABLE IF NOT EXISTS blocks (
	  block_id BIGINT NOT NULL,
	  writer_id STRING NOT NULL,
	  block_num BIGINT NOT NULL,
	  raw_bytes BYTES NOT NULL,
	  PRIMARY KEY (block_id, writer_id, block_num)
	)`); err != nil {
		return nil, err
	}

	return db, nil
}

var usage = func() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s <db URL>\n\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	flag.Usage = usage
	flag.Parse()

	if flag.NArg() != 1 {
		usage()
		os.Exit(2)
	}

	dbURL := flag.Arg(0)

	if *concurrency < 1 {
		log.Fatalf("Value of 'concurrency' flag (%d) must be greater than or equal to 1", *concurrency)
	}

	if max, min := *maxBlockSizeBytes, *minBlockSizeBytes; max < min {
		log.Fatalf("Value of 'max-block-bytes' (%d) must be greater than or equal to value of 'min-block-bytes' (%d)", max, min)
	}

	var db *sql.DB
	{
		var err error
		for err == nil || *tolerateErrors {
			db, err = setupDatabase(dbURL)
			if err == nil {
				break
			}
			if !*tolerateErrors {
				log.Fatal(err)
			}
		}
	}

	lastNow := time.Now()
	start := lastNow
	var lastNumDumps uint64
	writers := make([]blockWriter, *concurrency)

	errCh := make(chan error)
	for i := range writers {
		writers[i] = newBlockWriter(db)
		go writers[i].run(errCh)
	}

	var numErr int
	for range time.Tick(*outputInterval) {
		now := time.Now()
		elapsed := time.Since(lastNow)
		dumps := atomic.LoadUint64(&numBlocks)
		fmt.Printf("%6s: %6.1f/sec",
			time.Duration(time.Since(start).Seconds()+0.5)*time.Second,
			float64(dumps-lastNumDumps)/elapsed.Seconds())
		if numErr > 0 {
			fmt.Printf(" (%d total errors)\n", numErr)
		}
		fmt.Printf("\n")
		for {
			select {
			case err := <-errCh:
				numErr++
				if !*tolerateErrors {
					log.Fatal(err)
				} else {
					log.Print(err)
				}
				continue
			default:
			}
			break
		}
		lastNumDumps = dumps
		lastNow = now
	}
}
<|endoftext|>


<commit_before>package via

import (
	"github.com/str1ngs/util/json"
	"os"
	"path"
	"sort"
	"strings"
)

var (
	cache  Cache
	home   = os.Getenv("HOME")
	cfile  = path.Join(home, "via", "plans", "config.json")
	config = new(Config)
)

func init() {
	err := json.Read(&config, cfile)
	if err != nil {
		elog.Fatal(err)
	}

	sort.Strings([]string(config.Flags))
	sort.Strings(config.Remove)
	// TODO: provide Lint for master config
	err = json.Write(&config, cfile)
	if err != nil {
		elog.Fatal(err)
	}
	cache = Cache(os.ExpandEnv(string(config.Cache)))
	config.Plans = os.ExpandEnv(config.Plans)
	config.Repo = os.ExpandEnv(config.Repo)
	for i, j := range config.Env {
		os.Setenv(i, j)
	}
}

type Config struct {
	Identity  string
	Arch      string
	OS        string
	Root      string
	PlansRepo string

	// Paths
	Cache Cache
	DB    DB
	Plans string
	Repo  string

	// Toolchain
	Flags Flags

	Env    map[string]string
	Remove []string
}

type Flags []string

func (f Flags) String() string {
	return strings.Join(f, " ")
}

type DB string

func (d DB) Installed() string {
	return join(config.Root, string(d), "installed")
}

func (d DB) Plans() string {
	return join(config.Root, string(d), "plans")
}

type Cache string

func (c Cache) Pkgs() string {
	return path.Join(string(c), "pkg")
}

func (c Cache) Srcs() string {
	return path.Join(string(c), "src")
}

func (c Cache) Builds() string {
	return path.Join(string(c), "bld")
}

func (c Cache) Stages() string {
	return path.Join(string(c), "stg")
}
<commit_msg>expand Env variables<commit_after>package via

import (
	"github.com/str1ngs/util/json"
	"os"
	"path"
	"sort"
	"strings"
)

var (
	cache  Cache
	home   = os.Getenv("HOME")
	cfile  = path.Join(home, "via", "plans", "config.json")
	config = new(Config)
)

func init() {
	err := json.Read(&config, cfile)
	if err != nil {
		elog.Fatal(err)
	}

	sort.Strings([]string(config.Flags))
	sort.Strings(config.Remove)
	// TODO: provide Lint for master config
	err = json.Write(&config, cfile)
	if err != nil {
		elog.Fatal(err)
	}
	cache = Cache(os.ExpandEnv(string(config.Cache)))
	config.Plans = os.ExpandEnv(config.Plans)
	config.Repo = os.ExpandEnv(config.Repo)
	for i, j := range config.Env {
		os.Setenv(i, os.ExpandEnv(j))
	}
}

type Config struct {
	Identity  string
	Arch      string
	OS        string
	Root      string
	PlansRepo string

	// Paths
	Cache Cache
	DB    DB
	Plans string
	Repo  string

	// Toolchain
	Flags Flags

	Env    map[string]string
	Remove []string
}

type Flags []string

func (f Flags) String() string {
	return strings.Join(f, " ")
}

type DB string

func (d DB) Installed() string {
	return join(config.Root, string(d), "installed")
}

func (d DB) Plans() string {
	return join(config.Root, string(d), "plans")
}

type Cache string

func (c Cache) Pkgs() string {
	return path.Join(string(c), "pkg")
}

func (c Cache) Srcs() string {
	return path.Join(string(c), "src")
}

func (c Cache) Builds() string {
	return path.Join(string(c), "bld")
}

func (c Cache) Stages() string {
	return path.Join(string(c), "stg")
}
<|endoftext|>


<commit_before>// Copyright Istio Authors. 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 pkg

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path"
	"strings"

	"istio.io/release-builder/pkg/model"
	"istio.io/release-builder/pkg/util"

	"istio.io/pkg/log"
)

// Sources will copy all dependencies require, pulling from Github if required, and set up the working tree.
// This includes locally tagging all git repos with the version being built, so that the right version is present in binaries.
func Sources(manifest model.Manifest) error {
	// Clone istio first, as it is needed to determine which other dependencies to use
	if err := cloneRepo(manifest, "istio", manifest.Dependencies.Istio); err != nil {
		return err
	}

	for repo, dependency := range manifest.Dependencies.Get() {
		if repo == "istio" {
			continue
		}
		if dependency == nil {
			log.Warnf("skipping clone of missing dependency: %v", repo)
			continue
		}
		if err := cloneRepo(manifest, repo, dependency); err != nil {
			return err
		}
	}
	return nil
}

func cloneRepo(manifest model.Manifest, repo string, dependency *model.Dependency) error {
	src := path.Join(manifest.SourceDir(), repo)
	// Fetch the dependency
	if err := util.Clone(repo, *dependency, src); err != nil {
		return fmt.Errorf("failed to resolve %+v: %v", dependency, err)
	}
	log.Infof("Resolved %v", repo)
	// Also copy it to the working directory
	if err := util.CopyDir(src, manifest.RepoDir(repo)); err != nil {
		return fmt.Errorf("failed to copy dependency %v to working directory: %v", repo, err)
	}
	// Tag the repo. This allows the build process to look at the git tag for version information
	if err := TagRepo(manifest, manifest.RepoDir(repo)); err != nil {
		return fmt.Errorf("failed to tag repo %v: %v", repo, err)
	}
	return nil
}

// The release expects a working directory with:
// * sources/ contains all of the sources to build from. These should not be modified
// * work/ initially contains all the sources, but may be modified during the build
// * out/ contains all final artifacts
func SetupWorkDir(dir string) error {
	if err := os.Mkdir(path.Join(dir, "sources"), 0750); err != nil {
		return fmt.Errorf("failed to set up working directory: %v", err)
	}
	if err := os.Mkdir(path.Join(dir, "work"), 0750); err != nil {
		return fmt.Errorf("failed to set up working directory: %v", err)
	}
	if err := os.Mkdir(path.Join(dir, "out"), 0750); err != nil {
		return fmt.Errorf("failed to set up working directory: %v", err)
	}
	return nil
}

// TagRepo tags a given git repo with the version from the manifest.
func TagRepo(manifest model.Manifest, repo string) error {
	headSha, err := GetSha(repo, "HEAD")
	if err != nil {
		return fmt.Errorf("failed to get HEAD SHA: %v", err)
	}
	currentTagSha, _ := GetSha(repo, manifest.Version)
	if currentTagSha != "" {
		if currentTagSha == headSha {
			log.Infof("Tag %v already exists, but points to the right place.", manifest.Version)
			return nil
		}
		return fmt.Errorf("tag %v already exists, retagging would move from %v to %v", manifest.Version, currentTagSha, headSha)
	}
	cmd := util.VerboseCommand("git", "tag", manifest.Version)
	cmd.Dir = repo
	return cmd.Run()
}

// GetSha returns the SHA for a given reference, or error if sha is not found
func GetSha(repo string, ref string) (string, error) {
	buf := bytes.Buffer{}
	cmd := exec.Command("git", "rev-parse", ref)
	cmd.Stdout = &buf
	cmd.Dir = repo
	if err := cmd.Run(); err != nil {
		return "", err
	}
	return buf.String(), nil
}

// StandardizeManifest will convert a manifest to a fixed SHA, rather than a branch
// This allows outputting the exact version used after the build is complete
func StandardizeManifest(manifest *model.Manifest) error {
	for repo, dep := range manifest.Dependencies.Get() {
		if dep == nil {
			continue
		}
		sha, err := GetSha(manifest.RepoDir(repo), "HEAD")
		if err != nil {
			return fmt.Errorf("failed to get SHA for %v: %v", repo, err)
		}
		newDep := model.Dependency{
			Sha: strings.TrimSpace(sha),
		}
		manifest.Dependencies.Set(repo, newDep)
	}
	return nil
}
<commit_msg>Fix client go not being tagged with go version (#113)<commit_after>// Copyright Istio Authors. 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 pkg

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path"
	"strings"

	"istio.io/release-builder/pkg/model"
	"istio.io/release-builder/pkg/util"

	"istio.io/pkg/log"
)

// Sources will copy all dependencies require, pulling from Github if required, and set up the working tree.
// This includes locally tagging all git repos with the version being built, so that the right version is present in binaries.
func Sources(manifest model.Manifest) error {
	// Clone istio first, as it is needed to determine which other dependencies to use
	if err := cloneRepo(manifest, "istio", manifest.Dependencies.Istio); err != nil {
		return err
	}

	for repo, dependency := range manifest.Dependencies.Get() {
		if repo == "istio" {
			continue
		}
		if dependency == nil {
			log.Warnf("skipping clone of missing dependency: %v", repo)
			continue
		}
		if err := cloneRepo(manifest, repo, dependency); err != nil {
			return err
		}
	}
	return nil
}

func cloneRepo(manifest model.Manifest, repo string, dependency *model.Dependency) error {
	src := path.Join(manifest.SourceDir(), repo)
	// Fetch the dependency
	if err := util.Clone(repo, *dependency, src); err != nil {
		return fmt.Errorf("failed to resolve %+v: %v", dependency, err)
	}
	log.Infof("Resolved %v", repo)
	// Also copy it to the working directory
	if err := util.CopyDir(src, manifest.RepoDir(repo)); err != nil {
		return fmt.Errorf("failed to copy dependency %v to working directory: %v", repo, err)
	}
	// Tag the repo. This allows the build process to look at the git tag for version information
	if err := TagRepo(manifest, manifest.RepoDir(repo)); err != nil {
		return fmt.Errorf("failed to tag repo %v: %v", repo, err)
	}
	return nil
}

// The release expects a working directory with:
// * sources/ contains all of the sources to build from. These should not be modified
// * work/ initially contains all the sources, but may be modified during the build
// * out/ contains all final artifacts
func SetupWorkDir(dir string) error {
	if err := os.Mkdir(path.Join(dir, "sources"), 0750); err != nil {
		return fmt.Errorf("failed to set up working directory: %v", err)
	}
	if err := os.Mkdir(path.Join(dir, "work"), 0750); err != nil {
		return fmt.Errorf("failed to set up working directory: %v", err)
	}
	if err := os.Mkdir(path.Join(dir, "out"), 0750); err != nil {
		return fmt.Errorf("failed to set up working directory: %v", err)
	}
	return nil
}

// TagRepo tags a given git repo with the version from the manifest.
func TagRepo(manifest model.Manifest, repo string) error {
	headSha, err := GetSha(repo, "HEAD")
	if err != nil {
		return fmt.Errorf("failed to get HEAD SHA: %v", err)
	}
	currentTagSha, _ := GetSha(repo, manifest.Version)
	if currentTagSha != "" {
		if currentTagSha == headSha {
			log.Infof("Tag %v already exists, but points to the right place.", manifest.Version)
			return nil
		}
		return fmt.Errorf("tag %v already exists, retagging would move from %v to %v", manifest.Version, currentTagSha, headSha)
	}
	cmd := util.VerboseCommand("git", "tag", manifest.Version)
	cmd.Dir = repo
	return cmd.Run()
}

// GetSha returns the SHA for a given reference, or error if sha is not found
func GetSha(repo string, ref string) (string, error) {
	buf := bytes.Buffer{}
	cmd := exec.Command("git", "rev-parse", ref)
	cmd.Stdout = &buf
	cmd.Dir = repo
	if err := cmd.Run(); err != nil {
		return "", err
	}
	return buf.String(), nil
}

// StandardizeManifest will convert a manifest to a fixed SHA, rather than a branch
// This allows outputting the exact version used after the build is complete
func StandardizeManifest(manifest *model.Manifest) error {
	for repo, dep := range manifest.Dependencies.Get() {
		if dep == nil {
			continue
		}
		sha, err := GetSha(manifest.RepoDir(repo), "HEAD")
		if err != nil {
			return fmt.Errorf("failed to get SHA for %v: %v", repo, err)
		}
		newDep := model.Dependency{
			Sha:              strings.TrimSpace(sha),
			GoVersionEnabled: dep.GoVersionEnabled,
		}
		manifest.Dependencies.Set(repo, newDep)
	}
	return nil
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
package wmf

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/url"
	"testing"
)

func Test_digitsOnly(t *testing.T) {
	if digitsOnly('0') != '0' {
		t.Errorf("Failed to recognize digit")
	}
	if digitsOnly('a') != -1 {
		t.Errorf("Failed to discard non-digit")
	}
}

func Test_asciiOnly(t *testing.T) {
	if asciiOnly(' ') != ' ' {
		t.Error("Failed to recognize ASCII")
	}
	if asciiOnly('\a') != -1 {
		t.Error("Failed to discard control character")
	}
	if asciiOnly('☃') != -1 {
		t.Error("Failed to discard UTF8 extended character")
	}
}

func Test_deviceIdFilter(t *testing.T) {
	if deviceIdFilter('A') != 'A' {
		t.Error("Failed to accept valid DeviceID character")
	}
	if deviceIdFilter('+') != rune(-1) {
		t.Error("Failed to reject invalid DeviceID character")
	}
}

func Test_assertionFilter(t *testing.T) {
	if deviceIdFilter('A') != 'A' {
		t.Error("Failed to accept valid Assertion character")
	}
	if deviceIdFilter('+') != rune(-1) {
		t.Error("Failed to reject invalid Assertion character")
	}
}

func Test_parseBody(t *testing.T) {
	t_good := struct {
		Foo int64 `json:"foo"`
	}{123}
	ts, _ := json.Marshal(t_good)
	tr := ioutil.NopCloser(bytes.NewBuffer(ts))

	reply, body, err := parseBody(tr)
	if reply["foo"].(float64) != 123 {
		t.Error("foo not found or invalid")
	}
	if body != string(ts) {
		t.Error("body does not match expected string")
	}
	t_bad := ioutil.NopCloser(bytes.NewBuffer([]byte("{\"InvalidJson\",}")))
	reply, body, err = parseBody(t_bad)
	if err == nil {
		t.Error("Failed to catch bad JSON")
	}
}

func Test_isTrue(t *testing.T) {
	if !isTrue("true") {
		t.Error("\"True\" not true")
	}
	if !isTrue(1) {
		t.Error("1 not true")
	}
	if !isTrue(true) {
		t.Error("true not true")
	}
	if isTrue("Banana") {
		t.Error("\"Banana\" is true")
	}
	if isTrue("  False") {
		t.Error("\"  False\" is true")
	}
}

func Test_minInt(t *testing.T) {
	if minInt(10, 100) != 10 {
		t.Error("minInt returned wrong value")
	}
}

func Test_getDevFromUrl(t *testing.T) {
	var r string
	u, _ := url.Parse("http://")
	if r = getDevFromUrl(u); r != "" {
		t.Error("Bad parse of short url.")
	}
	u, _ = url.Parse("http://box/0123456789abcdef/")
	if r = getDevFromUrl(u); r != "0123456789abcdef" {
		t.Error("Failed to find slashed deviceid. %s", r)
	}
	u, _ = url.Parse("http://box/0123456789abcdef")
	if r = getDevFromUrl(u); r != "0123456789abcdef" {
		t.Error("Failed to find deviceid. %s", r)
	}
	u, _ = url.Parse("http://box/1234567890123456789012345678901234567890")
	if r = getDevFromUrl(u); r != "12345678901234567890123456789012" {
		t.Error("Failed to truncate to 32 characters:%s", r)
	}
	u, _ = url.Parse("http://box/DeadBeefWRONG")
	if r = getDevFromUrl(u); r != "DeadBeef" {
		t.Error("Failed to trim bad characters:%s", r)
	}
}
<commit_msg>wmf: vet fixes<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
package wmf

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/url"
	"testing"
)

func Test_digitsOnly(t *testing.T) {
	if digitsOnly('0') != '0' {
		t.Errorf("Failed to recognize digit")
	}
	if digitsOnly('a') != -1 {
		t.Errorf("Failed to discard non-digit")
	}
}

func Test_asciiOnly(t *testing.T) {
	if asciiOnly(' ') != ' ' {
		t.Error("Failed to recognize ASCII")
	}
	if asciiOnly('\a') != -1 {
		t.Error("Failed to discard control character")
	}
	if asciiOnly('☃') != -1 {
		t.Error("Failed to discard UTF8 extended character")
	}
}

func Test_deviceIdFilter(t *testing.T) {
	if deviceIdFilter('A') != 'A' {
		t.Error("Failed to accept valid DeviceID character")
	}
	if deviceIdFilter('+') != rune(-1) {
		t.Error("Failed to reject invalid DeviceID character")
	}
}

func Test_assertionFilter(t *testing.T) {
	if deviceIdFilter('A') != 'A' {
		t.Error("Failed to accept valid Assertion character")
	}
	if deviceIdFilter('+') != rune(-1) {
		t.Error("Failed to reject invalid Assertion character")
	}
}

func Test_parseBody(t *testing.T) {
	t_good := struct {
		Foo int64 `json:"foo"`
	}{123}
	ts, _ := json.Marshal(t_good)
	tr := ioutil.NopCloser(bytes.NewBuffer(ts))

	reply, body, err := parseBody(tr)
	if reply["foo"].(float64) != 123 {
		t.Error("foo not found or invalid")
	}
	if body != string(ts) {
		t.Error("body does not match expected string")
	}
	t_bad := ioutil.NopCloser(bytes.NewBuffer([]byte("{\"InvalidJson\",}")))
	reply, body, err = parseBody(t_bad)
	if err == nil {
		t.Error("Failed to catch bad JSON")
	}
}

func Test_isTrue(t *testing.T) {
	if !isTrue("true") {
		t.Error("\"True\" not true")
	}
	if !isTrue(1) {
		t.Error("1 not true")
	}
	if !isTrue(true) {
		t.Error("true not true")
	}
	if isTrue("Banana") {
		t.Error("\"Banana\" is true")
	}
	if isTrue("  False") {
		t.Error("\"  False\" is true")
	}
}

func Test_minInt(t *testing.T) {
	if minInt(10, 100) != 10 {
		t.Error("minInt returned wrong value")
	}
}

func Test_getDevFromUrl(t *testing.T) {
	var r string
	u, _ := url.Parse("http://")
	if r = getDevFromUrl(u); r != "" {
		t.Error("Bad parse of short url.")
	}
	u, _ = url.Parse("http://box/0123456789abcdef/")
	if r = getDevFromUrl(u); r != "0123456789abcdef" {
		t.Errorf("Failed to find slashed deviceid: %s", r)
	}
	u, _ = url.Parse("http://box/0123456789abcdef")
	if r = getDevFromUrl(u); r != "0123456789abcdef" {
		t.Errorf("Failed to find deviceid: %s", r)
	}
	u, _ = url.Parse("http://box/1234567890123456789012345678901234567890")
	if r = getDevFromUrl(u); r != "12345678901234567890123456789012" {
		t.Errorf("Failed to truncate to 32 characters: %s", r)
	}
	u, _ = url.Parse("http://box/DeadBeefWRONG")
	if r = getDevFromUrl(u); r != "DeadBeef" {
		t.Errorf("Failed to trim bad characters: %s", r)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/gorilla/websocket"
)

func IsWebsocket(r *http.Request) bool {
	contains := func(headers []string, part string) bool {
		for _, value := range headers {
			for _, token := range strings.Split(value, ",") {
				if strings.EqualFold(strings.TrimSpace(token), part) {
					return true
				}
			}
		}
		return false
	}

	if !contains(r.Header["Connection"], "upgrade") {
		return false
	}

	if !contains(r.Header["Upgrade"], "websocket") {
		return false
	}

	return true
}

type WebsocketCapableReverseProxy struct {
	*httputil.ReverseProxy

	target *url.URL
}

func NewWebsocketCapableReverseProxy(url *url.URL) *WebsocketCapableReverseProxy {
	return &WebsocketCapableReverseProxy{
		httputil.NewSingleHostReverseProxy(url),
		url,
	}
}

func (p *WebsocketCapableReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if IsWebsocket(r) {
		p.ServeWebsocket(w, r)
	} else {
		p.ReverseProxy.ServeHTTP(w, r)
	}
}

func (p *WebsocketCapableReverseProxy) ServeWebsocket(w http.ResponseWriter, r *http.Request) {

	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *r // includes shallow copies of maps, but okay

	p.Director(outreq)

	// Note: Director rewrites outreq.URL.Host, but we need it to be the
	// internal host for the websocket dial. The Host: header gets set to the
	// inbound http request's `Host` header.
	outreq.URL.Host = p.target.Host

	switch outreq.URL.Scheme {
	case "http", "":
		outreq.URL.Scheme = "ws"
	case "https":
		outreq.URL.Scheme = "wss"
	}

	if clientIP, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	outreq.Header.Set("Host", r.Host)

	log.Printf("Establishing outbound websocket to %v", outreq.URL.String())

	dial := websocket.DefaultDialer.Dial
	outConn, resp, err := dial(outreq.URL.String(), outreq.Header)
	if err != nil {
		if resp != nil {
			log.Printf("outbound websocket dial error, status: %v, err: %v",
				resp.StatusCode, err)
			w.WriteHeader(resp.StatusCode)
			_, err := io.Copy(w, resp.Body)
			if err != nil {
				log.Printf("error copying outbound body to response. err: %v", err)
			}
		} else {
			log.Printf("outbound websocket dial error, err: %v", err)
			w.WriteHeader(502)
			fmt.Fprintln(w, "Bad Gateway")
		}
		return
	}
	defer outConn.Close()

	inConn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Failed to upgrade: %v", err)
		// Don't send any response here, Upgrade already does that on error.
		return
	}
	defer inConn.Close()

	finish := make(chan struct{})
	defer func() { <-finish }()

	rawIn := inConn.UnderlyingConn()
	rawOut := outConn.UnderlyingConn()

	go func() {
		defer close(finish)
		_, _ = io.Copy(rawOut, rawIn)
	}()

	_, _ = io.Copy(rawIn, rawOut)
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}
<commit_msg>Use http.Error instead<commit_after>package main

import (
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/gorilla/websocket"
)

func IsWebsocket(r *http.Request) bool {
	contains := func(headers []string, part string) bool {
		for _, value := range headers {
			for _, token := range strings.Split(value, ",") {
				if strings.EqualFold(strings.TrimSpace(token), part) {
					return true
				}
			}
		}
		return false
	}

	if !contains(r.Header["Connection"], "upgrade") {
		return false
	}

	if !contains(r.Header["Upgrade"], "websocket") {
		return false
	}

	return true
}

type WebsocketCapableReverseProxy struct {
	*httputil.ReverseProxy

	target *url.URL
}

func NewWebsocketCapableReverseProxy(url *url.URL) *WebsocketCapableReverseProxy {
	return &WebsocketCapableReverseProxy{
		httputil.NewSingleHostReverseProxy(url),
		url,
	}
}

func (p *WebsocketCapableReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if IsWebsocket(r) {
		p.ServeWebsocket(w, r)
	} else {
		p.ReverseProxy.ServeHTTP(w, r)
	}
}

func (p *WebsocketCapableReverseProxy) ServeWebsocket(w http.ResponseWriter, r *http.Request) {

	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *r // includes shallow copies of maps, but okay

	p.Director(outreq)

	// Note: Director rewrites outreq.URL.Host, but we need it to be the
	// internal host for the websocket dial. The Host: header gets set to the
	// inbound http request's `Host` header.
	outreq.URL.Host = p.target.Host

	switch outreq.URL.Scheme {
	case "http", "":
		outreq.URL.Scheme = "ws"
	case "https":
		outreq.URL.Scheme = "wss"
	}

	if clientIP, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	outreq.Header.Set("Host", r.Host)

	log.Printf("Establishing outbound websocket to %v", outreq.URL.String())

	dial := websocket.DefaultDialer.Dial
	outConn, resp, err := dial(outreq.URL.String(), outreq.Header)
	if err != nil {
		if resp != nil {
			log.Printf("outbound websocket dial error, status: %v, err: %v",
				resp.StatusCode, err)
			w.WriteHeader(resp.StatusCode)
			_, err := io.Copy(w, resp.Body)
			if err != nil {
				log.Printf("error copying outbound body to response. err: %v", err)
			}
		} else {
			log.Printf("outbound websocket dial error, err: %v", err)
			http.Error(w, "502 Bad Gateway", http.StatusBadGateway)
		}
		return
	}
	defer outConn.Close()

	inConn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Failed to upgrade: %v", err)
		// Don't send any response here, Upgrade already does that on error.
		return
	}
	defer inConn.Close()

	finish := make(chan struct{})
	defer func() { <-finish }()

	rawIn := inConn.UnderlyingConn()
	rawOut := outConn.UnderlyingConn()

	go func() {
		defer close(finish)
		_, _ = io.Copy(rawOut, rawIn)
	}()

	_, _ = io.Copy(rawIn, rawOut)
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}
<|endoftext|>


<commit_before><commit_msg>fix route controller tryCreateRoute bug<commit_after><|endoftext|>