<commit_before>package isolated

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"regexp"

	. "code.cloudfoundry.org/cli/cf/util/testhelpers/matchers"
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("apply-manifest command", func() {
	var (
		orgName      string
		spaceName    string
		appName      string
		manifestPath string
		appDir       string
	)

	BeforeEach(func() {
		orgName = helpers.NewOrgName()
		spaceName = helpers.NewSpaceName()
		appName = helpers.PrefixedRandomName("app")

		appDir = helpers.TempDirAbsolutePath("", "simple-app")

		manifestPath = filepath.Join(appDir, "manifest.yml")
		// Ensure the file exists at the minimum
		helpers.WriteManifest(manifestPath, map[string]interface{}{})
	})

	AfterEach(func() {
		Expect(os.RemoveAll(appDir)).ToNot(HaveOccurred())
	})

	Describe("help", func() {
		It("appears in cf help -a", func() {
			session := helpers.CF("help", "-a")
			Eventually(session).Should(Exit(0))
			Expect(session).To(HaveCommandInCategoryWithDescription("apply-manifest", "SPACES", "Apply manifest properties to a space"))
		})

		When("--help flag is set", func() {
			It("displays command usage to output", func() {
				session := helpers.CF("apply-manifest", "--help")

				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Say("apply-manifest - Apply manifest properties to a space"))
				Eventually(session).Should(Say("USAGE:"))
				Eventually(session).Should(Say("cf apply-manifest -f APP_MANIFEST_PATH"))
				Eventually(session).Should(Say("SEE ALSO:"))
				Eventually(session).Should(Say("create-app, create-app-manifest, push"))

				Eventually(session).Should(Exit(0))
			})
		})
	})

	When("the environment is not setup correctly", func() {
		It("fails with the appropriate errors", func() {
			helpers.CheckEnvironmentTargetedCorrectly(true, true, ReadOnlyOrg, "apply-manifest", "-f", manifestPath)
		})
	})

	When("the environment is set up correctly", func() {
		BeforeEach(func() {
			helpers.SetupCF(orgName, spaceName)
		})

		AfterEach(func() {
			helpers.QuickDeleteOrg(orgName)
		})

		When("the app name in the manifest is missing", func() {
			BeforeEach(func() {
				helpers.WriteManifest(manifestPath, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"instances": 3,
						},
					},
				})
			})

			It("reports an error", func() {
				session := helpers.CF("apply-manifest", "-f", manifestPath)
				Eventually(session.Err).Should(Say("For application at index 0: Name must not be empty"))
				Eventually(session).Should(Say("FAILED"))
				Eventually(session).Should(Exit(1))
			})
		})

		When("there is a CC error", func() {
			BeforeEach(func() {
				helpers.WriteManifest(manifestPath, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"name":      appName,
							"instances": -1,
						},
					},
				})
			})

			It("displays the error", func() {
				session := helpers.CF("apply-manifest", "-f", manifestPath)
				Eventually(session.Err).Should(Say("Instances must be greater than or equal to 0"))
				Eventually(session).Should(Say("FAILED"))

				Eventually(session).Should(Exit(1))
			})
		})

		When("-f is provided", func() {
			When("the -f flag is not given an arg", func() {
				It("tells the user that the flag requires an arg, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f")

					Eventually(session.Err).Should(Say("Incorrect Usage: expected argument for flag `-f'"))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the -f flag points to a directory that does not have a manifest.yml file", func() {
				var (
					emptyDir string
				)

				BeforeEach(func() {
					emptyDir = helpers.TempDirAbsolutePath("", "empty")
				})

				AfterEach(func() {
					Expect(os.RemoveAll(emptyDir)).ToNot(HaveOccurred())
				})

				It("tells the user that the provided path doesn't exist, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f", emptyDir)

					Eventually(session.Err).Should(helpers.SayPath("Incorrect Usage: The specified directory '%s' does not contain a file named 'manifest.yml'.", emptyDir))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the -f flag points to a file that does not exist", func() {
				It("tells the user that the provided path doesn't exist, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f", "path/that/does/not/exist")

					Eventually(session.Err).Should(Say("Incorrect Usage: The specified path 'path/that/does/not/exist' does not exist."))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the manifest exists where -f points", func() {
				It("applies the manifest successfully", func() {
					userName, _ := helpers.GetCredentials()
					helpers.WriteManifest(filepath.Join(appDir, "manifest.yml"), map[string]interface{}{
						"applications": []map[string]interface{}{
							{
								"name":      appName,
								"instances": 3,
							},
						},
					})

					session := helpers.CF("apply-manifest", "-f", appDir)
					Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
					Eventually(session).Should(Exit())

					session = helpers.CF("app", appName)
					Eventually(session).Should(Say(`instances:\s+%s`, `\d/3`))
					Eventually(session).Should(Exit())
				})
			})

		})

		When("-f is not provided", func() {
			When("a properly formatted manifest is present in the pwd", func() {
				It("autodetects and applies the manifest", func() {
					userName, _ := helpers.GetCredentials()
					helpers.WriteManifest(filepath.Join(appDir, "manifest.yml"), map[string]interface{}{
						"applications": []map[string]interface{}{
							{
								"name":      appName,
								"instances": 3,
							},
						},
					})

					session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: appDir}, "apply-manifest")
					formatString := fmt.Sprintf("Applying manifest %%s in org %s / space %s as %s...", orgName, spaceName, userName)
					Eventually(session).Should(helpers.SayPath(formatString, manifestPath))
					Eventually(session).Should(Exit())

					session = helpers.CF("app", appName)
					Eventually(session).Should(Say(`instances:\s+%s`, `\d/3`))
					Eventually(session).Should(Exit())
				})
			})

			When("the current directory does not have a manifest", func() {
				It("fails nicely", func() {
					currentDir, err := os.Getwd()
					Expect(err).NotTo(HaveOccurred())
					session := helpers.CF("apply-manifest")

					Eventually(session.Err).Should(helpers.SayPath(`Could not find 'manifest.yml' file in %s`, currentDir))
					Eventually(session).Should(Say("FAILED"))
					Eventually(session).Should(Exit(1))
				})
			})
		})

		When("testing manifest diffing output", func() {
			var (
				userName string
			)

			BeforeEach(func() {
				userName, _ = helpers.GetCredentials()
			})

			When("there are no changes in the manifest", func() {
				It("shows no changes", func() {
					helpers.WithHelloWorldApp(func(dir string) {
						manifest, manifestPath := pushAppAndGenerateManifest(appName, dir)
						helpers.WriteManifest(filepath.Join(dir, "manifest.yml"), manifest)
						session := helpers.CF("apply-manifest", "-f", manifestPath)
						Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
						Eventually(session).Should(Say("Updating with these attributes..."))
						Consistently(session).ShouldNot(Say(`^\+ `))
						Consistently(session).ShouldNot(Say(`^- `))
						Eventually(session).Should(Exit(0))
					})
				})
			})

			When("there are changes in the manifest", func() {
				It("shows changes", func() {
					helpers.WithHelloWorldApp(func(dir string) {
						manifest, manifestPath := pushAppAndGenerateManifest(appName, dir)
						helpers.WriteManifest(filepath.Join(dir, "manifest.yml"), manifest)

						session := helpers.CF("scale", appName, "-i", "3")
						Eventually(session).Should(Exit(0))

						session = helpers.CF("apply-manifest", "-f", manifestPath)
						Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
						Eventually(session).Should(Say("Updating with these attributes..."))
						Eventually(session).Should(Say("applications"))
						Eventually(session).Should(Say(`\n-\s+instances: 3`))
						Eventually(session).Should(Say(`\n\+\s+instances: 1`))
						Eventually(session).Should(Exit(0))
					})
				})
			})
		})

		When("--vars are provided", func() {
			var (
				tempDir        string
				pathToManifest string
			)

			BeforeEach(func() {
				var err error
				tempDir, err = ioutil.TempDir("", "simple-manifest-test")
				Expect(err).ToNot(HaveOccurred())
				pathToManifest = filepath.Join(tempDir, "manifest.yml")
				helpers.WriteManifest(pathToManifest, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"name": appName,
							"env": map[string]interface{}{
								"key1": "((var1))",
								"key4": "((var2))",
							},
						},
					},
				})
			})

			It("uses the manifest with substituted variables", func() {
				helpers.WithHelloWorldApp(func(dir string) {
					session := helpers.CF("apply-manifest", "-f", pathToManifest, "--var=var1=secret-key", "--var=var2=foobar")
					Eventually(session).Should(Exit(0))
				})

				session := helpers.CF("env", appName)
				Eventually(session).Should(Say(`key1:\s+secret-key`))
				Eventually(session).Should(Say(`key4:\s+foobar`))
				Eventually(session).Should(Exit(0))
			})
		})
	})
})

func pushAppAndGenerateManifest(appName, dir string) (map[string]interface{}, string) {
	session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: dir}, PushCommandName, appName)
	Eventually(session).Should(Exit(0))
	manifestPath := path.Join(dir, "manifest.yml")
	session = helpers.CF("create-app-manifest", appName, "-p", manifestPath)
	Eventually(session).Should(Say(fmt.Sprintf("Manifest file created successfully at %s", manifestPath)))
	Eventually(session).Should(Exit(0))
	manifest := helpers.ReadManifest(manifestPath)

	return manifest, manifestPath
}
<commit_msg>filepath.Join is necessary for windows separators<commit_after>package isolated

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"

	. "code.cloudfoundry.org/cli/cf/util/testhelpers/matchers"
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("apply-manifest command", func() {
	var (
		orgName      string
		spaceName    string
		appName      string
		manifestPath string
		appDir       string
	)

	BeforeEach(func() {
		orgName = helpers.NewOrgName()
		spaceName = helpers.NewSpaceName()
		appName = helpers.PrefixedRandomName("app")

		appDir = helpers.TempDirAbsolutePath("", "simple-app")

		manifestPath = filepath.Join(appDir, "manifest.yml")
		// Ensure the file exists at the minimum
		helpers.WriteManifest(manifestPath, map[string]interface{}{})
	})

	AfterEach(func() {
		Expect(os.RemoveAll(appDir)).ToNot(HaveOccurred())
	})

	Describe("help", func() {
		It("appears in cf help -a", func() {
			session := helpers.CF("help", "-a")
			Eventually(session).Should(Exit(0))
			Expect(session).To(HaveCommandInCategoryWithDescription("apply-manifest", "SPACES", "Apply manifest properties to a space"))
		})

		When("--help flag is set", func() {
			It("displays command usage to output", func() {
				session := helpers.CF("apply-manifest", "--help")

				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Say("apply-manifest - Apply manifest properties to a space"))
				Eventually(session).Should(Say("USAGE:"))
				Eventually(session).Should(Say("cf apply-manifest -f APP_MANIFEST_PATH"))
				Eventually(session).Should(Say("SEE ALSO:"))
				Eventually(session).Should(Say("create-app, create-app-manifest, push"))

				Eventually(session).Should(Exit(0))
			})
		})
	})

	When("the environment is not setup correctly", func() {
		It("fails with the appropriate errors", func() {
			helpers.CheckEnvironmentTargetedCorrectly(true, true, ReadOnlyOrg, "apply-manifest", "-f", manifestPath)
		})
	})

	When("the environment is set up correctly", func() {
		BeforeEach(func() {
			helpers.SetupCF(orgName, spaceName)
		})

		AfterEach(func() {
			helpers.QuickDeleteOrg(orgName)
		})

		When("the app name in the manifest is missing", func() {
			BeforeEach(func() {
				helpers.WriteManifest(manifestPath, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"instances": 3,
						},
					},
				})
			})

			It("reports an error", func() {
				session := helpers.CF("apply-manifest", "-f", manifestPath)
				Eventually(session.Err).Should(Say("For application at index 0: Name must not be empty"))
				Eventually(session).Should(Say("FAILED"))
				Eventually(session).Should(Exit(1))
			})
		})

		When("there is a CC error", func() {
			BeforeEach(func() {
				helpers.WriteManifest(manifestPath, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"name":      appName,
							"instances": -1,
						},
					},
				})
			})

			It("displays the error", func() {
				session := helpers.CF("apply-manifest", "-f", manifestPath)
				Eventually(session.Err).Should(Say("Instances must be greater than or equal to 0"))
				Eventually(session).Should(Say("FAILED"))

				Eventually(session).Should(Exit(1))
			})
		})

		When("-f is provided", func() {
			When("the -f flag is not given an arg", func() {
				It("tells the user that the flag requires an arg, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f")

					Eventually(session.Err).Should(Say("Incorrect Usage: expected argument for flag `-f'"))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the -f flag points to a directory that does not have a manifest.yml file", func() {
				var (
					emptyDir string
				)

				BeforeEach(func() {
					emptyDir = helpers.TempDirAbsolutePath("", "empty")
				})

				AfterEach(func() {
					Expect(os.RemoveAll(emptyDir)).ToNot(HaveOccurred())
				})

				It("tells the user that the provided path doesn't exist, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f", emptyDir)

					Eventually(session.Err).Should(helpers.SayPath("Incorrect Usage: The specified directory '%s' does not contain a file named 'manifest.yml'.", emptyDir))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the -f flag points to a file that does not exist", func() {
				It("tells the user that the provided path doesn't exist, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f", "path/that/does/not/exist")

					Eventually(session.Err).Should(Say("Incorrect Usage: The specified path 'path/that/does/not/exist' does not exist."))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the manifest exists where -f points", func() {
				It("applies the manifest successfully", func() {
					userName, _ := helpers.GetCredentials()
					helpers.WriteManifest(filepath.Join(appDir, "manifest.yml"), map[string]interface{}{
						"applications": []map[string]interface{}{
							{
								"name":      appName,
								"instances": 3,
							},
						},
					})

					session := helpers.CF("apply-manifest", "-f", appDir)
					Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
					Eventually(session).Should(Exit())

					session = helpers.CF("app", appName)
					Eventually(session).Should(Say(`instances:\s+%s`, `\d/3`))
					Eventually(session).Should(Exit())
				})
			})

		})

		When("-f is not provided", func() {
			When("a properly formatted manifest is present in the pwd", func() {
				It("autodetects and applies the manifest", func() {
					userName, _ := helpers.GetCredentials()
					helpers.WriteManifest(filepath.Join(appDir, "manifest.yml"), map[string]interface{}{
						"applications": []map[string]interface{}{
							{
								"name":      appName,
								"instances": 3,
							},
						},
					})

					session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: appDir}, "apply-manifest")
					formatString := fmt.Sprintf("Applying manifest %%s in org %s / space %s as %s...", orgName, spaceName, userName)
					Eventually(session).Should(helpers.SayPath(formatString, manifestPath))
					Eventually(session).Should(Exit())

					session = helpers.CF("app", appName)
					Eventually(session).Should(Say(`instances:\s+%s`, `\d/3`))
					Eventually(session).Should(Exit())
				})
			})

			When("the current directory does not have a manifest", func() {
				It("fails nicely", func() {
					currentDir, err := os.Getwd()
					Expect(err).NotTo(HaveOccurred())
					session := helpers.CF("apply-manifest")

					Eventually(session.Err).Should(helpers.SayPath(`Could not find 'manifest.yml' file in %s`, currentDir))
					Eventually(session).Should(Say("FAILED"))
					Eventually(session).Should(Exit(1))
				})
			})
		})

		When("testing manifest diffing output", func() {
			var (
				userName string
			)

			BeforeEach(func() {
				userName, _ = helpers.GetCredentials()
			})

			When("there are no changes in the manifest", func() {
				It("shows no changes", func() {
					helpers.WithHelloWorldApp(func(dir string) {
						manifest, manifestPath := pushAppAndGenerateManifest(appName, dir)
						helpers.WriteManifest(filepath.Join(dir, "manifest.yml"), manifest)
						session := helpers.CF("apply-manifest", "-f", manifestPath)
						Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
						Eventually(session).Should(Say("Updating with these attributes..."))
						Consistently(session).ShouldNot(Say(`^\+ `))
						Consistently(session).ShouldNot(Say(`^- `))
						Eventually(session).Should(Exit(0))
					})
				})
			})

			When("there are changes in the manifest", func() {
				It("shows changes", func() {
					helpers.WithHelloWorldApp(func(dir string) {
						manifest, manifestPath := pushAppAndGenerateManifest(appName, dir)
						helpers.WriteManifest(filepath.Join(dir, "manifest.yml"), manifest)

						session := helpers.CF("scale", appName, "-i", "3")
						Eventually(session).Should(Exit(0))

						session = helpers.CF("apply-manifest", "-f", manifestPath)
						Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
						Eventually(session).Should(Say("Updating with these attributes..."))
						Eventually(session).Should(Say("applications"))
						Eventually(session).Should(Say(`\n-\s+instances: 3`))
						Eventually(session).Should(Say(`\n\+\s+instances: 1`))
						Eventually(session).Should(Exit(0))
					})
				})
			})
		})

		When("--vars are provided", func() {
			var (
				tempDir        string
				pathToManifest string
			)

			BeforeEach(func() {
				var err error
				tempDir, err = ioutil.TempDir("", "simple-manifest-test")
				Expect(err).ToNot(HaveOccurred())
				pathToManifest = filepath.Join(tempDir, "manifest.yml")
				helpers.WriteManifest(pathToManifest, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"name": appName,
							"env": map[string]interface{}{
								"key1": "((var1))",
								"key4": "((var2))",
							},
						},
					},
				})
			})

			It("uses the manifest with substituted variables", func() {
				helpers.WithHelloWorldApp(func(dir string) {
					session := helpers.CF("apply-manifest", "-f", pathToManifest, "--var=var1=secret-key", "--var=var2=foobar")
					Eventually(session).Should(Exit(0))
				})

				session := helpers.CF("env", appName)
				Eventually(session).Should(Say(`key1:\s+secret-key`))
				Eventually(session).Should(Say(`key4:\s+foobar`))
				Eventually(session).Should(Exit(0))
			})
		})
	})
})

func pushAppAndGenerateManifest(appName, dir string) (map[string]interface{}, string) {
	session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: dir}, PushCommandName, appName)
	Eventually(session).Should(Exit(0))
	manifestPath := filepath.Join(dir, "manifest.yml")
	session = helpers.CF("create-app-manifest", appName, "-p", manifestPath)
	Eventually(session).Should(Say(fmt.Sprintf("Manifest file created successfully at %s", manifestPath)))
	Eventually(session).Should(Exit(0))
	manifest := helpers.ReadManifest(manifestPath)

	return manifest, manifestPath
}
<|endoftext|>


<commit_before>package juju

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/state"
	"net/url"
	"os"
)

// NewService creates a new service with the given name to run the given
// charm.  If svcName is empty, the charm name will be used.
func (conn *Conn) AddService(name string, ch *state.Charm) (*state.Service, error) {
	st, err := conn.State()
	if err != nil {
		return nil, err
	}
	if name == "" {
		name = ch.URL().Name // TODO sch.Meta().Name ?
	}
	svc, err := st.AddService(name, ch)
	if err != nil {
		return nil, err
	}
	meta := ch.Meta()
	for rname, rel := range meta.Peers {
		ep := state.RelationEndpoint{
			name,
			rel.Interface,
			rname,
			state.RolePeer,
			state.RelationScope(rel.Scope),
		}
		if err := st.AddRelation(ep); err != nil {
			return nil, fmt.Errorf("cannot add peer relation %q to service %q: %v", rname, name, err)
		}
	}
	return svc, nil
}

// PutCharm uploads the given charm to provider storage, and adds a
// state.Charm to the state.  The charm is not uploaded if a charm with
// the same URL already exists in the state.
// If bumpRevision is true, the charm must be a local directory,
// and the revision number will be incremented before pushing.
// Local charms will be interpreted relative to the repoPath directory.
func (conn *Conn) PutCharm(curl *charm.URL, repoPath string, bumpRevision bool) (*state.Charm, error) {
	repo, err := charm.InferRepository(curl, repoPath)
	if err != nil {
		return nil, fmt.Errorf("cannot infer charm repository: %v", err)
	}
	if curl.Revision == -1 {
		rev, err := repo.Latest(curl)
		if err != nil {
			return nil, fmt.Errorf("cannot get latest charm revision: %v", err)
		}
		curl = curl.WithRevision(rev)
	}
	ch, err := repo.Get(curl)
	if err != nil {
		return nil, fmt.Errorf("cannot get charm: %v", err)
	}
	if bumpRevision {
		chd, ok := ch.(*charm.Dir)
		if !ok {
			return nil, fmt.Errorf("cannot increment version of charm %q: not a directory", curl)
		}
		if err = chd.SetDiskRevision(chd.Revision() + 1); err != nil {
			return nil, fmt.Errorf("cannot increment version of charm %q: %v", curl, err)
		}
		curl = curl.WithRevision(chd.Revision())
	}
	st, err := conn.State()
	if err != nil {
		return nil, err
	}
	if sch, err := st.Charm(curl); err == nil {
		return sch, nil
	}
	var buf bytes.Buffer
	switch ch := ch.(type) {
	case *charm.Dir:
		if err := ch.BundleTo(&buf); err != nil {
			return nil, fmt.Errorf("cannot bundle charm: %v", err)
		}
	case *charm.Bundle:
		f, err := os.Open(ch.Path)
		if err != nil {
			return nil, fmt.Errorf("cannot open charm bundle path: %v", err)
		}
		defer f.Close()
		if _, err := io.Copy(&buf, f); err != nil {
			return nil, fmt.Errorf("cannot read charm from bundle: %v", err)
		}
	default:
		return nil, fmt.Errorf("unknown charm type %T", ch)
	}
	h := sha256.New()
	h.Write(buf.Bytes())
	digest := hex.EncodeToString(h.Sum(nil))
	storage := conn.Environ.Storage()
	name := charm.Quote(curl.String())
	if err := storage.Put(name, &buf, int64(len(buf.Bytes()))); err != nil {
		return nil, fmt.Errorf("cannot put charm: %v", err)
	}
	ustr, err := storage.URL(name)
	if err != nil {
		return nil, fmt.Errorf("cannot get storage URL for charm: %v", err)
	}
	u, err := url.Parse(ustr)
	if err != nil {
		return nil, fmt.Errorf("cannot parse storage URL: %v", err)
	}
	sch, err := st.AddCharm(ch, curl, u, digest)
	if err != nil {
		return nil, fmt.Errorf("cannot add charm: %v", err)
	}
	return sch, nil
}

// AddUnits starts n units of the given service and allocates machines
// to them as necessary.
func (conn *Conn) AddUnits(svc *state.Service, n int) ([]*state.Unit, error) {
	st, err := conn.State()
	if err != nil {
		return nil, err
	}
	units := make([]*state.Unit, n)
	// TODO what do we do if we fail half-way through this process?
	for i := 0; i < n; i++ {
		policy := conn.Environ.AssignmentPolicy()
		unit, err := svc.AddUnit()
		if err != nil {
			return nil, fmt.Errorf("cannot add unit %d/%d to service %q: %v", i+1, n, svc.Name(), err)
		}
		if err := st.AssignUnit(unit, policy); err != nil {
			return nil, fmt.Errorf("cannot assign machine to unit %s of service %q: %v", unit.Name(), svc.Name(), err)
		}
		units[i] = unit
	}
	return units, nil
}
<commit_msg>juju: fix comment<commit_after>package juju

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/state"
	"net/url"
	"os"
)

// AddService creates a new service with the given name to run the given
// charm.  If svcName is empty, the charm name will be used.
func (conn *Conn) AddService(name string, ch *state.Charm) (*state.Service, error) {
	st, err := conn.State()
	if err != nil {
		return nil, err
	}
	if name == "" {
		name = ch.URL().Name // TODO sch.Meta().Name ?
	}
	svc, err := st.AddService(name, ch)
	if err != nil {
		return nil, err
	}
	meta := ch.Meta()
	for rname, rel := range meta.Peers {
		ep := state.RelationEndpoint{
			name,
			rel.Interface,
			rname,
			state.RolePeer,
			state.RelationScope(rel.Scope),
		}
		if err := st.AddRelation(ep); err != nil {
			return nil, fmt.Errorf("cannot add peer relation %q to service %q: %v", rname, name, err)
		}
	}
	return svc, nil
}

// PutCharm uploads the given charm to provider storage, and adds a
// state.Charm to the state.  The charm is not uploaded if a charm with
// the same URL already exists in the state.
// If bumpRevision is true, the charm must be a local directory,
// and the revision number will be incremented before pushing.
// Local charms will be interpreted relative to the repoPath directory.
func (conn *Conn) PutCharm(curl *charm.URL, repoPath string, bumpRevision bool) (*state.Charm, error) {
	repo, err := charm.InferRepository(curl, repoPath)
	if err != nil {
		return nil, fmt.Errorf("cannot infer charm repository: %v", err)
	}
	if curl.Revision == -1 {
		rev, err := repo.Latest(curl)
		if err != nil {
			return nil, fmt.Errorf("cannot get latest charm revision: %v", err)
		}
		curl = curl.WithRevision(rev)
	}
	ch, err := repo.Get(curl)
	if err != nil {
		return nil, fmt.Errorf("cannot get charm: %v", err)
	}
	if bumpRevision {
		chd, ok := ch.(*charm.Dir)
		if !ok {
			return nil, fmt.Errorf("cannot increment version of charm %q: not a directory", curl)
		}
		if err = chd.SetDiskRevision(chd.Revision() + 1); err != nil {
			return nil, fmt.Errorf("cannot increment version of charm %q: %v", curl, err)
		}
		curl = curl.WithRevision(chd.Revision())
	}
	st, err := conn.State()
	if err != nil {
		return nil, err
	}
	if sch, err := st.Charm(curl); err == nil {
		return sch, nil
	}
	var buf bytes.Buffer
	switch ch := ch.(type) {
	case *charm.Dir:
		if err := ch.BundleTo(&buf); err != nil {
			return nil, fmt.Errorf("cannot bundle charm: %v", err)
		}
	case *charm.Bundle:
		f, err := os.Open(ch.Path)
		if err != nil {
			return nil, fmt.Errorf("cannot open charm bundle path: %v", err)
		}
		defer f.Close()
		if _, err := io.Copy(&buf, f); err != nil {
			return nil, fmt.Errorf("cannot read charm from bundle: %v", err)
		}
	default:
		return nil, fmt.Errorf("unknown charm type %T", ch)
	}
	h := sha256.New()
	h.Write(buf.Bytes())
	digest := hex.EncodeToString(h.Sum(nil))
	storage := conn.Environ.Storage()
	name := charm.Quote(curl.String())
	if err := storage.Put(name, &buf, int64(len(buf.Bytes()))); err != nil {
		return nil, fmt.Errorf("cannot put charm: %v", err)
	}
	ustr, err := storage.URL(name)
	if err != nil {
		return nil, fmt.Errorf("cannot get storage URL for charm: %v", err)
	}
	u, err := url.Parse(ustr)
	if err != nil {
		return nil, fmt.Errorf("cannot parse storage URL: %v", err)
	}
	sch, err := st.AddCharm(ch, curl, u, digest)
	if err != nil {
		return nil, fmt.Errorf("cannot add charm: %v", err)
	}
	return sch, nil
}

// AddUnits starts n units of the given service and allocates machines
// to them as necessary.
func (conn *Conn) AddUnits(svc *state.Service, n int) ([]*state.Unit, error) {
	st, err := conn.State()
	if err != nil {
		return nil, err
	}
	units := make([]*state.Unit, n)
	// TODO what do we do if we fail half-way through this process?
	for i := 0; i < n; i++ {
		policy := conn.Environ.AssignmentPolicy()
		unit, err := svc.AddUnit()
		if err != nil {
			return nil, fmt.Errorf("cannot add unit %d/%d to service %q: %v", i+1, n, svc.Name(), err)
		}
		if err := st.AssignUnit(unit, policy); err != nil {
			return nil, fmt.Errorf("cannot assign machine to unit %s of service %q: %v", unit.Name(), svc.Name(), err)
		}
		units[i] = unit
	}
	return units, nil
}
<|endoftext|>


<commit_before>package kafka

import (
	"fmt"
	"github.com/Shopify/sarama"
	log "github.com/Sirupsen/logrus"
	"github.com/krallistic/kafka-operator/spec"
	kazoo "github.com/krallistic/kazoo-go"
	"github.com/krallistic/kafka-operator/util"
)

var (
	logger = log.WithFields(log.Fields{
		"package": "kafka",
	})
)

type KafkaUtil struct {
	KafkaClient sarama.Client
	BrokerList  []string
	ClusterName string
	KazooClient *kazoo.Kazoo
}

func New(clusterSpec spec.KafkaCluster) (*KafkaUtil, error) {
	brokerList := util.GetBrokerAdressess(clusterSpec)

	methodLogger := log.WithFields(log.Fields{
		"method":      "new",
		"clusterName": clusterSpec.Metadata.Name,
		"brokers":     brokerList,
	})
	config := sarama.NewConfig()

	methodLogger.Info("Creating KafkaUtil")

	client, err := sarama.NewClient(brokerList, config)
	if err != nil {
		methodLogger.WithFields(log.Fields{
			"error": err,
		}).Error("Error creating sarama kafka Client")
		return nil, err
	}

	kz, err := kazoo.NewKazooFromConnectionString(clusterSpec.Spec.ZookeeperConnect, nil)
	if err != nil {
		methodLogger.WithFields(log.Fields{
			"error": err,
			"zookeeperConnect": clusterSpec.Spec.ZookeeperConnect,
		}).Error("Cant create kazoo client")
		return nil, err
	}

	k := &KafkaUtil{
		KafkaClient: client,
		ClusterName: clusterSpec.Metadata.Name,
		BrokerList:  brokerList,
		KazooClient: kz,
	}

	methodLogger.Info("Initilized Kafka CLient, KazooClient and created KafkaUtil")
	k.ListTopics()
	return k, nil
}

func (k *KafkaUtil) ListTopics() ([]string, error) {
	fmt.Println("Listing KafkaTopics")
	topics, err := k.KafkaClient.Topics()
	if err != nil {
		return nil, err
	}

	for _, t := range topics {
		fmt.Println("Current topic:", t)
	}
	return topics, nil
}

func (k *KafkaUtil) GetPartitions(topic string) ([]int32, error) {
	partitions, err := k.KafkaClient.Partitions(topic)
	if err != nil {
		return nil, err
	}
	return partitions, nil
}

func (k *KafkaUtil) PrintFullStats() error {
	topics, err := k.ListTopics()
	if err != nil {
		return err
	}
	for _, topic := range topics {
		partitions, err := k.GetPartitions(topic)
		if err != nil {
			return err
		}
		fmt.Println("Topic: %s, Partitions %s", topic, partitions)
	}

	return nil
}

func (k *KafkaUtil) RemoveTopicsFromBrokers(cluster spec.KafkaCluster, brokerToDelete int32) (spec.KafkaReassignmentConfig, error) {
	methodLogger := log.WithFields(log.Fields{
		"method":      "GenerateReassign",
		"clusterName": cluster.Metadata.Name,
	})
	topics, err := k.KafkaClient.Topics()
	if err != nil {
		methodLogger.Error("Error Listing Topics")
		return spec.KafkaReassignmentConfig{}, err
	}

	//TODO it should be possible to Delete multiple Brokers
	brokersToDelete := []int32{brokerToDelete}
	for _, topic := range topics {
		k.KazooClient.RemoveTopicFromBrokers(topic, brokersToDelete)
		partitions, err := k.KafkaClient.Partitions(topic)
		if err != nil {
			methodLogger.Error("Error Listing Partitions")
			return spec.KafkaReassignmentConfig{}, err
		}
		for _, partition := range partitions {
			partition, err := k.KafkaClient.Replicas(topic, partition)
			if err != nil {
				methodLogger.Error("Error listing partitions")
				return spec.KafkaReassignmentConfig{}, err
			}
			fmt.Println(partition)
		}
	}

	return spec.KafkaReassignmentConfig{}, nil
}

func (k *KafkaUtil) AllTopicsInSync() bool {
	//TODO error checking
	topics, _ := k.KazooClient.Topics()
	for _, topic := range topics {
		partitions, _ := topic.Partitions()
		for _, partition := range partitions {
			underReplicated, _ := partition.UnderReplicated()
			if underReplicated{
				return false
			}
		}
	}
	return true
}

func (k *KafkaUtil) CreateTopic(topicSpec spec.KafkaTopicSpec) error {
	fmt.Println("Creating Kafka Topics: ", topicSpec)
	broker, _ := k.KafkaClient.Coordinator("operatorConsumerGroup")
	request := sarama.MetadataRequest{Topics: []string{topicSpec.Name}}
	metadataPartial, err := broker.GetMetadata(&request)
	if err != nil {
		return err
	}

	replicas := []int32{0}
	isr := []int32{0}

	metadataResponse := &sarama.MetadataResponse{}
	metadataResponse.AddBroker(broker.Addr(), broker.ID())

	metadataPartial.AddTopic(topicSpec.Name, sarama.ErrNoError)
	//TODO dynamic partitions
	metadataPartial.AddTopicPartition(topicSpec.Name, 0, broker.ID(), replicas, isr, sarama.ErrNoError)
	metadataPartial.AddTopicPartition(topicSpec.Name, 1, broker.ID(), replicas, isr, sarama.ErrNoError)

	return nil
}
<commit_msg>Error Handling<commit_after>package kafka

import (
	"fmt"
	"github.com/Shopify/sarama"
	log "github.com/Sirupsen/logrus"
	"github.com/krallistic/kafka-operator/spec"
	kazoo "github.com/krallistic/kazoo-go"
	"github.com/krallistic/kafka-operator/util"
)

var (
	logger = log.WithFields(log.Fields{
		"package": "kafka",
	})
)

type KafkaUtil struct {
	KafkaClient sarama.Client
	BrokerList  []string
	ClusterName string
	KazooClient *kazoo.Kazoo
}

func New(clusterSpec spec.KafkaCluster) (*KafkaUtil, error) {
	brokerList := util.GetBrokerAdressess(clusterSpec)

	methodLogger := log.WithFields(log.Fields{
		"method":      "new",
		"clusterName": clusterSpec.Metadata.Name,
		"brokers":     brokerList,
	})
	config := sarama.NewConfig()

	methodLogger.Info("Creating KafkaUtil")

	client, err := sarama.NewClient(brokerList, config)
	if err != nil {
		methodLogger.WithFields(log.Fields{
			"error": err,
		}).Error("Error creating sarama kafka Client")
		return nil, err
	}

	kz, err := kazoo.NewKazooFromConnectionString(clusterSpec.Spec.ZookeeperConnect, nil)
	if err != nil {
		methodLogger.WithFields(log.Fields{
			"error": err,
			"zookeeperConnect": clusterSpec.Spec.ZookeeperConnect,
		}).Error("Cant create kazoo client")
		return nil, err
	}

	k := &KafkaUtil{
		KafkaClient: client,
		ClusterName: clusterSpec.Metadata.Name,
		BrokerList:  brokerList,
		KazooClient: kz,
	}

	methodLogger.Info("Initilized Kafka CLient, KazooClient and created KafkaUtil")
	k.ListTopics()
	return k, nil
}

func (k *KafkaUtil) ListTopics() ([]string, error) {
	fmt.Println("Listing KafkaTopics")
	topics, err := k.KafkaClient.Topics()
	if err != nil {
		return nil, err
	}

	for _, t := range topics {
		fmt.Println("Current topic:", t)
	}
	return topics, nil
}

func (k *KafkaUtil) GetPartitions(topic string) ([]int32, error) {
	partitions, err := k.KafkaClient.Partitions(topic)
	if err != nil {
		return nil, err
	}
	return partitions, nil
}

func (k *KafkaUtil) PrintFullStats() error {
	topics, err := k.ListTopics()
	if err != nil {
		return err
	}
	for _, topic := range topics {
		partitions, err := k.GetPartitions(topic)
		if err != nil {
			return err
		}
		fmt.Println("Topic: %s, Partitions %s", topic, partitions)
	}

	return nil
}

func (k *KafkaUtil) RemoveTopicsFromBrokers(cluster spec.KafkaCluster, brokerToDelete int32) (spec.KafkaReassignmentConfig, error) {
	methodLogger := log.WithFields(log.Fields{
		"method":      "GenerateReassign",
		"clusterName": cluster.Metadata.Name,
	})
	topics, err := k.KafkaClient.Topics()
	if err != nil {
		methodLogger.Error("Error Listing Topics")
		return spec.KafkaReassignmentConfig{}, err
	}

	//TODO it should be possible to Delete multiple Brokers
	brokersToDelete := []int32{brokerToDelete}
	for _, topic := range topics {
		k.KazooClient.RemoveTopicFromBrokers(topic, brokersToDelete)
		partitions, err := k.KafkaClient.Partitions(topic)
		if err != nil {
			methodLogger.Error("Error Listing Partitions")
			return spec.KafkaReassignmentConfig{}, err
		}
		for _, partition := range partitions {
			partition, err := k.KafkaClient.Replicas(topic, partition)
			if err != nil {
				methodLogger.Error("Error listing partitions")
				return spec.KafkaReassignmentConfig{}, err
			}
			fmt.Println(partition)
		}
	}

	return spec.KafkaReassignmentConfig{}, nil
}

func (k *KafkaUtil) AllTopicsInSync() (bool, error) {
	//TODO error checking
	topics, err := k.KazooClient.Topics()
	if err != nil {
		return nil, err
	}
	for _, topic := range topics {
		partitions, err := topic.Partitions()
		if err != nil {
			return nil, err
		}
		for _, partition := range partitions {
			underReplicated, err := partition.UnderReplicated()
			if err != nil {
				return nil, err
			}
			if underReplicated{
				return false, nil
			}
		}
	}
	return true, nil
}

func (k *KafkaUtil) CreateTopic(topicSpec spec.KafkaTopicSpec) error {
	fmt.Println("Creating Kafka Topics: ", topicSpec)
	broker, _ := k.KafkaClient.Coordinator("operatorConsumerGroup")
	request := sarama.MetadataRequest{Topics: []string{topicSpec.Name}}
	metadataPartial, err := broker.GetMetadata(&request)
	if err != nil {
		return err
	}

	replicas := []int32{0}
	isr := []int32{0}

	metadataResponse := &sarama.MetadataResponse{}
	metadataResponse.AddBroker(broker.Addr(), broker.ID())

	metadataPartial.AddTopic(topicSpec.Name, sarama.ErrNoError)
	//TODO dynamic partitions
	metadataPartial.AddTopicPartition(topicSpec.Name, 0, broker.ID(), replicas, isr, sarama.ErrNoError)
	metadataPartial.AddTopicPartition(topicSpec.Name, 1, broker.ID(), replicas, isr, sarama.ErrNoError)

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"

	"github.com/beatgammit/stein"
	"github.com/wsxiaoys/terminal"
)

const (
	timeFormat = "2006-01-02 at 15:04:05"
)

func newKegParser() *stein.Parser {
	return &stein.Parser{
		&kegHandler{
			currentLevel: -1,
		},
	}
}

type kegHandler struct {
	*stein.DefaultHandler
	currentLevel int
	afterCase    bool
}

func (h *kegHandler) getIndentation(level int) string {
	indentation := ""
	for i := 0; i < level; i++ {
		indentation += "    "
	}
	return indentation
}

func (h *kegHandler) HandleSuite(s *stein.Suite) {
	fmt.Printf("Test suite started on %s\n", s.Start.Local().Format(timeFormat))
	if s.Count > 0 {
		fmt.Printf("Running %d tests...\n\n", s.Count)
	}
}

func (h *kegHandler) HandleCase(c *stein.Case) {
	clearLine()
	fmt.Printf("%s%s\n", h.getIndentation(h.currentLevel), c.Label)
	h.currentLevel = c.Level
}

func (h *kegHandler) HandleTest(t *stein.Test) {
	status := t.Status
	if strings.ToLower(status) == "omit" {
		status = "skip"
	}
	clearLine()
	testLine := fmt.Sprintf("%s%s ... %s", h.getIndentation(h.currentLevel+1), status, t.Label)
	printColor(statusColors[status], testLine)
	if *onlyFail {
		if status == "fail" || status == "error" {
			fmt.Print("\n")
		} else {
			fmt.Print("\r")
		}
	} else {
		fmt.Print("\n")
	}
}

func (h *kegHandler) HandleFinal(t *stein.Tally) {
	clearLine()
	fmt.Printf("\nResults: %d pass, %d fail, %d error, %d skip\n",
		t.Counts.Pass, t.Counts.Fail, t.Counts.Error, t.Counts.Omit)
}

func clearLine() {
	if *onlyFail {
		terminal.Stdout.ClearLine()
	}
}

var statusColors = map[string]string{
	"pass":  "g",
	"skip":  "y",
	"fail":  "r",
	"error": "m",
	"todo":  "b",
}

func printColor(color, text string) {
	terminal.Stdout.Color(color).Print(text).Reset()
}
<commit_msg>Show file name and line number for fail/error<commit_after>package main

import (
	"fmt"
	"strings"

	"github.com/beatgammit/stein"
	"github.com/wsxiaoys/terminal"
)

const (
	timeFormat = "2006-01-02 at 15:04:05"
)

func newKegParser() *stein.Parser {
	return &stein.Parser{
		&kegHandler{
			currentLevel: -1,
		},
	}
}

type kegHandler struct {
	*stein.DefaultHandler
	currentLevel int
	afterCase    bool
}

func (h *kegHandler) getIndentation(level int) string {
	indentation := ""
	for i := 0; i < level; i++ {
		indentation += "    "
	}
	return indentation
}

func (h *kegHandler) HandleSuite(s *stein.Suite) {
	fmt.Printf("Test suite started on %s\n", s.Start.Local().Format(timeFormat))
	if s.Count > 0 {
		fmt.Printf("Running %d tests...\n\n", s.Count)
	}
}

func (h *kegHandler) HandleCase(c *stein.Case) {
	clearLine()
	fmt.Printf("%s%s\n", h.getIndentation(h.currentLevel), c.Label)
	h.currentLevel = c.Level
}

func (h *kegHandler) HandleTest(t *stein.Test) {
	status := t.Status
	if strings.ToLower(status) == "omit" {
		status = "skip"
	}
	clearLine()
    indent := h.getIndentation(h.currentLevel+1)
    var testLine string
    var postIndent = ""
    for len(status) + len(postIndent) < longestStatus {
        postIndent += " "
    }
    testLine = fmt.Sprintf("%s%s%s ... %s", indent, status, postIndent, t.Label)
	printColor(statusColors[status], testLine)
    if t.Exception != nil && t.Exception.Line > 0 && t.Exception.File != "" {
        for i := 0; i < longestStatus + len(" ... "); i++ {
            indent += " "
        }
        printColor(statusColors["note"], fmt.Sprintf("\n%s%s:%d", indent, t.Exception.File, t.Exception.Line))
    }
	if *onlyFail {
		if status == "fail" || status == "error" {
			fmt.Print("\n")
		} else {
			fmt.Print("\r")
		}
	} else {
		fmt.Print("\n")
	}
}

func (h *kegHandler) HandleFinal(t *stein.Tally) {
	clearLine()
	fmt.Printf("\nResults: %d pass, %d fail, %d error, %d skip\n",
		t.Counts.Pass, t.Counts.Fail, t.Counts.Error, t.Counts.Omit)
}

func clearLine() {
	if *onlyFail {
		terminal.Stdout.ClearLine()
	}
}

var statusColors = map[string]string{
	"pass":  "g",
	"skip":  "y",
	"fail":  "r",
	"error": "m",
	"todo":  "b",
    "note": "w",
}

var longestStatus int

func init() {
    for status := range statusColors {
        if len(status) > longestStatus {
            longestStatus = len(status)
        }
    }
}

func printColor(color, text string) {
	terminal.Stdout.Color(color).Print(text).Reset()
}
<|endoftext|>


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

package api

import (
	"fmt"
	"os"
	"syscall"

	dockerclient "github.com/control-center/serviced/commons/docker"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/utils"
)

// TODO: what to do about logging?

// GetServiceInstances returns all instances running on a service
func (a *api) GetServiceInstances(serviceID string) ([]service.Instance, error) {
	client, err := a.connectMaster()
	if err != nil {
		return nil, err
	}
	return client.GetServiceInstances(serviceID)
}

// StopServiceInstance stops a running instance of a service.
func (a *api) StopServiceInstance(serviceID string, instanceID int) error {
	client, err := a.connectMaster()
	if err != nil {
		return err
	}
	return client.StopServiceInstance(serviceID, instanceID)
}

// AttachServiceInstance locates and attaches to a running instance of a service
func (a *api) AttachServiceInstance(serviceID string, instanceID int, command string, args []string) error {
	var (
		targetHost      string
		targetIP        string
		targetContainer string
	)

	hostID, err := utils.HostID()
	if err != nil {
		return err
	}

	// Check to see if serviceID is actually a dockerID
	_, err = dockerclient.FindContainer(serviceID)
	if err == nil {
		targetHost = hostID
		targetContainer = serviceID
	} else {
		client, err := a.connectMaster()
		if err != nil {
			return err
		}

		// get the location of the running instance
		location, err := client.LocateServiceInstance(serviceID, instanceID)
		if err != nil {
			return err
		}

		targetHost = location.HostID
		targetIP = location.HostIP
		targetContainer = location.ContainerID
	}

	// attach to the container
	cmd := []string{}
	if targetHost != hostID {
		cmd := []string{
			"/usr/bin/ssh",
			"-t", targetIP, "--",
			"serviced", "--endpoint", GetOptionsRPCEndpoint(),
			"service", "attach", fmt.Sprintf("%s", targetContainer),
		}
		cmd = append(cmd, command)
		cmd = append(cmd, args...)
		return syscall.Exec(cmd[0], cmd[0:], os.Environ())
	} else {
		if command == "" {
			cmd = append(cmd, "/bin/bash")
		} else {
			cmd = append(cmd, command)
			cmd = append(cmd, args...)
		}
		return utils.AttachAndExec(targetContainer, cmd)
	}
}

// LogsForServiceInstance returns the logs for the service instance
func (a *api) LogsForServiceInstance(serviceID string, instanceID int, command string, args []string) error {
	client, err := a.connectMaster()
	if err != nil {
		return err
	}

	// get the location of the running instance
	location, err := client.LocateServiceInstance(serviceID, instanceID)
	if err != nil {
		return err
	}

	// check to see if it is running on this host
	hostID, err := utils.HostID()
	if err != nil {
		return err
	}

	// report container logs
	cmd := []string{}
	if location.HostID != hostID {
		cmd := []string{
			"/usr/bin/ssh",
			"-t", location.HostIP, "--",
			"serviced", "--endpoint", GetOptionsRPCEndpoint(),
			"service", "logs", fmt.Sprintf("%s/%d", serviceID, instanceID),
		}
		if command != "" {
			cmd = append(cmd, command)
			cmd = append(cmd, args...)
		}
		return syscall.Exec(cmd[0], cmd[0:], os.Environ())
	} else {
		if command != "" {
			cmd = append(cmd, command)
			cmd = append(cmd, args...)
		}
		return dockerclient.Logs(location.ContainerID, cmd)
	}
}

// SendDockerAction submits an action to a running service instance
func (a *api) SendDockerAction(serviceID string, instanceID int, action string, args []string) error {
	client, err := a.connectMaster()
	if err != nil {
		return err
	}

	return client.SendDockerAction(serviceID, instanceID, action, args)
}
<commit_msg>call docker exec directly when attaching to a remote container<commit_after>// Copyright 2016 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package api

import (
	"fmt"
	"os"
	"syscall"

	dockerclient "github.com/control-center/serviced/commons/docker"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/utils"
)

// TODO: what to do about logging?

// GetServiceInstances returns all instances running on a service
func (a *api) GetServiceInstances(serviceID string) ([]service.Instance, error) {
	client, err := a.connectMaster()
	if err != nil {
		return nil, err
	}
	return client.GetServiceInstances(serviceID)
}

// StopServiceInstance stops a running instance of a service.
func (a *api) StopServiceInstance(serviceID string, instanceID int) error {
	client, err := a.connectMaster()
	if err != nil {
		return err
	}
	return client.StopServiceInstance(serviceID, instanceID)
}

// AttachServiceInstance locates and attaches to a running instance of a service
func (a *api) AttachServiceInstance(serviceID string, instanceID int, command string, args []string) error {
	var (
		targetHost      string
		targetIP        string
		targetContainer string
	)

	hostID, err := utils.HostID()
	if err != nil {
		return err
	}

	// Check to see if serviceID is actually a dockerID
	_, err = dockerclient.FindContainer(serviceID)
	if err == nil {
		targetHost = hostID
		targetContainer = serviceID
	} else {
		client, err := a.connectMaster()
		if err != nil {
			return err
		}

		// get the location of the running instance
		location, err := client.LocateServiceInstance(serviceID, instanceID)
		if err != nil {
			return err
		}

		targetHost = location.HostID
		targetIP = location.HostIP
		targetContainer = location.ContainerID
	}

	// attach to the container
	cmd := []string{}
	if targetHost != hostID {
		cmd := []string{
			"/usr/bin/ssh",
			"-t", targetIP, "--",
			"docker", "exec", "-it", fmt.Sprintf("%s", targetContainer),
		}
		cmd = append(cmd, command)
		cmd = append(cmd, args...)
		return syscall.Exec(cmd[0], cmd[0:], os.Environ())
	} else {
		if command == "" {
			cmd = append(cmd, "/bin/bash")
		} else {
			cmd = append(cmd, command)
			cmd = append(cmd, args...)
		}
		return utils.AttachAndExec(targetContainer, cmd)
	}
}

// LogsForServiceInstance returns the logs for the service instance
func (a *api) LogsForServiceInstance(serviceID string, instanceID int, command string, args []string) error {
	client, err := a.connectMaster()
	if err != nil {
		return err
	}

	// get the location of the running instance
	location, err := client.LocateServiceInstance(serviceID, instanceID)
	if err != nil {
		return err
	}

	// check to see if it is running on this host
	hostID, err := utils.HostID()
	if err != nil {
		return err
	}

	// report container logs
	cmd := []string{}
	if location.HostID != hostID {
		cmd := []string{
			"/usr/bin/ssh",
			"-t", location.HostIP, "--",
			"serviced", "--endpoint", GetOptionsRPCEndpoint(),
			"service", "logs", fmt.Sprintf("%s/%d", serviceID, instanceID),
		}
		if command != "" {
			cmd = append(cmd, command)
			cmd = append(cmd, args...)
		}
		return syscall.Exec(cmd[0], cmd[0:], os.Environ())
	} else {
		if command != "" {
			cmd = append(cmd, command)
			cmd = append(cmd, args...)
		}
		return dockerclient.Logs(location.ContainerID, cmd)
	}
}

// SendDockerAction submits an action to a running service instance
func (a *api) SendDockerAction(serviceID string, instanceID int, action string, args []string) error {
	client, err := a.connectMaster()
	if err != nil {
		return err
	}

	return client.SendDockerAction(serviceID, instanceID, action, args)
}
<|endoftext|>


<commit_before>package command

import (
	"context"
	"encoding/gob"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"syscall"
	"unsafe"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/hashicorp/yamux"
	"github.com/kr/pty"
	"github.com/spolu/wrp"
	"github.com/spolu/wrp/cli"
	"github.com/spolu/wrp/lib/errors"
	"github.com/spolu/wrp/lib/out"
	"github.com/spolu/wrp/lib/token"
)

const (
	// CmdNmOpen is the command name.
	CmdNmOpen cli.CmdName = "open"
)

func init() {
	cli.Registrar[CmdNmOpen] = NewOpen
}

// Open spawns a new shared terminal.
type Open struct {
	shell string

	address string
	session string
	user    wrp.User

	username string

	cmd *exec.Cmd
	pty *os.File

	dataC   net.Conn
	stateC  net.Conn
	stateR  *gob.Decoder
	updateC net.Conn
	updateW *gob.Encoder
	hostC   net.Conn
	hostW   *gob.Encoder
}

// NewOpen constructs and initializes the command.
func NewOpen() cli.Command {
	return &Open{}
}

// Name returns the command name.
func (c *Open) Name() cli.CmdName {
	return CmdNmOpen
}

// Help prints out the help message for the command.
func (c *Open) Help(
	ctx context.Context,
) {
	out.Normf("\nUsage: ")
	out.Boldf("wrp open [<id>]\n")
	out.Normf("\n")
	out.Normf("  Spawns a shared terminal with the provided id. Others can use the id to connect.\n")
	out.Normf("  If no id is provided a random one is generated.\n")
	out.Normf("\n")
	out.Normf("Arguments:\n")
	out.Boldf("  id\n")
	out.Normf("    The id to assign to the newly shared terminal.\n")
	out.Valuf("    spolu-dev\n")
	out.Normf("\n")
	out.Normf("Examples:\n")
	out.Valuf("  wrp open\n")
	out.Valuf("  wrp open spolu-dev\n")
	out.Normf("\n")
}

// Parse parses the arguments passed to the command.
func (c *Open) Parse(
	ctx context.Context,
	args []string,
) error {
	if len(args) == 0 {
		c.session = token.RandStr()
	} else {
		c.session = args[0]
	}

	c.address = wrp.DefaultAddress

	c.shell = "/bin/bash"
	if os.Getenv("SHELL") != "" {
		c.shell = os.Getenv("SHELL")
	}

	user, err := user.Current()
	if err != nil {
		return errors.Trace(
			errors.Newf("Error retrieving current user: %v", err),
		)
	}
	c.username = user.Username

	c.user = wrp.User{
		Token:   token.New("guest"),
		Secret:  "",
		Session: token.New("session"),
	}

	return nil
}

// Execute the command or return a human-friendly error.
func (c *Open) Execute(
	ctx context.Context,
) error {
	ctx, cancel := context.WithCancel(ctx)

	out.Normf("\n")
	out.Normf("wrp id: ")
	out.Boldf("%s\n", c.session)
	out.Normf("\n")

	conn, err := net.Dial("tcp", c.address)
	if err != nil {
		return errors.Trace(
			errors.Newf("Connection error: %v", err),
		)
	}

	mux, err := yamux.Client(conn, nil)
	if err != nil {
		return errors.Trace(
			errors.Newf("Session error: %v", err),
		)
	}
	// Closes stateC, updateC, hostC, dataC, mux and conn.
	defer mux.Close()

	// Setup pty
	c.cmd = exec.Command(c.shell)
	c.pty, err = pty.Start(c.cmd)
	if err != nil {
		return errors.Trace(
			errors.Newf("PTY error: %v", err),
		)
	}
	go func() {
		if err := c.cmd.Wait(); err != nil {
			out.Errof("[Error] Cmd wait error: %v\n", err)
		}
		cancel()
	}()
	// Closes the newly created pty.
	defer c.pty.Close()

	// Setup local term.
	stdin := int(os.Stdin.Fd())
	if !terminal.IsTerminal(stdin) {
		return errors.Trace(
			errors.Newf("Not running in a terminal."),
		)
	}
	old, err := terminal.MakeRaw(stdin)
	if err != nil {
		return errors.Trace(
			errors.Newf("Unable to make terminal raw: %v", err),
		)
	}
	// Restores the terminal once we're done.
	defer terminal.Restore(stdin, old)

	// Opens state channel stateC.
	c.stateC, err = mux.Open()
	if err != nil {
		return errors.Trace(
			errors.Newf("State channel open error: %v", err),
		)
	}
	c.stateR = gob.NewDecoder(c.stateC)

	// Open update channel updateC.
	c.updateC, err = mux.Open()
	if err != nil {
		return errors.Trace(
			errors.Newf("Update channel open error: %v", err),
		)
	}
	c.updateW = gob.NewEncoder(c.updateC)

	// Send initial client update.
	if err := c.updateW.Encode(wrp.ClientUpdate{
		Session:  c.session,
		From:     c.user,
		Hosting:  true,
		Username: c.username,
	}); err != nil {
		return errors.Trace(
			errors.Newf("Send client update error: %v", err),
		)
	}

	// Open data channel dataC.
	c.dataC, err = mux.Open()
	if err != nil {
		return errors.Trace(
			errors.Newf("Data channel open error: %v", err),
		)
	}

	// Open host channel hostC.
	c.hostC, err = mux.Open()
	if err != nil {
		return errors.Trace(
			errors.Newf("Host channel open error: %v", err),
		)
	}
	c.hostW = gob.NewEncoder(c.hostC)

	// Send initial host update.
	cols, rows, err := terminal.GetSize(stdin)
	if err != nil {
		return errors.Trace(
			errors.Newf("Getsize error: %v", err),
		)
	}

	if err := c.updateW.Encode(wrp.HostUpdate{
		Session:    c.session,
		From:       c.user,
		WindowSize: wrp.Size{Rows: rows, Cols: cols},
	}); err != nil {
		return errors.Trace(
			errors.Newf("Send host update error: %v", err),
		)
	}

	// Main loops.

	// Forward window resizes to pty and stateChannel
	go func() {
		ch := make(chan os.Signal, 1)
		signal.Notify(ch, syscall.SIGWINCH)
		for {
			cols, rows, err := terminal.GetSize(stdin)
			if err != nil {
				out.Errof("[Error] Getsize error: %v\n", err)
				break
			}
			if err := Setsize(c.pty, rows, cols); err != nil {
				out.Errof("[Error] Setsize error: %v\n", err)
				break
			}
			if err := syscall.Kill(c.cmd.Process.Pid, syscall.SIGWINCH); err != nil {
				out.Errof("[Error] Sigwinch error: %v\n", err)
				break
			}

			if err := c.hostW.Encode(wrp.HostUpdate{
				Session:    c.session,
				From:       c.user,
				WindowSize: wrp.Size{Rows: rows, Cols: cols},
			}); err != nil {
				out.Errof("[Error] Send host update error: %v\n", err)
				break
			}
			<-ch
		}
		cancel()
	}()

	// Multiplex shell to dataC, Stdout
	go func() {
		cli.Multiplex(ctx, []io.Writer{c.dataC, os.Stdout}, c.pty)
		cancel()
	}()

	// Multiplex dataC to pty
	go func() {
		cli.Multiplex(ctx, []io.Writer{c.pty}, c.dataC)
		cancel()
	}()

	// Multiplex Stdin to pty
	go func() {
		cli.Multiplex(ctx, []io.Writer{c.pty}, os.Stdin)
		cancel()
	}()

	<-ctx.Done()

	return nil
}

type winsize struct {
	ws_row    uint16
	ws_col    uint16
	ws_xpixel uint16
	ws_ypixel uint16
}

func Setsize(f *os.File, rows, cols int) error {
	ws := winsize{ws_row: uint16(rows), ws_col: uint16(cols)}
	_, _, errno := syscall.Syscall(
		syscall.SYS_IOCTL,
		f.Fd(),
		syscall.TIOCSWINSZ,
		uintptr(unsafe.Pointer(&ws)),
	)
	if errno != 0 {
		return syscall.Errno(errno)
	}
	return nil
}
<commit_msg>PS1 support for bash<commit_after>package command

import (
	"context"
	"encoding/gob"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"syscall"
	"unsafe"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/hashicorp/yamux"
	"github.com/kr/pty"
	"github.com/spolu/wrp"
	"github.com/spolu/wrp/cli"
	"github.com/spolu/wrp/lib/errors"
	"github.com/spolu/wrp/lib/out"
	"github.com/spolu/wrp/lib/token"
)

const (
	// CmdNmOpen is the command name.
	CmdNmOpen cli.CmdName = "open"
)

func init() {
	cli.Registrar[CmdNmOpen] = NewOpen
}

// Open spawns a new shared terminal.
type Open struct {
	shell string

	address string
	session string
	user    wrp.User

	username string

	cmd *exec.Cmd
	pty *os.File

	dataC   net.Conn
	stateC  net.Conn
	stateR  *gob.Decoder
	updateC net.Conn
	updateW *gob.Encoder
	hostC   net.Conn
	hostW   *gob.Encoder
}

// NewOpen constructs and initializes the command.
func NewOpen() cli.Command {
	return &Open{}
}

// Name returns the command name.
func (c *Open) Name() cli.CmdName {
	return CmdNmOpen
}

// Help prints out the help message for the command.
func (c *Open) Help(
	ctx context.Context,
) {
	out.Normf("\nUsage: ")
	out.Boldf("wrp open [<id>]\n")
	out.Normf("\n")
	out.Normf("  Spawns a shared terminal with the provided id. Others can use the id to connect.\n")
	out.Normf("  If no id is provided a random one is generated.\n")
	out.Normf("\n")
	out.Normf("Arguments:\n")
	out.Boldf("  id\n")
	out.Normf("    The id to assign to the newly shared terminal.\n")
	out.Valuf("    spolu-dev\n")
	out.Normf("\n")
	out.Normf("Examples:\n")
	out.Valuf("  wrp open\n")
	out.Valuf("  wrp open spolu-dev\n")
	out.Normf("\n")
}

// Parse parses the arguments passed to the command.
func (c *Open) Parse(
	ctx context.Context,
	args []string,
) error {
	if len(args) == 0 {
		c.session = token.RandStr()
	} else {
		c.session = args[0]
	}

	c.address = wrp.DefaultAddress

	c.shell = "/bin/bash"
	if os.Getenv("SHELL") != "" {
		c.shell = os.Getenv("SHELL")
	}

	user, err := user.Current()
	if err != nil {
		return errors.Trace(
			errors.Newf("Error retrieving current user: %v", err),
		)
	}
	c.username = user.Username

	// Sets the BASH prompt
	os.Setenv(
		"PS1",
		fmt.Sprintf(
			//"\\[\033[01;31m\\][wrp:%s]\\[\033[00m\\] \\[\\033[01;32m\\]\\h\\[\033[00m\\]:\\[\033[01;34m\\]\\W\\[\033[00m\\]\\$ ",
			"\\[\033[01;31m\\][wrp:%s]\\[\033[00m\\] \\[\033[01;34m\\]\\W\\[\033[00m\\]\\$ ",
			c.session,
		),
	)

	c.user = wrp.User{
		Token:   token.New("guest"),
		Secret:  "",
		Session: token.New("session"),
	}

	return nil
}

// Execute the command or return a human-friendly error.
func (c *Open) Execute(
	ctx context.Context,
) error {
	ctx, cancel := context.WithCancel(ctx)

	out.Normf("\n")
	out.Normf("wrp id: ")
	out.Boldf("%s\n", c.session)
	out.Normf("\n")

	conn, err := net.Dial("tcp", c.address)
	if err != nil {
		return errors.Trace(
			errors.Newf("Connection error: %v", err),
		)
	}

	mux, err := yamux.Client(conn, nil)
	if err != nil {
		return errors.Trace(
			errors.Newf("Session error: %v", err),
		)
	}
	// Closes stateC, updateC, hostC, dataC, mux and conn.
	defer mux.Close()

	// Setup pty
	c.cmd = exec.Command(c.shell)
	c.pty, err = pty.Start(c.cmd)
	if err != nil {
		return errors.Trace(
			errors.Newf("PTY error: %v", err),
		)
	}
	go func() {
		if err := c.cmd.Wait(); err != nil {
			out.Errof("[Error] Cmd wait error: %v\n", err)
		}
		cancel()
	}()
	// Closes the newly created pty.
	defer c.pty.Close()

	// Setup local term.
	stdin := int(os.Stdin.Fd())
	if !terminal.IsTerminal(stdin) {
		return errors.Trace(
			errors.Newf("Not running in a terminal."),
		)
	}
	old, err := terminal.MakeRaw(stdin)
	if err != nil {
		return errors.Trace(
			errors.Newf("Unable to make terminal raw: %v", err),
		)
	}
	// Restores the terminal once we're done.
	defer terminal.Restore(stdin, old)

	// Opens state channel stateC.
	c.stateC, err = mux.Open()
	if err != nil {
		return errors.Trace(
			errors.Newf("State channel open error: %v", err),
		)
	}
	c.stateR = gob.NewDecoder(c.stateC)

	// Open update channel updateC.
	c.updateC, err = mux.Open()
	if err != nil {
		return errors.Trace(
			errors.Newf("Update channel open error: %v", err),
		)
	}
	c.updateW = gob.NewEncoder(c.updateC)

	// Send initial client update.
	if err := c.updateW.Encode(wrp.ClientUpdate{
		Session:  c.session,
		From:     c.user,
		Hosting:  true,
		Username: c.username,
	}); err != nil {
		return errors.Trace(
			errors.Newf("Send client update error: %v", err),
		)
	}

	// Open data channel dataC.
	c.dataC, err = mux.Open()
	if err != nil {
		return errors.Trace(
			errors.Newf("Data channel open error: %v", err),
		)
	}

	// Open host channel hostC.
	c.hostC, err = mux.Open()
	if err != nil {
		return errors.Trace(
			errors.Newf("Host channel open error: %v", err),
		)
	}
	c.hostW = gob.NewEncoder(c.hostC)

	// Send initial host update.
	cols, rows, err := terminal.GetSize(stdin)
	if err != nil {
		return errors.Trace(
			errors.Newf("Getsize error: %v", err),
		)
	}

	if err := c.updateW.Encode(wrp.HostUpdate{
		Session:    c.session,
		From:       c.user,
		WindowSize: wrp.Size{Rows: rows, Cols: cols},
	}); err != nil {
		return errors.Trace(
			errors.Newf("Send host update error: %v", err),
		)
	}

	// Main loops.

	// Forward window resizes to pty and stateChannel
	go func() {
		ch := make(chan os.Signal, 1)
		signal.Notify(ch, syscall.SIGWINCH)
		for {
			cols, rows, err := terminal.GetSize(stdin)
			if err != nil {
				out.Errof("[Error] Getsize error: %v\n", err)
				break
			}
			if err := Setsize(c.pty, rows, cols); err != nil {
				out.Errof("[Error] Setsize error: %v\n", err)
				break
			}
			if err := syscall.Kill(c.cmd.Process.Pid, syscall.SIGWINCH); err != nil {
				out.Errof("[Error] Sigwinch error: %v\n", err)
				break
			}

			if err := c.hostW.Encode(wrp.HostUpdate{
				Session:    c.session,
				From:       c.user,
				WindowSize: wrp.Size{Rows: rows, Cols: cols},
			}); err != nil {
				out.Errof("[Error] Send host update error: %v\n", err)
				break
			}
			<-ch
		}
		cancel()
	}()

	// Multiplex shell to dataC, Stdout
	go func() {
		cli.Multiplex(ctx, []io.Writer{c.dataC, os.Stdout}, c.pty)
		cancel()
	}()

	// Multiplex dataC to pty
	go func() {
		cli.Multiplex(ctx, []io.Writer{c.pty}, c.dataC)
		cancel()
	}()

	// Multiplex Stdin to pty
	go func() {
		cli.Multiplex(ctx, []io.Writer{c.pty}, os.Stdin)
		cancel()
	}()

	<-ctx.Done()

	return nil
}

type winsize struct {
	ws_row    uint16
	ws_col    uint16
	ws_xpixel uint16
	ws_ypixel uint16
}

func Setsize(f *os.File, rows, cols int) error {
	ws := winsize{ws_row: uint16(rows), ws_col: uint16(cols)}
	_, _, errno := syscall.Syscall(
		syscall.SYS_IOCTL,
		f.Fd(),
		syscall.TIOCSWINSZ,
		uintptr(unsafe.Pointer(&ws)),
	)
	if errno != 0 {
		return syscall.Errno(errno)
	}
	return nil
}
<|endoftext|>


<commit_before>package autonat

import (
	"context"
	"errors"
	"math/rand"
	"sync"
	"time"

	host "github.com/libp2p/go-libp2p-host"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

type NATStatus int

const (
	NATStatusUnknown NATStatus = iota
	NATStatusPublic
	NATStatusPrivate
)

type AutoNAT interface {
	Status() NATStatus
	PublicAddr() (ma.Multiaddr, error)
}

type AutoNATState struct {
	ctx    context.Context
	host   host.Host
	peers  map[peer.ID]struct{}
	status NATStatus
	addr   ma.Multiaddr
	mx     sync.Mutex
}

func NewAutoNAT(ctx context.Context, h host.Host) AutoNAT {
	as := &AutoNATState{
		ctx:    ctx,
		host:   h,
		peers:  make(map[peer.ID]struct{}),
		status: NATStatusUnknown,
	}

	h.Network().Notify(as)
	go as.background()

	return as
}

func (as *AutoNATState) Status() NATStatus {
	return as.status
}

func (as *AutoNATState) PublicAddr() (ma.Multiaddr, error) {
	as.mx.Lock()
	defer as.mx.Unlock()

	if as.status != NATStatusPublic {
		return nil, errors.New("NAT Status is not public")
	}

	return as.addr, nil
}

func (as *AutoNATState) background() {
	// wait a bit for the node to come online and establish some connections
	// before starting autodetection
	time.Sleep(10 * time.Second)
	for {
		as.autodetect()
		select {
		case <-time.After(15 * time.Minute):
		case <-as.ctx.Done():
			return
		}
	}
}

func (as *AutoNATState) autodetect() {
	if len(as.peers) == 0 {
		log.Debugf("skipping NAT auto detection; no autonat peers")
		return
	}

	as.mx.Lock()
	peers := make([]peer.ID, 0, len(as.peers))
	for p := range as.peers {
		if len(as.host.Network().ConnsToPeer(p)) > 0 {
			peers = append(peers, p)
		}
	}

	if len(peers) == 0 {
		// we don't have any open connections, try any autonat peer that we know about
		for p := range as.peers {
			peers = append(peers, p)
		}
	}

	as.mx.Unlock()

	shufflePeers(peers)

	for _, p := range peers {
		cli := NewAutoNATClient(as.host, p)
		ctx, cancel := context.WithTimeout(as.ctx, 60*time.Second)
		a, err := cli.Dial(ctx)
		cancel()

		switch {
		case err == nil:
			log.Debugf("NAT status is public; address through %s: %s", p.Pretty(), a.String())
			as.mx.Lock()
			as.addr = a
			as.status = NATStatusPublic
			as.mx.Unlock()
			return

		case IsDialError(err):
			log.Debugf("NAT status is private; dial error through %s: %s", p.Pretty(), err.Error())
			as.mx.Lock()
			as.status = NATStatusPrivate
			as.mx.Unlock()
			return

		default:
			log.Debugf("Error dialing through %s: %s", p.Pretty(), err.Error())
		}
	}

	as.mx.Lock()
	as.status = NATStatusUnknown
	as.mx.Unlock()
}

func shufflePeers(peers []peer.ID) {
	for i := range peers {
		j := rand.Intn(i + 1)
		peers[i], peers[j] = peers[j], peers[i]
	}
}
<commit_msg>bump initial autodiscovery delay to 15s<commit_after>package autonat

import (
	"context"
	"errors"
	"math/rand"
	"sync"
	"time"

	host "github.com/libp2p/go-libp2p-host"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

type NATStatus int

const (
	NATStatusUnknown NATStatus = iota
	NATStatusPublic
	NATStatusPrivate
)

type AutoNAT interface {
	Status() NATStatus
	PublicAddr() (ma.Multiaddr, error)
}

type AutoNATState struct {
	ctx    context.Context
	host   host.Host
	peers  map[peer.ID]struct{}
	status NATStatus
	addr   ma.Multiaddr
	mx     sync.Mutex
}

func NewAutoNAT(ctx context.Context, h host.Host) AutoNAT {
	as := &AutoNATState{
		ctx:    ctx,
		host:   h,
		peers:  make(map[peer.ID]struct{}),
		status: NATStatusUnknown,
	}

	h.Network().Notify(as)
	go as.background()

	return as
}

func (as *AutoNATState) Status() NATStatus {
	return as.status
}

func (as *AutoNATState) PublicAddr() (ma.Multiaddr, error) {
	as.mx.Lock()
	defer as.mx.Unlock()

	if as.status != NATStatusPublic {
		return nil, errors.New("NAT Status is not public")
	}

	return as.addr, nil
}

func (as *AutoNATState) background() {
	// wait a bit for the node to come online and establish some connections
	// before starting autodetection
	time.Sleep(15 * time.Second)
	for {
		as.autodetect()
		select {
		case <-time.After(15 * time.Minute):
		case <-as.ctx.Done():
			return
		}
	}
}

func (as *AutoNATState) autodetect() {
	if len(as.peers) == 0 {
		log.Debugf("skipping NAT auto detection; no autonat peers")
		return
	}

	as.mx.Lock()
	peers := make([]peer.ID, 0, len(as.peers))
	for p := range as.peers {
		if len(as.host.Network().ConnsToPeer(p)) > 0 {
			peers = append(peers, p)
		}
	}

	if len(peers) == 0 {
		// we don't have any open connections, try any autonat peer that we know about
		for p := range as.peers {
			peers = append(peers, p)
		}
	}

	as.mx.Unlock()

	shufflePeers(peers)

	for _, p := range peers {
		cli := NewAutoNATClient(as.host, p)
		ctx, cancel := context.WithTimeout(as.ctx, 60*time.Second)
		a, err := cli.Dial(ctx)
		cancel()

		switch {
		case err == nil:
			log.Debugf("NAT status is public; address through %s: %s", p.Pretty(), a.String())
			as.mx.Lock()
			as.addr = a
			as.status = NATStatusPublic
			as.mx.Unlock()
			return

		case IsDialError(err):
			log.Debugf("NAT status is private; dial error through %s: %s", p.Pretty(), err.Error())
			as.mx.Lock()
			as.status = NATStatusPrivate
			as.mx.Unlock()
			return

		default:
			log.Debugf("Error dialing through %s: %s", p.Pretty(), err.Error())
		}
	}

	as.mx.Lock()
	as.status = NATStatusUnknown
	as.mx.Unlock()
}

func shufflePeers(peers []peer.ID) {
	for i := range peers {
		j := rand.Intn(i + 1)
		peers[i], peers[j] = peers[j], peers[i]
	}
}
<|endoftext|>


<commit_before>package relay

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

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	_ "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]struct{}
	status autonat.NATStatus
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]struct{}),
		disconnect: make(chan struct{}, 1),
		status:     autonat.NATStatusUnknown,
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	return ar.relayAddrs(ar.addrsF(addrs))
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			ar.mx.Lock()
			ar.status = autonat.NATStatusUnknown
			ar.mx.Unlock()
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			ar.mx.Lock()
			if ar.status != autonat.NATStatusPublic {
				push = true
			}
			ar.status = autonat.NATStatusPublic
			ar.mx.Unlock()

		case autonat.NATStatusPrivate:
			update := ar.findRelays(ctx)
			ar.mx.Lock()
			if update || ar.status != autonat.NATStatusPrivate {
				push = true
			}
			ar.status = autonat.NATStatusPrivate
			ar.mx.Unlock()
		}

		if push {
			push = false
			ar.host.PushIdentify()
		}

		select {
		case <-ar.disconnect:
			push = true
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) bool {
	if ar.numRelays() >= DesiredRelays {
		return false
	}

	update := false
	for retry := 0; retry < 5; retry++ {
		if retry > 0 {
			log.Debug("no relays connected; retrying in 30s")
			select {
			case <-time.After(30 * time.Second):
			case <-ctx.Done():
				return update
			}
		}

		update = ar.findRelaysOnce(ctx) || update
		if ar.numRelays() > 0 {
			return update
		}
	}
	return update
}

func (ar *AutoRelay) findRelaysOnce(ctx context.Context) bool {
	pis, err := ar.discoverRelays(ctx)
	if err != nil {
		log.Debugf("error discovering relays: %s", err)
		return false
	}
	log.Debugf("discovered %d relays", len(pis))
	pis = ar.selectRelays(ctx, pis)
	log.Debugf("selected %d relays", len(pis))

	update := false
	for _, pi := range pis {
		update = ar.tryRelay(ctx, pi) || update
	}
	return update
}

func (ar *AutoRelay) numRelays() int {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	return len(ar.relays)
}

// usingRelay returns if we're currently using the given relay.
func (ar *AutoRelay) usingRelay(p peer.ID) bool {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	_, ok := ar.relays[p]
	return ok
}

// addRelay adds the given relay to our set of relays.
// returns true when we add a new relay
func (ar *AutoRelay) tryRelay(ctx context.Context, pi pstore.PeerInfo) bool {
	if ar.usingRelay(pi.ID) {
		return false
	}

	if !ar.connect(ctx, pi) {
		return false
	}

	ar.mx.Lock()
	defer ar.mx.Unlock()

	// make sure we're still connected.
	if ar.host.Network().Connectedness(pi.ID) != inet.Connected {
		return false
	}
	ar.relays[pi.ID] = struct{}{}

	return true
}

func (ar *AutoRelay) connect(ctx context.Context, pi pstore.PeerInfo) bool {
	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	if len(pi.Addrs) == 0 {
		var err error
		pi, err = ar.router.FindPeer(ctx, pi.ID)
		if err != nil {
			log.Debugf("error finding relay peer %s: %s", pi.ID, err.Error())
			return false
		}
	}

	err := ar.host.Connect(ctx, pi)
	if err != nil {
		log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
		return false
	}

	// tag the connection as very important
	ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)
	return true
}

func (ar *AutoRelay) discoverRelays(ctx context.Context) ([]pstore.PeerInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	return discovery.FindPeers(ctx, ar.discover, RelayRendezvous, 1000)
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []pstore.PeerInfo) []pstore.PeerInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	shuffleRelays(pis)
	return pis
}

// This function is computes the NATed relay addrs when our status is private:
// - The public addrs are removed from the address set.
// - The non-public addrs are included verbatim so that peers behind the same NAT/firewall
//   can still dial us directly.
// - On top of those, we add the relay-specific addrs for the relays to which we are
//   connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
//   through which we can be dialed.
func (ar *AutoRelay) relayAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	if ar.status != autonat.NATStatusPrivate {
		ar.mx.Unlock()
		return addrs
	}

	relays := make([]peer.ID, 0, len(ar.relays))
	for p := range ar.relays {
		relays = append(relays, p)
	}
	ar.mx.Unlock()

	raddrs := make([]ma.Multiaddr, 0, 4*len(relays)+2)

	// only keep private addrs from the original addr set
	for _, addr := range addrs {
		if manet.IsPrivateAddr(addr) {
			raddrs = append(raddrs, addr)
		}
	}

	// add relay specific addrs to the list
	for _, p := range relays {
		addrs := cleanupAddressSet(ar.host.Peerstore().Addrs(p))

		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", p.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range addrs {
			pub := addr.Encapsulate(circuit)
			raddrs = append(raddrs, pub)
		}
	}

	return raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

// Notifee
func (ar *AutoRelay) Listen(inet.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(inet.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(inet.Network, inet.Conn)      {}

func (ar *AutoRelay) Disconnected(net inet.Network, c inet.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == inet.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(inet.Network, inet.Stream) {}
func (ar *AutoRelay) ClosedStream(inet.Network, inet.Stream) {}
<commit_msg>fix bug in findRelaysOnce: it connects to all relays<commit_after>package relay

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

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	_ "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]struct{}
	status autonat.NATStatus
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]struct{}),
		disconnect: make(chan struct{}, 1),
		status:     autonat.NATStatusUnknown,
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	return ar.relayAddrs(ar.addrsF(addrs))
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			ar.mx.Lock()
			ar.status = autonat.NATStatusUnknown
			ar.mx.Unlock()
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			ar.mx.Lock()
			if ar.status != autonat.NATStatusPublic {
				push = true
			}
			ar.status = autonat.NATStatusPublic
			ar.mx.Unlock()

		case autonat.NATStatusPrivate:
			update := ar.findRelays(ctx)
			ar.mx.Lock()
			if update || ar.status != autonat.NATStatusPrivate {
				push = true
			}
			ar.status = autonat.NATStatusPrivate
			ar.mx.Unlock()
		}

		if push {
			push = false
			ar.host.PushIdentify()
		}

		select {
		case <-ar.disconnect:
			push = true
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) bool {
	if ar.numRelays() >= DesiredRelays {
		return false
	}

	update := false
	for retry := 0; retry < 5; retry++ {
		if retry > 0 {
			log.Debug("no relays connected; retrying in 30s")
			select {
			case <-time.After(30 * time.Second):
			case <-ctx.Done():
				return update
			}
		}

		update = ar.findRelaysOnce(ctx) || update
		if ar.numRelays() > 0 {
			return update
		}
	}
	return update
}

func (ar *AutoRelay) findRelaysOnce(ctx context.Context) bool {
	pis, err := ar.discoverRelays(ctx)
	if err != nil {
		log.Debugf("error discovering relays: %s", err)
		return false
	}
	log.Debugf("discovered %d relays", len(pis))
	pis = ar.selectRelays(ctx, pis)
	log.Debugf("selected %d relays", len(pis))

	update := false
	for _, pi := range pis {
		update = ar.tryRelay(ctx, pi) || update
		if ar.numRelays() >= DesiredRelays {
			break
		}
	}
	return update
}

func (ar *AutoRelay) numRelays() int {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	return len(ar.relays)
}

// usingRelay returns if we're currently using the given relay.
func (ar *AutoRelay) usingRelay(p peer.ID) bool {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	_, ok := ar.relays[p]
	return ok
}

// addRelay adds the given relay to our set of relays.
// returns true when we add a new relay
func (ar *AutoRelay) tryRelay(ctx context.Context, pi pstore.PeerInfo) bool {
	if ar.usingRelay(pi.ID) {
		return false
	}

	if !ar.connect(ctx, pi) {
		return false
	}

	ar.mx.Lock()
	defer ar.mx.Unlock()

	// make sure we're still connected.
	if ar.host.Network().Connectedness(pi.ID) != inet.Connected {
		return false
	}
	ar.relays[pi.ID] = struct{}{}

	return true
}

func (ar *AutoRelay) connect(ctx context.Context, pi pstore.PeerInfo) bool {
	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	if len(pi.Addrs) == 0 {
		var err error
		pi, err = ar.router.FindPeer(ctx, pi.ID)
		if err != nil {
			log.Debugf("error finding relay peer %s: %s", pi.ID, err.Error())
			return false
		}
	}

	err := ar.host.Connect(ctx, pi)
	if err != nil {
		log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
		return false
	}

	// tag the connection as very important
	ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)
	return true
}

func (ar *AutoRelay) discoverRelays(ctx context.Context) ([]pstore.PeerInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	return discovery.FindPeers(ctx, ar.discover, RelayRendezvous, 1000)
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []pstore.PeerInfo) []pstore.PeerInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	shuffleRelays(pis)
	return pis
}

// This function is computes the NATed relay addrs when our status is private:
// - The public addrs are removed from the address set.
// - The non-public addrs are included verbatim so that peers behind the same NAT/firewall
//   can still dial us directly.
// - On top of those, we add the relay-specific addrs for the relays to which we are
//   connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
//   through which we can be dialed.
func (ar *AutoRelay) relayAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	if ar.status != autonat.NATStatusPrivate {
		ar.mx.Unlock()
		return addrs
	}

	relays := make([]peer.ID, 0, len(ar.relays))
	for p := range ar.relays {
		relays = append(relays, p)
	}
	ar.mx.Unlock()

	raddrs := make([]ma.Multiaddr, 0, 4*len(relays)+2)

	// only keep private addrs from the original addr set
	for _, addr := range addrs {
		if manet.IsPrivateAddr(addr) {
			raddrs = append(raddrs, addr)
		}
	}

	// add relay specific addrs to the list
	for _, p := range relays {
		addrs := cleanupAddressSet(ar.host.Peerstore().Addrs(p))

		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", p.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range addrs {
			pub := addr.Encapsulate(circuit)
			raddrs = append(raddrs, pub)
		}
	}

	return raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

// Notifee
func (ar *AutoRelay) Listen(inet.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(inet.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(inet.Network, inet.Conn)      {}

func (ar *AutoRelay) Disconnected(net inet.Network, c inet.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == inet.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(inet.Network, inet.Stream) {}
func (ar *AutoRelay) ClosedStream(inet.Network, inet.Stream) {}
<|endoftext|>


<commit_before>package dialect

import (
	"database/sql"
	"fmt"
	"time"
)

type postgres struct {
}

func (s *postgres) BinVar(i int) string {
	return fmt.Sprintf("$%v", i)
}

func (s *postgres) SupportLastInsertId() bool {
	return false
}

func (d *postgres) SqlTag(column interface{}, size int) string {
	switch column.(type) {
	case time.Time:
		return "timestamp with time zone"
	case bool, sql.NullBool:
		return "boolean"
	case int, int8, int16, int32, uint, uint8, uint16, uint32:
		return "integer"
	case int64, uint64, sql.NullInt64:
		return "bigint"
	case float32, float64, sql.NullFloat64:
		return "double precision"
	case []byte:
		return "bytea"
	case string, sql.NullString:
		if size > 0 && size < 65532 {
			return fmt.Sprintf("varchar(%d)", size)
		} else {
			return "text"
		}
	default:
		panic("Invalid sql type for postgres")
	}
}

func (s *postgres) PrimaryKeyTag(column interface{}, size int) string {
	switch column.(type) {
	case int, int8, int16, int32, uint, uint8, uint16, uint32:
		return "sehrial"
	case int64, uint64:
		return "bigserial"
	default:
		panic("Invalid primary key type")
	}
}

func (s *postgres) ReturningStr(key string) (str string) {
	return fmt.Sprintf("RETURNING \"%v\"", key)
}

func (s *postgres) Quote(key string) (str string) {
	return fmt.Sprintf("\"%s\"", key)
}
<commit_msg>Postgresql Dialet: fix primary tag for int types<commit_after>package dialect

import (
	"database/sql"
	"fmt"
	"time"
)

type postgres struct {
}

func (s *postgres) BinVar(i int) string {
	return fmt.Sprintf("$%v", i)
}

func (s *postgres) SupportLastInsertId() bool {
	return false
}

func (d *postgres) SqlTag(column interface{}, size int) string {
	switch column.(type) {
	case time.Time:
		return "timestamp with time zone"
	case bool, sql.NullBool:
		return "boolean"
	case int, int8, int16, int32, uint, uint8, uint16, uint32:
		return "integer"
	case int64, uint64, sql.NullInt64:
		return "bigint"
	case float32, float64, sql.NullFloat64:
		return "double precision"
	case []byte:
		return "bytea"
	case string, sql.NullString:
		if size > 0 && size < 65532 {
			return fmt.Sprintf("varchar(%d)", size)
		} else {
			return "text"
		}
	default:
		panic("Invalid sql type for postgres")
	}
}

func (s *postgres) PrimaryKeyTag(column interface{}, size int) string {
	switch column.(type) {
	case int, int8, int16, int32, uint, uint8, uint16, uint32:
		return "serial"
	case int64, uint64:
		return "bigserial"
	default:
		panic("Invalid primary key type")
	}
}

func (s *postgres) ReturningStr(key string) (str string) {
	return fmt.Sprintf("RETURNING \"%v\"", key)
}

func (s *postgres) Quote(key string) (str string) {
	return fmt.Sprintf("\"%s\"", key)
}
<|endoftext|>


<commit_before>// Copyright 2012 James Cooper. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// Package gorp provides a simple way to marshal Go structs to and from
// SQL databases.  It uses the database/sql package, and should work with any
// compliant database/sql driver.
//
// Source code and project home:
// https://github.com/go-gorp/gorp

package gorp

import (
	"fmt"
	"reflect"
	"strings"
)

type PostgresDialect struct {
	suffix string
}

func (d PostgresDialect) QuerySuffix() string { return ";" }

func (d PostgresDialect) ToSqlType(val reflect.Type, maxsize int, isAutoIncr bool) string {
	switch val.Kind() {
	case reflect.Ptr:
		return d.ToSqlType(val.Elem(), maxsize, isAutoIncr)
	case reflect.Bool:
		return "boolean"
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32:
		if isAutoIncr {
			return "serial"
		}
		return "integer"
	case reflect.Int64, reflect.Uint64:
		if isAutoIncr {
			return "bigserial"
		}
		return "bigint"
	case reflect.Float64:
		return "double precision"
	case reflect.Float32:
		return "real"
	case reflect.Slice:
		if val.Elem().Kind() == reflect.Uint8 {
			return "bytea"
		}
	}

	switch val.Name() {
	case "NullInt64":
		return "bigint"
	case "NullFloat64":
		return "double precision"
	case "NullBool":
		return "boolean"
	case "Time", "NullTime":
		return "timestamp with time zone"
	}

	if maxsize > 0 {
		return fmt.Sprintf("varchar(%d)", maxsize)
	} else {
		return "text"
	}

}

// Returns empty string
func (d PostgresDialect) AutoIncrStr() string {
	return ""
}

func (d PostgresDialect) AutoIncrBindValue() string {
	return "default"
}

func (d PostgresDialect) AutoIncrInsertSuffix(col *ColumnMap) string {
	return " returning " + col.ColumnName
}

// Returns suffix
func (d PostgresDialect) CreateTableSuffix() string {
	return d.suffix
}

func (d PostgresDialect) CreateIndexSuffix() string {
	return "using"
}

func (d PostgresDialect) DropIndexSuffix() string {
	return ""
}

func (d PostgresDialect) TruncateClause() string {
	return "truncate"
}

// Returns "$(i+1)"
func (d PostgresDialect) BindVar(i int) string {
	return fmt.Sprintf("$%d", i+1)
}

func (d PostgresDialect) InsertAutoIncrToTarget(exec SqlExecutor, insertSql string, target interface{}, params ...interface{}) error {
	rows, err := exec.query(insertSql, params...)
	if err != nil {
		return err
	}
	defer rows.Close()

	if !rows.Next() {
		return fmt.Errorf("No serial value returned for insert: %s Encountered error: %s", insertSql, rows.Err())
	}
	if err := rows.Scan(target); err != nil {
		return err
	}
	if rows.Next() {
		return fmt.Errorf("more than two serial value returned for insert: %s", insertSql)
	}
	return rows.Err()
}

func (d PostgresDialect) QuoteField(f string) string {
	return `"` + strings.ToLower(f) + `"`
}

func (d PostgresDialect) QuotedTableForQuery(schema string, table string) string {
	if strings.TrimSpace(schema) == "" {
		return d.QuoteField(table)
	}

	return schema + "." + d.QuoteField(table)
}

func (d PostgresDialect) IfSchemaNotExists(command, schema string) string {
	return fmt.Sprintf("%s if not exists", command)
}

func (d PostgresDialect) IfTableExists(command, schema, table string) string {
	return fmt.Sprintf("%s if exists", command)
}

func (d PostgresDialect) IfTableNotExists(command, schema, table string) string {
	return fmt.Sprintf("%s if not exists", command)
}
<commit_msg>Do not force convert Postgres's table name and field name to lower case. Fixes #119, #192. To use lower case table name, do something like this: dbmap.AddTableWithName(Post{}, "posts") To use lower case field name, specify the field name in type definition: type Post struct { Title   string `db:"title"` }<commit_after>// Copyright 2012 James Cooper. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// Package gorp provides a simple way to marshal Go structs to and from
// SQL databases.  It uses the database/sql package, and should work with any
// compliant database/sql driver.
//
// Source code and project home:
// https://github.com/go-gorp/gorp

package gorp

import (
	"fmt"
	"reflect"
	"strings"
)

type PostgresDialect struct {
	suffix string
}

func (d PostgresDialect) QuerySuffix() string { return ";" }

func (d PostgresDialect) ToSqlType(val reflect.Type, maxsize int, isAutoIncr bool) string {
	switch val.Kind() {
	case reflect.Ptr:
		return d.ToSqlType(val.Elem(), maxsize, isAutoIncr)
	case reflect.Bool:
		return "boolean"
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32:
		if isAutoIncr {
			return "serial"
		}
		return "integer"
	case reflect.Int64, reflect.Uint64:
		if isAutoIncr {
			return "bigserial"
		}
		return "bigint"
	case reflect.Float64:
		return "double precision"
	case reflect.Float32:
		return "real"
	case reflect.Slice:
		if val.Elem().Kind() == reflect.Uint8 {
			return "bytea"
		}
	}

	switch val.Name() {
	case "NullInt64":
		return "bigint"
	case "NullFloat64":
		return "double precision"
	case "NullBool":
		return "boolean"
	case "Time", "NullTime":
		return "timestamp with time zone"
	}

	if maxsize > 0 {
		return fmt.Sprintf("varchar(%d)", maxsize)
	} else {
		return "text"
	}

}

// Returns empty string
func (d PostgresDialect) AutoIncrStr() string {
	return ""
}

func (d PostgresDialect) AutoIncrBindValue() string {
	return "default"
}

func (d PostgresDialect) AutoIncrInsertSuffix(col *ColumnMap) string {
	return " returning " + col.ColumnName
}

// Returns suffix
func (d PostgresDialect) CreateTableSuffix() string {
	return d.suffix
}

func (d PostgresDialect) CreateIndexSuffix() string {
	return "using"
}

func (d PostgresDialect) DropIndexSuffix() string {
	return ""
}

func (d PostgresDialect) TruncateClause() string {
	return "truncate"
}

// Returns "$(i+1)"
func (d PostgresDialect) BindVar(i int) string {
	return fmt.Sprintf("$%d", i+1)
}

func (d PostgresDialect) InsertAutoIncrToTarget(exec SqlExecutor, insertSql string, target interface{}, params ...interface{}) error {
	rows, err := exec.query(insertSql, params...)
	if err != nil {
		return err
	}
	defer rows.Close()

	if !rows.Next() {
		return fmt.Errorf("No serial value returned for insert: %s Encountered error: %s", insertSql, rows.Err())
	}
	if err := rows.Scan(target); err != nil {
		return err
	}
	if rows.Next() {
		return fmt.Errorf("more than two serial value returned for insert: %s", insertSql)
	}
	return rows.Err()
}

func (d PostgresDialect) QuoteField(f string) string {
	return `"` + f + `"`
}

func (d PostgresDialect) QuotedTableForQuery(schema string, table string) string {
	if strings.TrimSpace(schema) == "" {
		return d.QuoteField(table)
	}

	return schema + "." + d.QuoteField(table)
}

func (d PostgresDialect) IfSchemaNotExists(command, schema string) string {
	return fmt.Sprintf("%s if not exists", command)
}

func (d PostgresDialect) IfTableExists(command, schema, table string) string {
	return fmt.Sprintf("%s if exists", command)
}

func (d PostgresDialect) IfTableNotExists(command, schema, table string) string {
	return fmt.Sprintf("%s if not exists", command)
}
<|endoftext|>