<commit_before>package db_test

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

	"github.com/concourse/concourse/atc"
	. "github.com/concourse/concourse/atc/db"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/types"
)

var _ = Describe("Worker", func() {
	var (
		atcWorker atc.Worker
		worker    Worker
	)

	BeforeEach(func() {
		atcWorker = atc.Worker{
			GardenAddr:       "some-garden-addr",
			BaggageclaimURL:  "some-bc-url",
			HTTPProxyURL:     "some-http-proxy-url",
			HTTPSProxyURL:    "some-https-proxy-url",
			NoProxy:          "some-no-proxy",
			Ephemeral:        true,
			ActiveContainers: 140,
			ResourceTypes: []atc.WorkerResourceType{
				{
					Type:    "some-resource-type",
					Image:   "some-image",
					Version: "some-version",
				},
				{
					Type:    "other-resource-type",
					Image:   "other-image",
					Version: "other-version",
				},
			},
			Platform:  "some-platform",
			Tags:      atc.Tags{"some", "tags"},
			Name:      "some-name",
			StartTime: 55912945,
		}
	})

	Describe("Land", func() {
		BeforeEach(func() {
			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())
		})

		Context("when the worker is present", func() {
			It("marks the worker as `landing`", func() {
				err := worker.Land()
				Expect(err).NotTo(HaveOccurred())

				_, err = worker.Reload()
				Expect(err).NotTo(HaveOccurred())
				Expect(worker.Name()).To(Equal(atcWorker.Name))
				Expect(worker.State()).To(Equal(WorkerStateLanding))
			})

			Context("when worker is already landed", func() {
				BeforeEach(func() {
					err := worker.Land()
					Expect(err).NotTo(HaveOccurred())
					_, err = workerLifecycle.LandFinishedLandingWorkers()
					Expect(err).NotTo(HaveOccurred())
				})

				It("keeps worker state as landed", func() {
					err := worker.Land()
					Expect(err).NotTo(HaveOccurred())
					_, err = worker.Reload()
					Expect(err).NotTo(HaveOccurred())

					Expect(worker.Name()).To(Equal(atcWorker.Name))
					Expect(worker.State()).To(Equal(WorkerStateLanded))
				})
			})
		})

		Context("when the worker is not present", func() {
			It("returns an error", func() {
				err := worker.Delete()
				Expect(err).NotTo(HaveOccurred())

				err = worker.Land()
				Expect(err).To(HaveOccurred())
				Expect(err).To(Equal(ErrWorkerNotPresent))
			})
		})
	})

	Describe("Retire", func() {
		BeforeEach(func() {
			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())
		})

		Context("when the worker is present", func() {
			It("marks the worker as `retiring`", func() {
				err := worker.Retire()
				Expect(err).NotTo(HaveOccurred())

				_, err = worker.Reload()
				Expect(err).NotTo(HaveOccurred())
				Expect(worker.Name()).To(Equal(atcWorker.Name))
				Expect(worker.State()).To(Equal(WorkerStateRetiring))
			})
		})

		Context("when the worker is not present", func() {
			BeforeEach(func() {
				err := worker.Delete()
				Expect(err).NotTo(HaveOccurred())
			})

			It("returns an error", func() {
				err := worker.Retire()
				Expect(err).To(HaveOccurred())
				Expect(err).To(Equal(ErrWorkerNotPresent))
			})
		})
	})

	Describe("Delete", func() {
		BeforeEach(func() {
			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())
		})

		It("deletes the record for the worker", func() {
			err := worker.Delete()
			Expect(err).NotTo(HaveOccurred())

			_, found, err := workerFactory.GetWorker(atcWorker.Name)
			Expect(err).ToNot(HaveOccurred())
			Expect(found).To(BeFalse())
		})
	})

	Describe("Prune", func() {
		Context("when worker exists", func() {
			DescribeTable("worker in state",
				func(workerState string, errMatch GomegaMatcher) {
					worker, err := workerFactory.SaveWorker(atc.Worker{
						Name:       "worker-to-prune",
						GardenAddr: "1.2.3.4",
						State:      workerState,
					}, 5*time.Minute)
					Expect(err).NotTo(HaveOccurred())

					err = worker.Prune()
					Expect(err).To(errMatch)
				},

				Entry("running", "running", Equal(ErrCannotPruneRunningWorker)),
				Entry("landing", "landing", BeNil()),
				Entry("retiring", "retiring", BeNil()),
			)

			Context("when worker is stalled", func() {
				var pruneErr error
				BeforeEach(func() {
					worker, err := workerFactory.SaveWorker(atc.Worker{
						Name:       "worker-to-prune",
						GardenAddr: "1.2.3.4",
						State:      "running",
					}, -5*time.Minute)
					Expect(err).NotTo(HaveOccurred())

					_, err = workerLifecycle.StallUnresponsiveWorkers()
					Expect(err).NotTo(HaveOccurred())
					pruneErr = worker.Prune()
				})

				It("does not return error", func() {
					Expect(pruneErr).NotTo(HaveOccurred())
				})
			})
		})

		Context("when worker does not exist", func() {
			BeforeEach(func() {
				var err error
				worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
				Expect(err).NotTo(HaveOccurred())
				err = worker.Delete()
				Expect(err).NotTo(HaveOccurred())
			})

			It("raises ErrWorkerNotPresent", func() {
				err := worker.Prune()
				Expect(err).To(Equal(ErrWorkerNotPresent))
			})
		})
	})

	Describe("FindContainer/CreateContainer", func() {
		var (
			containerMetadata ContainerMetadata
			containerOwner    ContainerOwner

			foundCreatingContainer CreatingContainer
			foundCreatedContainer  CreatedContainer
			worker                 Worker
		)

		expiries := ContainerOwnerExpiries{
			Min: 5 * time.Minute,
			Max: 1 * time.Hour,
		}

		BeforeEach(func() {
			containerMetadata = ContainerMetadata{
				Type: "check",
			}

			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())

			atcWorker2 := atcWorker
			atcWorker2.Name = "some-name2"
			atcWorker2.GardenAddr = "some-garden-addr-other"
			otherWorker, err = workerFactory.SaveWorker(atcWorker2, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())

			resourceConfig, err := resourceConfigFactory.FindOrCreateResourceConfig(
				"some-resource-type",
				atc.Source{"some": "source"},
				atc.VersionedResourceTypes{},
			)
			Expect(err).ToNot(HaveOccurred())

			containerOwner = NewResourceConfigCheckSessionContainerOwner(
				resourceConfig.ID(),
				resourceConfig.OriginBaseResourceType().ID,
				expiries,
			)
		})

		JustBeforeEach(func() {
			var err error
			foundCreatingContainer, foundCreatedContainer, err = worker.FindContainer(containerOwner)
			Expect(err).ToNot(HaveOccurred())
		})

		Context("when there is a creating container", func() {
			var creatingContainer CreatingContainer

			BeforeEach(func() {
				var err error
				creatingContainer, err = worker.CreateContainer(containerOwner, containerMetadata)
				Expect(err).ToNot(HaveOccurred())
			})

			It("returns it", func() {
				Expect(foundCreatedContainer).To(BeNil())
				Expect(foundCreatingContainer).ToNot(BeNil())
			})

			Context("when finding on another worker", func() {
				BeforeEach(func() {
					worker = otherWorker
				})

				It("does not find it", func() {
					Expect(foundCreatingContainer).To(BeNil())
					Expect(foundCreatedContainer).To(BeNil())
				})
			})

			Context("when there is a created container", func() {
				BeforeEach(func() {
					_, err := creatingContainer.Created()
					Expect(err).ToNot(HaveOccurred())
				})

				It("returns it", func() {
					Expect(foundCreatedContainer).ToNot(BeNil())
					Expect(foundCreatingContainer).To(BeNil())
				})

				Context("when finding on another worker", func() {
					BeforeEach(func() {
						worker = otherWorker
					})

					It("does not find it", func() {
						Expect(foundCreatingContainer).To(BeNil())
						Expect(foundCreatedContainer).To(BeNil())
					})
				})
			})

			Context("when the creating container is failed and gced", func() {
				BeforeEach(func() {
					var err error
					_, err = creatingContainer.Failed()
					Expect(err).ToNot(HaveOccurred())

					containerRepository := NewContainerRepository(dbConn)
					containersDestroyed, err := containerRepository.DestroyFailedContainers()
					Expect(containersDestroyed).To(Equal(1))
					Expect(err).ToNot(HaveOccurred())

					var checkSessions int
					err = dbConn.QueryRow("SELECT COUNT(*) FROM resource_config_check_sessions").Scan(&checkSessions)
					Expect(err).ToNot(HaveOccurred())
					Expect(checkSessions).To(Equal(1))
				})

				Context("and we create a new container", func() {
					BeforeEach(func() {
						_, err := worker.CreateContainer(containerOwner, containerMetadata)
						Expect(err).ToNot(HaveOccurred())
					})

					It("does not duplicate the resource config check session", func() {
						var checkSessions int
						err := dbConn.QueryRow("SELECT COUNT(*) FROM resource_config_check_sessions").Scan(&checkSessions)
						Expect(err).ToNot(HaveOccurred())
						Expect(checkSessions).To(Equal(1))
					})
				})
			})
		})

		Context("when there is no container", func() {
			It("returns nil", func() {
				Expect(foundCreatedContainer).To(BeNil())
				Expect(foundCreatingContainer).To(BeNil())
			})
		})

		Context("when the container has a meta type", func() {
			var container CreatingContainer

			Context("when the meta type is check", func() {
				BeforeEach(func() {
					containerMetadata = ContainerMetadata{
						Type: "check",
					}

					var err error
					container, err = worker.CreateContainer(containerOwner, containerMetadata)
					Expect(err).ToNot(HaveOccurred())
				})

				It("returns a container with empty team id", func() {
					var teamID sql.NullString

					err := dbConn.QueryRow(fmt.Sprintf("SELECT team_id FROM containers WHERE id='%d'", container.ID())).Scan(&teamID)
					Expect(err).ToNot(HaveOccurred())
					Expect(teamID.Valid).To(BeFalse())
				})
			})

			Context("when the meta type is not check", func() {
				BeforeEach(func() {
					containerMetadata = ContainerMetadata{
						Type: "get",
					}

					oneOffBuild, err := defaultTeam.CreateOneOffBuild()
					Expect(err).ToNot(HaveOccurred())

					container, err = worker.CreateContainer(NewBuildStepContainerOwner(oneOffBuild.ID(), atc.PlanID("1"), 1), containerMetadata)
					Expect(err).ToNot(HaveOccurred())
				})

				It("returns a container with a team id", func() {
					var teamID sql.NullString

					err := dbConn.QueryRow(fmt.Sprintf("SELECT team_id FROM containers WHERE id='%d'", container.ID())).Scan(&teamID)
					Expect(err).ToNot(HaveOccurred())
					Expect(teamID.Valid).To(BeTrue())
				})
			})
		})
	})

	Describe("Active tasks", func() {
		BeforeEach(func() {
			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())
		})

		Context("when the worker registers", func() {
			It("has no active tasks", func() {
				at, err := worker.ActiveTasks()
				Expect(err).ToNot(HaveOccurred())
				Expect(at).To(Equal(0))
			})
		})

		Context("when the active task is increased", func() {
			BeforeEach(func() {
				at, err := worker.IncreaseActiveTasks()
				Expect(err).ToNot(HaveOccurred())
				Expect(at).To(Equal(1))
			})

			It("increase the active tasks counter", func() {
				at, err := worker.ActiveTasks()
				Expect(err).ToNot(HaveOccurred())
				Expect(at).To(Equal(1))
			})

			Context("when the active task is decreased", func() {
				BeforeEach(func() {
					at, err := worker.DecreaseActiveTasks()
					Expect(err).ToNot(HaveOccurred())
					Expect(at).To(Equal(0))
				})

				It("reset the active tasks to 0", func() {
					at, err := worker.ActiveTasks()
					Expect(err).ToNot(HaveOccurred())
					Expect(at).To(Equal(0))
				})
			})
		})

		Context("when the active task is decreased below 0", func() {
			It("raise an error", func() {
				at, err := worker.DecreaseActiveTasks()
				Expect(err).To(HaveOccurred())
				Expect(at).To(Equal(0))
			})
		})
	})
})
<commit_msg>fix gomega compilation error<commit_after>package db_test

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

	"github.com/concourse/concourse/atc"
	. "github.com/concourse/concourse/atc/db"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/types"
)

var _ = Describe("Worker", func() {
	var (
		atcWorker atc.Worker
		worker    Worker
	)

	BeforeEach(func() {
		atcWorker = atc.Worker{
			GardenAddr:       "some-garden-addr",
			BaggageclaimURL:  "some-bc-url",
			HTTPProxyURL:     "some-http-proxy-url",
			HTTPSProxyURL:    "some-https-proxy-url",
			NoProxy:          "some-no-proxy",
			Ephemeral:        true,
			ActiveContainers: 140,
			ResourceTypes: []atc.WorkerResourceType{
				{
					Type:    "some-resource-type",
					Image:   "some-image",
					Version: "some-version",
				},
				{
					Type:    "other-resource-type",
					Image:   "other-image",
					Version: "other-version",
				},
			},
			Platform:  "some-platform",
			Tags:      atc.Tags{"some", "tags"},
			Name:      "some-name",
			StartTime: 55912945,
		}
	})

	Describe("Land", func() {
		BeforeEach(func() {
			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())
		})

		Context("when the worker is present", func() {
			It("marks the worker as `landing`", func() {
				err := worker.Land()
				Expect(err).NotTo(HaveOccurred())

				_, err = worker.Reload()
				Expect(err).NotTo(HaveOccurred())
				Expect(worker.Name()).To(Equal(atcWorker.Name))
				Expect(worker.State()).To(Equal(WorkerStateLanding))
			})

			Context("when worker is already landed", func() {
				BeforeEach(func() {
					err := worker.Land()
					Expect(err).NotTo(HaveOccurred())
					_, err = workerLifecycle.LandFinishedLandingWorkers()
					Expect(err).NotTo(HaveOccurred())
				})

				It("keeps worker state as landed", func() {
					err := worker.Land()
					Expect(err).NotTo(HaveOccurred())
					_, err = worker.Reload()
					Expect(err).NotTo(HaveOccurred())

					Expect(worker.Name()).To(Equal(atcWorker.Name))
					Expect(worker.State()).To(Equal(WorkerStateLanded))
				})
			})
		})

		Context("when the worker is not present", func() {
			It("returns an error", func() {
				err := worker.Delete()
				Expect(err).NotTo(HaveOccurred())

				err = worker.Land()
				Expect(err).To(HaveOccurred())
				Expect(err).To(Equal(ErrWorkerNotPresent))
			})
		})
	})

	Describe("Retire", func() {
		BeforeEach(func() {
			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())
		})

		Context("when the worker is present", func() {
			It("marks the worker as `retiring`", func() {
				err := worker.Retire()
				Expect(err).NotTo(HaveOccurred())

				_, err = worker.Reload()
				Expect(err).NotTo(HaveOccurred())
				Expect(worker.Name()).To(Equal(atcWorker.Name))
				Expect(worker.State()).To(Equal(WorkerStateRetiring))
			})
		})

		Context("when the worker is not present", func() {
			BeforeEach(func() {
				err := worker.Delete()
				Expect(err).NotTo(HaveOccurred())
			})

			It("returns an error", func() {
				err := worker.Retire()
				Expect(err).To(HaveOccurred())
				Expect(err).To(Equal(ErrWorkerNotPresent))
			})
		})
	})

	Describe("Delete", func() {
		BeforeEach(func() {
			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())
		})

		It("deletes the record for the worker", func() {
			err := worker.Delete()
			Expect(err).NotTo(HaveOccurred())

			_, found, err := workerFactory.GetWorker(atcWorker.Name)
			Expect(err).ToNot(HaveOccurred())
			Expect(found).To(BeFalse())
		})
	})

	Describe("Prune", func() {
		Context("when worker exists", func() {
			DescribeTable("worker in state",
				func(workerState string, errMatch types.GomegaMatcher) {
					worker, err := workerFactory.SaveWorker(atc.Worker{
						Name:       "worker-to-prune",
						GardenAddr: "1.2.3.4",
						State:      workerState,
					}, 5*time.Minute)
					Expect(err).NotTo(HaveOccurred())

					err = worker.Prune()
					Expect(err).To(errMatch)
				},

				Entry("running", "running", Equal(ErrCannotPruneRunningWorker)),
				Entry("landing", "landing", BeNil()),
				Entry("retiring", "retiring", BeNil()),
			)

			Context("when worker is stalled", func() {
				var pruneErr error
				BeforeEach(func() {
					worker, err := workerFactory.SaveWorker(atc.Worker{
						Name:       "worker-to-prune",
						GardenAddr: "1.2.3.4",
						State:      "running",
					}, -5*time.Minute)
					Expect(err).NotTo(HaveOccurred())

					_, err = workerLifecycle.StallUnresponsiveWorkers()
					Expect(err).NotTo(HaveOccurred())
					pruneErr = worker.Prune()
				})

				It("does not return error", func() {
					Expect(pruneErr).NotTo(HaveOccurred())
				})
			})
		})

		Context("when worker does not exist", func() {
			BeforeEach(func() {
				var err error
				worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
				Expect(err).NotTo(HaveOccurred())
				err = worker.Delete()
				Expect(err).NotTo(HaveOccurred())
			})

			It("raises ErrWorkerNotPresent", func() {
				err := worker.Prune()
				Expect(err).To(Equal(ErrWorkerNotPresent))
			})
		})
	})

	Describe("FindContainer/CreateContainer", func() {
		var (
			containerMetadata ContainerMetadata
			containerOwner    ContainerOwner

			foundCreatingContainer CreatingContainer
			foundCreatedContainer  CreatedContainer
			worker                 Worker
		)

		expiries := ContainerOwnerExpiries{
			Min: 5 * time.Minute,
			Max: 1 * time.Hour,
		}

		BeforeEach(func() {
			containerMetadata = ContainerMetadata{
				Type: "check",
			}

			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())

			atcWorker2 := atcWorker
			atcWorker2.Name = "some-name2"
			atcWorker2.GardenAddr = "some-garden-addr-other"
			otherWorker, err = workerFactory.SaveWorker(atcWorker2, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())

			resourceConfig, err := resourceConfigFactory.FindOrCreateResourceConfig(
				"some-resource-type",
				atc.Source{"some": "source"},
				atc.VersionedResourceTypes{},
			)
			Expect(err).ToNot(HaveOccurred())

			containerOwner = NewResourceConfigCheckSessionContainerOwner(
				resourceConfig.ID(),
				resourceConfig.OriginBaseResourceType().ID,
				expiries,
			)
		})

		JustBeforeEach(func() {
			var err error
			foundCreatingContainer, foundCreatedContainer, err = worker.FindContainer(containerOwner)
			Expect(err).ToNot(HaveOccurred())
		})

		Context("when there is a creating container", func() {
			var creatingContainer CreatingContainer

			BeforeEach(func() {
				var err error
				creatingContainer, err = worker.CreateContainer(containerOwner, containerMetadata)
				Expect(err).ToNot(HaveOccurred())
			})

			It("returns it", func() {
				Expect(foundCreatedContainer).To(BeNil())
				Expect(foundCreatingContainer).ToNot(BeNil())
			})

			Context("when finding on another worker", func() {
				BeforeEach(func() {
					worker = otherWorker
				})

				It("does not find it", func() {
					Expect(foundCreatingContainer).To(BeNil())
					Expect(foundCreatedContainer).To(BeNil())
				})
			})

			Context("when there is a created container", func() {
				BeforeEach(func() {
					_, err := creatingContainer.Created()
					Expect(err).ToNot(HaveOccurred())
				})

				It("returns it", func() {
					Expect(foundCreatedContainer).ToNot(BeNil())
					Expect(foundCreatingContainer).To(BeNil())
				})

				Context("when finding on another worker", func() {
					BeforeEach(func() {
						worker = otherWorker
					})

					It("does not find it", func() {
						Expect(foundCreatingContainer).To(BeNil())
						Expect(foundCreatedContainer).To(BeNil())
					})
				})
			})

			Context("when the creating container is failed and gced", func() {
				BeforeEach(func() {
					var err error
					_, err = creatingContainer.Failed()
					Expect(err).ToNot(HaveOccurred())

					containerRepository := NewContainerRepository(dbConn)
					containersDestroyed, err := containerRepository.DestroyFailedContainers()
					Expect(containersDestroyed).To(Equal(1))
					Expect(err).ToNot(HaveOccurred())

					var checkSessions int
					err = dbConn.QueryRow("SELECT COUNT(*) FROM resource_config_check_sessions").Scan(&checkSessions)
					Expect(err).ToNot(HaveOccurred())
					Expect(checkSessions).To(Equal(1))
				})

				Context("and we create a new container", func() {
					BeforeEach(func() {
						_, err := worker.CreateContainer(containerOwner, containerMetadata)
						Expect(err).ToNot(HaveOccurred())
					})

					It("does not duplicate the resource config check session", func() {
						var checkSessions int
						err := dbConn.QueryRow("SELECT COUNT(*) FROM resource_config_check_sessions").Scan(&checkSessions)
						Expect(err).ToNot(HaveOccurred())
						Expect(checkSessions).To(Equal(1))
					})
				})
			})
		})

		Context("when there is no container", func() {
			It("returns nil", func() {
				Expect(foundCreatedContainer).To(BeNil())
				Expect(foundCreatingContainer).To(BeNil())
			})
		})

		Context("when the container has a meta type", func() {
			var container CreatingContainer

			Context("when the meta type is check", func() {
				BeforeEach(func() {
					containerMetadata = ContainerMetadata{
						Type: "check",
					}

					var err error
					container, err = worker.CreateContainer(containerOwner, containerMetadata)
					Expect(err).ToNot(HaveOccurred())
				})

				It("returns a container with empty team id", func() {
					var teamID sql.NullString

					err := dbConn.QueryRow(fmt.Sprintf("SELECT team_id FROM containers WHERE id='%d'", container.ID())).Scan(&teamID)
					Expect(err).ToNot(HaveOccurred())
					Expect(teamID.Valid).To(BeFalse())
				})
			})

			Context("when the meta type is not check", func() {
				BeforeEach(func() {
					containerMetadata = ContainerMetadata{
						Type: "get",
					}

					oneOffBuild, err := defaultTeam.CreateOneOffBuild()
					Expect(err).ToNot(HaveOccurred())

					container, err = worker.CreateContainer(NewBuildStepContainerOwner(oneOffBuild.ID(), atc.PlanID("1"), 1), containerMetadata)
					Expect(err).ToNot(HaveOccurred())
				})

				It("returns a container with a team id", func() {
					var teamID sql.NullString

					err := dbConn.QueryRow(fmt.Sprintf("SELECT team_id FROM containers WHERE id='%d'", container.ID())).Scan(&teamID)
					Expect(err).ToNot(HaveOccurred())
					Expect(teamID.Valid).To(BeTrue())
				})
			})
		})
	})

	Describe("Active tasks", func() {
		BeforeEach(func() {
			var err error
			worker, err = workerFactory.SaveWorker(atcWorker, 5*time.Minute)
			Expect(err).NotTo(HaveOccurred())
		})

		Context("when the worker registers", func() {
			It("has no active tasks", func() {
				at, err := worker.ActiveTasks()
				Expect(err).ToNot(HaveOccurred())
				Expect(at).To(Equal(0))
			})
		})

		Context("when the active task is increased", func() {
			BeforeEach(func() {
				at, err := worker.IncreaseActiveTasks()
				Expect(err).ToNot(HaveOccurred())
				Expect(at).To(Equal(1))
			})

			It("increase the active tasks counter", func() {
				at, err := worker.ActiveTasks()
				Expect(err).ToNot(HaveOccurred())
				Expect(at).To(Equal(1))
			})

			Context("when the active task is decreased", func() {
				BeforeEach(func() {
					at, err := worker.DecreaseActiveTasks()
					Expect(err).ToNot(HaveOccurred())
					Expect(at).To(Equal(0))
				})

				It("reset the active tasks to 0", func() {
					at, err := worker.ActiveTasks()
					Expect(err).ToNot(HaveOccurred())
					Expect(at).To(Equal(0))
				})
			})
		})

		Context("when the active task is decreased below 0", func() {
			It("raise an error", func() {
				at, err := worker.DecreaseActiveTasks()
				Expect(err).To(HaveOccurred())
				Expect(at).To(Equal(0))
			})
		})
	})
})
<|endoftext|>


<commit_before>package policy

import (
	"fmt"
	"strings"

	"code.cloudfoundry.org/lager"
	"github.com/jessevdk/go-flags"
)

const ActionUsingImage = "UsingImage"

type PolicyCheckNotPass struct{}

func (e PolicyCheckNotPass) Error() string {
	return "policy check not pass"
}

type Filter struct {
	HttpMethods   []string `long:"policy-check-filter-http-methods" description:"API http method to go through policy check"`
	Actions       []string `long:"policy-check-filter-action" description:"Actions in the list will go through policy check"`
	ActionsToSkip []string `long:"policy-check-filter-action-skip" default:"UsingImage" description:"Actions the list will not go through policy check"`
}

func (f Filter) normalize() Filter {
	if len(f.HttpMethods) == 1 {
		f.HttpMethods = strings.Split(f.HttpMethods[0], ",")
	}

	if len(f.Actions) == 1 {
		f.Actions = strings.Split(f.Actions[0], ",")
	}

	if len(f.ActionsToSkip) == 1 {
		f.ActionsToSkip = strings.Split(f.ActionsToSkip[0], ",")
	}

	return f
}

type PolicyCheckInput struct {
	Service        string      `json:"service"`
	ClusterName    string      `json:"cluster_name"`
	ClusterVersion string      `json:"cluster_version"`
	HttpMethod     string      `json:"http_method,omitempty"`
	Action         string      `json:"action"`
	User           string      `json:"user"`
	Team           string      `json:"team,omitempty"`
	Pipeline       string      `json:"pipeline,omitempty"`
	Data           interface{} `json:"data,omitempty"`
}

//go:generate counterfeiter . Agent

// Agent should be implemented by policy agents.
type Agent interface {
	// Check returns true if passes policy check. If not goes through policy
	// check, just return true.
	Check(PolicyCheckInput) (bool, error)
}

//go:generate counterfeiter . AgentFactory

type AgentFactory interface {
	Description() string
	IsConfigured() bool
	NewAgent(lager.Logger) (Agent, error)
}

var agentFactories []AgentFactory

func RegisterAgent(factory AgentFactory) {
	agentFactories = append(agentFactories, factory)
}

func WireCheckers(group *flags.Group) {
	for _, factory := range agentFactories {
		_, err := group.AddGroup(fmt.Sprintf("Policy Check Agent (%s)", factory.Description()), "", factory)
		if err != nil {
			panic(err)
		}
	}
}

var (
	clusterName    string
	clusterVersion string
)

func Initialize(logger lager.Logger, cluster string, version string, filter Filter) (*Checker, error) {
	logger.Debug("policy-checker-initialize")

	clusterName = cluster
	clusterVersion = version

	var checkerDescriptions []string
	for _, factory := range agentFactories {
		if factory.IsConfigured() {
			checkerDescriptions = append(checkerDescriptions, factory.Description())
		}
	}
	if len(checkerDescriptions) > 1 {
		return nil, fmt.Errorf("Multiple policy checker configured: %s", strings.Join(checkerDescriptions, ", "))
	}

	for _, factory := range agentFactories {
		if factory.IsConfigured() {
			agent, err := factory.NewAgent(logger.Session("policy-checker"))
			if err != nil {
				return nil, err
			}
			return &Checker{
				filter: filter.normalize(),
				agent:  agent,
			}, nil
		}
	}

	// No policy checker configured.
	return nil, nil
}

type Checker struct {
	filter Filter
	agent  Agent
}

func (c *Checker) ShouldCheckHttpMethod(method string) bool {
	return inArray(c.filter.HttpMethods, method)
}

func (c *Checker) ShouldCheckAction(action string) bool {
	return inArray(c.filter.Actions, action)
}

func (c *Checker) ShouldSkipAction(action string) bool {
	return inArray(c.filter.ActionsToSkip, action)
}

func inArray(array []string, target string) bool {
	found := false
	for _, ele := range array {
		if ele == target {
			found = true
			break
		}
	}
	return found
}

func (c *Checker) Check(input PolicyCheckInput) (bool, error) {
	input.Service = "concourse"
	input.ClusterName = clusterName
	input.ClusterVersion = clusterVersion
	return c.agent.Check(input)
}
<commit_msg>based on vito's comment in the rfc, rename policy-check-filter-http-methods to policy-check-filter-http-method.<commit_after>package policy

import (
	"fmt"
	"strings"

	"code.cloudfoundry.org/lager"
	"github.com/jessevdk/go-flags"
)

const ActionUsingImage = "UsingImage"

type PolicyCheckNotPass struct{}

func (e PolicyCheckNotPass) Error() string {
	return "policy check not pass"
}

type Filter struct {
	HttpMethods   []string `long:"policy-check-filter-http-method" description:"API http method to go through policy check"`
	Actions       []string `long:"policy-check-filter-action" description:"Actions in the list will go through policy check"`
	ActionsToSkip []string `long:"policy-check-filter-action-skip" default:"UsingImage" description:"Actions the list will not go through policy check"`
}

func (f Filter) normalize() Filter {
	if len(f.HttpMethods) == 1 {
		f.HttpMethods = strings.Split(f.HttpMethods[0], ",")
	}

	if len(f.Actions) == 1 {
		f.Actions = strings.Split(f.Actions[0], ",")
	}

	if len(f.ActionsToSkip) == 1 {
		f.ActionsToSkip = strings.Split(f.ActionsToSkip[0], ",")
	}

	return f
}

type PolicyCheckInput struct {
	Service        string      `json:"service"`
	ClusterName    string      `json:"cluster_name"`
	ClusterVersion string      `json:"cluster_version"`
	HttpMethod     string      `json:"http_method,omitempty"`
	Action         string      `json:"action"`
	User           string      `json:"user"`
	Team           string      `json:"team,omitempty"`
	Pipeline       string      `json:"pipeline,omitempty"`
	Data           interface{} `json:"data,omitempty"`
}

//go:generate counterfeiter . Agent

// Agent should be implemented by policy agents.
type Agent interface {
	// Check returns true if passes policy check. If not goes through policy
	// check, just return true.
	Check(PolicyCheckInput) (bool, error)
}

//go:generate counterfeiter . AgentFactory

type AgentFactory interface {
	Description() string
	IsConfigured() bool
	NewAgent(lager.Logger) (Agent, error)
}

var agentFactories []AgentFactory

func RegisterAgent(factory AgentFactory) {
	agentFactories = append(agentFactories, factory)
}

func WireCheckers(group *flags.Group) {
	for _, factory := range agentFactories {
		_, err := group.AddGroup(fmt.Sprintf("Policy Check Agent (%s)", factory.Description()), "", factory)
		if err != nil {
			panic(err)
		}
	}
}

var (
	clusterName    string
	clusterVersion string
)

func Initialize(logger lager.Logger, cluster string, version string, filter Filter) (*Checker, error) {
	logger.Debug("policy-checker-initialize")

	clusterName = cluster
	clusterVersion = version

	var checkerDescriptions []string
	for _, factory := range agentFactories {
		if factory.IsConfigured() {
			checkerDescriptions = append(checkerDescriptions, factory.Description())
		}
	}
	if len(checkerDescriptions) > 1 {
		return nil, fmt.Errorf("Multiple policy checker configured: %s", strings.Join(checkerDescriptions, ", "))
	}

	for _, factory := range agentFactories {
		if factory.IsConfigured() {
			agent, err := factory.NewAgent(logger.Session("policy-checker"))
			if err != nil {
				return nil, err
			}
			return &Checker{
				filter: filter.normalize(),
				agent:  agent,
			}, nil
		}
	}

	// No policy checker configured.
	return nil, nil
}

type Checker struct {
	filter Filter
	agent  Agent
}

func (c *Checker) ShouldCheckHttpMethod(method string) bool {
	return inArray(c.filter.HttpMethods, method)
}

func (c *Checker) ShouldCheckAction(action string) bool {
	return inArray(c.filter.Actions, action)
}

func (c *Checker) ShouldSkipAction(action string) bool {
	return inArray(c.filter.ActionsToSkip, action)
}

func inArray(array []string, target string) bool {
	found := false
	for _, ele := range array {
		if ele == target {
			found = true
			break
		}
	}
	return found
}

func (c *Checker) Check(input PolicyCheckInput) (bool, error) {
	input.Service = "concourse"
	input.ClusterName = clusterName
	input.ClusterVersion = clusterVersion
	return c.agent.Check(input)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package lifecycle

import (
	"fmt"
	"regexp"
	"time"

	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/prow/github"
	"k8s.io/test-infra/prow/pluginhelp"
	"k8s.io/test-infra/prow/plugins"
)

const deprecatedWarn = true

var (
	deprecatedTick = time.Tick(time.Hour) // Warn once per hour
	lifecycleRe    = regexp.MustCompile(`(?mi)^/(remove-)?lifecycle (frozen|stale|putrid|rotten)\s*$`)
)

func init() {
	plugins.RegisterGenericCommentHandler("lifecycle", lifecycleHandleGenericComment, help)
	logrus.SetLevel(logrus.DebugLevel)
}

func help(config *plugins.Configuration, enabledRepos []string) (*pluginhelp.PluginHelp, error) {
	pluginHelp := &pluginhelp.PluginHelp{
		Description: "Close, reopen, flag and/or unflag an issue or PR as stale/putrid/rotten/frozen",
	}
	pluginHelp.AddCommand(pluginhelp.Command{
		Usage:       "/close",
		Description: "Closes an issue or PR.",
		Featured:    false,
		WhoCanUse:   "Authors and assignees can triggers this command.",
		Examples:    []string{"/close"},
	})
	pluginHelp.AddCommand(pluginhelp.Command{
		Usage:       "/reopen",
		Description: "Reopens an issue or PR",
		Featured:    false,
		WhoCanUse:   "Authors and assignees can trigger this command.",
		Examples:    []string{"/reopen"},
	})
	pluginHelp.AddCommand(pluginhelp.Command{
		Usage:       "/[remove-]lifecycle <frozen|stale|putrid|rotten>",
		Description: "Flags an issue or PR as frozen/stale/putrid/rotten",
		Featured:    false,
		WhoCanUse:   "Anyone can trigger this command.",
		Examples:    []string{"/lifecycle frozen", "/remove-lifecycle stale"},
	})
	return pluginHelp, nil
}

type commentClient interface {
	CreateComment(owner, repo string, number int, comment string) error
}

type lifecycleClient interface {
	CreateComment(owner, repo string, number int, comment string) error
	AddLabel(owner, repo string, number int, label string) error
	RemoveLabel(owner, repo string, number int, label string) error
}

func deprecate(gc commentClient, plugin, org, repo string, number int, e *github.GenericCommentEvent) error {
	select {
	case <-deprecatedTick:
		// Only warn once per tick
		return gc.CreateComment(org, repo, number, plugins.FormatResponseRaw(e.Body, e.HTMLURL, e.User.Login, fmt.Sprintf("The %s prow plugin is deprecated, please migrate to the lifecycle plugin before April 2018", plugin)))
	default:
		return nil
	}
}

func lifecycleHandleGenericComment(pc plugins.PluginClient, e github.GenericCommentEvent) error {
	gc := pc.GitHubClient
	log := pc.Logger
	if err := handleReopen(gc, log, &e, !deprecatedWarn); err != nil {
		return err
	}
	if err := handleClose(gc, log, &e, !deprecatedWarn); err != nil {
		return err
	}
	return handle(gc, log, &e)
}

func handle(gc lifecycleClient, log *logrus.Entry, e *github.GenericCommentEvent) error {
	// Only consider new comments.
	if e.Action != github.GenericCommentActionCreated {
		return nil
	}

	for _, mat := range lifecycleRe.FindAllStringSubmatch(e.Body, -1) {
		if err := handleOne(gc, log, e, mat); err != nil {
			return err
		}
	}
	return nil
}

func handleOne(gc lifecycleClient, log *logrus.Entry, e *github.GenericCommentEvent, mat []string) error {
	org := e.Repo.Owner.Login
	repo := e.Repo.Name
	number := e.Number

	remove := mat[1] != ""
	cmd := mat[2]
	lbl := "lifecycle/" + cmd
	// Let's start simple and allow anyone to add/remove frozen, stale, putrid, rotten labels.
	// Adjust if we find evidence of the community abusing these labels.
	if remove {
		log.Infof("/remove-%s", cmd)
		return gc.RemoveLabel(org, repo, number, lbl)
	}
	log.Infof("/%s", cmd)
	return gc.AddLabel(org, repo, number, lbl)
}
<commit_msg>Minor fixes to lifecycle plugin logging.<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package lifecycle

import (
	"fmt"
	"regexp"
	"time"

	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/prow/github"
	"k8s.io/test-infra/prow/pluginhelp"
	"k8s.io/test-infra/prow/plugins"
)

const deprecatedWarn = true

var (
	deprecatedTick = time.Tick(time.Hour) // Warn once per hour
	lifecycleRe    = regexp.MustCompile(`(?mi)^/(remove-)?lifecycle (frozen|stale|putrid|rotten)\s*$`)
)

func init() {
	plugins.RegisterGenericCommentHandler("lifecycle", lifecycleHandleGenericComment, help)
}

func help(config *plugins.Configuration, enabledRepos []string) (*pluginhelp.PluginHelp, error) {
	pluginHelp := &pluginhelp.PluginHelp{
		Description: "Close, reopen, flag and/or unflag an issue or PR as stale/putrid/rotten/frozen",
	}
	pluginHelp.AddCommand(pluginhelp.Command{
		Usage:       "/close",
		Description: "Closes an issue or PR.",
		Featured:    false,
		WhoCanUse:   "Authors and assignees can triggers this command.",
		Examples:    []string{"/close"},
	})
	pluginHelp.AddCommand(pluginhelp.Command{
		Usage:       "/reopen",
		Description: "Reopens an issue or PR",
		Featured:    false,
		WhoCanUse:   "Authors and assignees can trigger this command.",
		Examples:    []string{"/reopen"},
	})
	pluginHelp.AddCommand(pluginhelp.Command{
		Usage:       "/[remove-]lifecycle <frozen|stale|putrid|rotten>",
		Description: "Flags an issue or PR as frozen/stale/putrid/rotten",
		Featured:    false,
		WhoCanUse:   "Anyone can trigger this command.",
		Examples:    []string{"/lifecycle frozen", "/remove-lifecycle stale"},
	})
	return pluginHelp, nil
}

type commentClient interface {
	CreateComment(owner, repo string, number int, comment string) error
}

type lifecycleClient interface {
	CreateComment(owner, repo string, number int, comment string) error
	AddLabel(owner, repo string, number int, label string) error
	RemoveLabel(owner, repo string, number int, label string) error
}

func deprecate(gc commentClient, plugin, org, repo string, number int, e *github.GenericCommentEvent) error {
	select {
	case <-deprecatedTick:
		// Only warn once per tick
		return gc.CreateComment(org, repo, number, plugins.FormatResponseRaw(e.Body, e.HTMLURL, e.User.Login, fmt.Sprintf("The %s prow plugin is deprecated, please migrate to the lifecycle plugin before April 2018", plugin)))
	default:
		return nil
	}
}

func lifecycleHandleGenericComment(pc plugins.PluginClient, e github.GenericCommentEvent) error {
	gc := pc.GitHubClient
	log := pc.Logger
	if err := handleReopen(gc, log, &e, !deprecatedWarn); err != nil {
		return err
	}
	if err := handleClose(gc, log, &e, !deprecatedWarn); err != nil {
		return err
	}
	return handle(gc, log, &e)
}

func handle(gc lifecycleClient, log *logrus.Entry, e *github.GenericCommentEvent) error {
	// Only consider new comments.
	if e.Action != github.GenericCommentActionCreated {
		return nil
	}

	for _, mat := range lifecycleRe.FindAllStringSubmatch(e.Body, -1) {
		if err := handleOne(gc, log, e, mat); err != nil {
			return err
		}
	}
	return nil
}

func handleOne(gc lifecycleClient, log *logrus.Entry, e *github.GenericCommentEvent, mat []string) error {
	org := e.Repo.Owner.Login
	repo := e.Repo.Name
	number := e.Number

	remove := mat[1] != ""
	cmd := mat[2]
	lbl := "lifecycle/" + cmd
	// Let's start simple and allow anyone to add/remove frozen, stale, putrid, rotten labels.
	// Adjust if we find evidence of the community abusing these labels.
	if remove {
		return gc.RemoveLabel(org, repo, number, lbl)
	}
	return gc.AddLabel(org, repo, number, lbl)
}
<|endoftext|>


<commit_before>//   Copyright 2017 MSolution.IO
//
//   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 s3

import (
	"context"
	"database/sql"
	"strings"
	"sync"
	"time"

	"github.com/satori/go.uuid"
	"github.com/trackit/jsonlog"
	"gopkg.in/olivere/elastic.v5"

	"github.com/trackit/trackit2/aws"
	"github.com/trackit/trackit2/es"
)

const (
	kibibyte = 1 << 10
	mebibyte = 1 << 20
	gibibyte = 1 << 30

	esBulkInsertSize    = 8 * mebibyte
	esBulkInsertWorkers = 4

	opTypeIndex  = "index"
	opTypeCreate = "create"

	tagPrefix = `resourceTags/user:`
)

type ReportUpdateConclusion struct {
	BillRepository       BillRepository
	LastImportedManifest time.Time
	Error                error
}

func reportUpdateConclusionChanToSlice(rucc <-chan ReportUpdateConclusion, count int) (rucs []ReportUpdateConclusion) {
	rucs = make([]ReportUpdateConclusion, count)
	for i := range rucs {
		if r, ok := <-rucc; ok {
			rucs[i] = r
		} else {
			rucs = rucs[:i]
			return
		}
	}
	return
}

func UpdateDueReports(ctx context.Context, tx *sql.Tx) ([]ReportUpdateConclusion, error) {
	var wg sync.WaitGroup
	aas := make(map[int]aws.AwsAccount)
	brs, err := GetAwsBillRepositoriesWithDueUpdate(tx)
	if err != nil {
		return nil, err
	}
	wg.Add(len(brs))
	conclusionChan := make(chan ReportUpdateConclusion, len(brs))
	defer close(conclusionChan)
	for _, br := range brs {
		var aa aws.AwsAccount
		var ok bool
		var err error
		if aa, ok = aas[br.AwsAccountId]; !ok {
			aa, err = aws.GetAwsAccountWithId(br.AwsAccountId, tx)
			if err != nil {
				return nil, err
			}
			aas[br.AwsAccountId] = aa
		}
		go func(ctx context.Context, aa aws.AwsAccount, br BillRepository) {
			lim, err := UpdateReport(ctx, aa, br)
			conclusionChan <- ReportUpdateConclusion{
				BillRepository:       br,
				LastImportedManifest: lim,
				Error:                err,
			}
			wg.Done()
		}(ctx, aa, br)
	}
	wg.Wait()
	return reportUpdateConclusionChanToSlice(conclusionChan, len(brs)), nil
}

// contextKey is a key in a context, to prevent collision with other modules.
type contextKey uint

// ingestionContextKey is used to store an 'ingestionId' in a context.
const ingestionContextKey = contextKey(iota)

// contextWithIngestionId returns a context configured so that its logger logs
// an 'ingestionId'.
func contextWithIngestionId(ctx context.Context) context.Context {
	ingestionId := uuid.NewV1().String()
	ctx = context.WithValue(ctx, ingestionContextKey, ingestionId)
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	logger = logger.WithContextKey(ingestionContextKey, "ingestionId")
	logger = logger.WithContext(ctx)
	return jsonlog.ContextWithLogger(ctx, logger)
}

// UpdateReport updates the elasticsearch database with new data from usage and
// cost reports.
func UpdateReport(ctx context.Context, aa aws.AwsAccount, br BillRepository) (latestManifest time.Time, err error) {
	ctx = contextWithIngestionId(ctx)
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	logger.Info("Updating reports for AWS account.", map[string]interface{}{
		"awsAccount":     aa,
		"billRepository": br,
	})
	if bp, err := getBulkProcessor(ctx); err != nil {
		logger.Error("Failed to get bulk processor.", err.Error())
		return latestManifest, err
	} else {
		index := es.IndexNameForUserId(aa.UserId, IndexPrefixLineItem)
		latestManifest, err = ReadBills(
			ctx,
			aa,
			br,
			ingestLineItems(ctx, bp, index),
			manifestsModifiedAfter(br.LastImportedManifest),
		)
	}
	return
}

// getBulkProcessor builds a bulk processor for ElasticSearch.
func getBulkProcessor(ctx context.Context) (*elastic.BulkProcessor, error) {
	bps := elastic.NewBulkProcessorService(es.Client)
	bps = bps.BulkActions(-1)
	bps = bps.BulkSize(esBulkInsertSize)
	bps = bps.Workers(esBulkInsertWorkers)
	bps = bps.Before(beforeBulk(ctx))
	bps = bps.After(afterBulk(ctx))
	return bps.Do(context.Background()) // use of background context is not an error
}

// ingestLineItems returns an OnLineItem handler which ingests LineItems in an
// ElasticSearch index.
func ingestLineItems(ctx context.Context, bp *elastic.BulkProcessor, index string) OnLineItem {
	return func(li LineItem, ok bool) {
		if ok {
			li = extractTags(li)
			rq := elastic.NewBulkIndexRequest()
			rq = rq.Index(index)
			rq = rq.OpType(opTypeCreate)
			rq = rq.Type(TypeLineItem)
			rq = rq.Id(li.EsId())
			rq = rq.Doc(li)
			bp.Add(rq)
		} else {
			bp.Flush()
			bp.Close()
		}
	}
}

// manifestsStartingAfter returns a manifest predicate which is true for all
// manifests starting after a given date.
func manifestsModifiedAfter(t time.Time) ManifestPredicate {
	return func(m manifest) bool {
		if time.Time(m.LastModified).After(t) {
			return true
		} else {
			return false
		}
	}
}

// extractTags extracts tags from a LineItem's Any field. It retrieves user
// tags only and stores them in the Tags map with a clean key.
func extractTags(li LineItem) LineItem {
	tags := make(map[string]string)
	for k, v := range li.Any {
		if strings.HasPrefix(k, tagPrefix) {
			tags[strings.TrimPrefix(k, tagPrefix)] = v
		}
	}
	li.Tags = tags
	li.Any = nil
	return li
}

func beforeBulk(ctx context.Context) func(int64, []elastic.BulkableRequest) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	return func(execId int64, reqs []elastic.BulkableRequest) {
		logger.Info("Performing bulk ElasticSearch requests.", map[string]interface{}{
			"executionId":   execId,
			"requestsCount": len(reqs),
		})
	}
}

func afterBulk(ctx context.Context) func(int64, []elastic.BulkableRequest, *elastic.BulkResponse, error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	return func(execId int64, reqs []elastic.BulkableRequest, resp *elastic.BulkResponse, err error) {
		if err != nil {
			logger.Error("Failed bulk ElasticSearch requests.", map[string]interface{}{
				"executionId": execId,
				"error":       err.Error(),
				"took":        resp.Took,
			})
		} else {
			logger.Info("Finished bulk ElasticSearch requests.", map[string]interface{}{
				"executionId": execId,
				"took":        resp.Took,
			})
		}

	}
}
<commit_msg>aws/s3: log when done ingesting data<commit_after>//   Copyright 2017 MSolution.IO
//
//   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 s3

import (
	"context"
	"database/sql"
	"strings"
	"sync"
	"time"

	"github.com/satori/go.uuid"
	"github.com/trackit/jsonlog"
	"gopkg.in/olivere/elastic.v5"

	"github.com/trackit/trackit2/aws"
	"github.com/trackit/trackit2/es"
)

const (
	kibibyte = 1 << 10
	mebibyte = 1 << 20
	gibibyte = 1 << 30

	esBulkInsertSize    = 8 * mebibyte
	esBulkInsertWorkers = 4

	opTypeIndex  = "index"
	opTypeCreate = "create"

	tagPrefix = `resourceTags/user:`
)

type ReportUpdateConclusion struct {
	BillRepository       BillRepository
	LastImportedManifest time.Time
	Error                error
}

func reportUpdateConclusionChanToSlice(rucc <-chan ReportUpdateConclusion, count int) (rucs []ReportUpdateConclusion) {
	rucs = make([]ReportUpdateConclusion, count)
	for i := range rucs {
		if r, ok := <-rucc; ok {
			rucs[i] = r
		} else {
			rucs = rucs[:i]
			return
		}
	}
	return
}

func UpdateDueReports(ctx context.Context, tx *sql.Tx) ([]ReportUpdateConclusion, error) {
	var wg sync.WaitGroup
	aas := make(map[int]aws.AwsAccount)
	brs, err := GetAwsBillRepositoriesWithDueUpdate(tx)
	if err != nil {
		return nil, err
	}
	wg.Add(len(brs))
	conclusionChan := make(chan ReportUpdateConclusion, len(brs))
	defer close(conclusionChan)
	for _, br := range brs {
		var aa aws.AwsAccount
		var ok bool
		var err error
		if aa, ok = aas[br.AwsAccountId]; !ok {
			aa, err = aws.GetAwsAccountWithId(br.AwsAccountId, tx)
			if err != nil {
				return nil, err
			}
			aas[br.AwsAccountId] = aa
		}
		go func(ctx context.Context, aa aws.AwsAccount, br BillRepository) {
			lim, err := UpdateReport(ctx, aa, br)
			conclusionChan <- ReportUpdateConclusion{
				BillRepository:       br,
				LastImportedManifest: lim,
				Error:                err,
			}
			wg.Done()
		}(ctx, aa, br)
	}
	wg.Wait()
	return reportUpdateConclusionChanToSlice(conclusionChan, len(brs)), nil
}

// contextKey is a key in a context, to prevent collision with other modules.
type contextKey uint

// ingestionContextKey is used to store an 'ingestionId' in a context.
const ingestionContextKey = contextKey(iota)

// contextWithIngestionId returns a context configured so that its logger logs
// an 'ingestionId'.
func contextWithIngestionId(ctx context.Context) context.Context {
	ingestionId := uuid.NewV1().String()
	ctx = context.WithValue(ctx, ingestionContextKey, ingestionId)
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	logger = logger.WithContextKey(ingestionContextKey, "ingestionId")
	logger = logger.WithContext(ctx)
	return jsonlog.ContextWithLogger(ctx, logger)
}

// UpdateReport updates the elasticsearch database with new data from usage and
// cost reports.
func UpdateReport(ctx context.Context, aa aws.AwsAccount, br BillRepository) (latestManifest time.Time, err error) {
	ctx = contextWithIngestionId(ctx)
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	logger.Info("Updating reports for AWS account.", map[string]interface{}{
		"awsAccount":     aa,
		"billRepository": br,
	})
	if bp, err := getBulkProcessor(ctx); err != nil {
		logger.Error("Failed to get bulk processor.", err.Error())
		return latestManifest, err
	} else {
		index := es.IndexNameForUserId(aa.UserId, IndexPrefixLineItem)
		latestManifest, err = ReadBills(
			ctx,
			aa,
			br,
			ingestLineItems(ctx, bp, index),
			manifestsModifiedAfter(br.LastImportedManifest),
		)
	}
	logger.Info("Done ingesting data.", nil)
	return
}

// getBulkProcessor builds a bulk processor for ElasticSearch.
func getBulkProcessor(ctx context.Context) (*elastic.BulkProcessor, error) {
	bps := elastic.NewBulkProcessorService(es.Client)
	bps = bps.BulkActions(-1)
	bps = bps.BulkSize(esBulkInsertSize)
	bps = bps.Workers(esBulkInsertWorkers)
	bps = bps.Before(beforeBulk(ctx))
	bps = bps.After(afterBulk(ctx))
	return bps.Do(context.Background()) // use of background context is not an error
}

// ingestLineItems returns an OnLineItem handler which ingests LineItems in an
// ElasticSearch index.
func ingestLineItems(ctx context.Context, bp *elastic.BulkProcessor, index string) OnLineItem {
	return func(li LineItem, ok bool) {
		if ok {
			li = extractTags(li)
			rq := elastic.NewBulkIndexRequest()
			rq = rq.Index(index)
			rq = rq.OpType(opTypeCreate)
			rq = rq.Type(TypeLineItem)
			rq = rq.Id(li.EsId())
			rq = rq.Doc(li)
			bp.Add(rq)
		} else {
			bp.Flush()
			bp.Close()
		}
	}
}

// manifestsStartingAfter returns a manifest predicate which is true for all
// manifests starting after a given date.
func manifestsModifiedAfter(t time.Time) ManifestPredicate {
	return func(m manifest) bool {
		if time.Time(m.LastModified).After(t) {
			return true
		} else {
			return false
		}
	}
}

// extractTags extracts tags from a LineItem's Any field. It retrieves user
// tags only and stores them in the Tags map with a clean key.
func extractTags(li LineItem) LineItem {
	tags := make(map[string]string)
	for k, v := range li.Any {
		if strings.HasPrefix(k, tagPrefix) {
			tags[strings.TrimPrefix(k, tagPrefix)] = v
		}
	}
	li.Tags = tags
	li.Any = nil
	return li
}

func beforeBulk(ctx context.Context) func(int64, []elastic.BulkableRequest) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	return func(execId int64, reqs []elastic.BulkableRequest) {
		logger.Info("Performing bulk ElasticSearch requests.", map[string]interface{}{
			"executionId":   execId,
			"requestsCount": len(reqs),
		})
	}
}

func afterBulk(ctx context.Context) func(int64, []elastic.BulkableRequest, *elastic.BulkResponse, error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	return func(execId int64, reqs []elastic.BulkableRequest, resp *elastic.BulkResponse, err error) {
		if err != nil {
			logger.Error("Failed bulk ElasticSearch requests.", map[string]interface{}{
				"executionId": execId,
				"error":       err.Error(),
				"took":        resp.Took,
			})
		} else {
			logger.Info("Finished bulk ElasticSearch requests.", map[string]interface{}{
				"executionId": execId,
				"took":        resp.Took,
			})
		}

	}
}
<|endoftext|>


<commit_before>package boilingcore

import (
	"path/filepath"
	"strings"

	"github.com/spf13/cast"
	"github.com/volatiletech/sqlboiler/drivers"
	"github.com/volatiletech/sqlboiler/importers"
)

// Config for the running of the commands
type Config struct {
	DriverName   string         `toml:"driver_name,omitempty" json:"driver_name,omitempty"`
	DriverConfig drivers.Config `toml:"driver_config,omitempty" json:"driver_config,omitempty"`

	PkgName          string   `toml:"pkg_name,omitempty" json:"pkg_name,omitempty"`
	OutFolder        string   `toml:"out_folder,omitempty" json:"out_folder,omitempty"`
	TemplateDirs     []string `toml:"template_dirs,omitempty" json:"template_dirs,omitempty"`
	Tags             []string `toml:"tags,omitempty" json:"tags,omitempty"`
	Replacements     []string `toml:"replacements,omitempty" json:"replacements,omitempty"`
	Debug            bool     `toml:"debug,omitempty" json:"debug,omitempty"`
	AddGlobal        bool     `toml:"add_global,omitempty" json:"add_global,omitempty"`
	AddPanic         bool     `toml:"add_panic,omitempty" json:"add_panic,omitempty"`
	NoContext        bool     `toml:"no_context,omitempty" json:"no_context,omitempty"`
	NoTests          bool     `toml:"no_tests,omitempty" json:"no_tests,omitempty"`
	NoHooks          bool     `toml:"no_hooks,omitempty" json:"no_hooks,omitempty"`
	NoAutoTimestamps bool     `toml:"no_auto_timestamps,omitempty" json:"no_auto_timestamps,omitempty"`
	NoRowsAffected   bool     `toml:"no_rows_affected,omitempty" json:"no_rows_affected,omitempty"`
	Wipe             bool     `toml:"wipe,omitempty" json:"wipe,omitempty"`
	StructTagCasing  string   `toml:"struct_tag_casing,omitempty" json:"struct_tag_casing,omitempty"`

	Imports importers.Collection `toml:"imports,omitempty" json:"imports,omitempty"`

	Aliases      Aliases       `toml:"aliases,omitempty" json:"aliases,omitempty"`
	TypeReplaces []TypeReplace `toml:"type_replaces,omitempty" json:"type_replaces,omitempty"`
}

// TypeReplace replaces a column type with something else
type TypeReplace struct {
	Match   drivers.Column `toml:"match,omitempty" json:"match,omitempty"`
	Replace drivers.Column `toml:"replace,omitempty" json:"replace,omitempty"`
	Imports importers.Set  `toml:"imports,omitempty" json:"imports,omitempty"`
}

// OutputDirDepth returns depth of output directory
func (c *Config) OutputDirDepth() int {
	d := filepath.ToSlash(filepath.Clean(c.OutFolder))
	if d == "." {
		return 0
	}

	return strings.Count(d, "/") + 1
}

// ConvertAliases is necessary because viper
//
// It also supports two different syntaxes, because of viper:
//
//   [aliases.tables.table_name]
//   fields... = "values"
//     [aliases.tables.columns]
//     colname = "alias"
//     [aliases.tables.relationships.fkey_name]
//     local   = "x"
//     foreign = "y"
//
// Or alternatively (when toml key names or viper's
// lowercasing of key names gets in the way):
//
//   [[aliases.tables]]
//   name = "table_name"
//   fields... = "values"
//     [[aliases.tables.columns]]
//     name  = "colname"
//     alias = "alias"
//     [[aliases.tables.relationships]]
//     name    = "fkey_name"
//     local   = "x"
//     foreign = "y"
func ConvertAliases(i interface{}) (a Aliases) {
	if i == nil {
		return a
	}

	topLevel := cast.ToStringMap(i)

	tablesIntf := topLevel["tables"]

	iterateMapOrSlice(tablesIntf, func(name string, tIntf interface{}) {
		if a.Tables == nil {
			a.Tables = make(map[string]TableAlias)
		}

		t := cast.ToStringMap(tIntf)

		var ta TableAlias

		if s := t["up_plural"]; s != nil {
			ta.UpPlural = s.(string)
		}
		if s := t["up_singular"]; s != nil {
			ta.UpSingular = s.(string)
		}
		if s := t["down_plural"]; s != nil {
			ta.DownPlural = s.(string)
		}
		if s := t["down_singular"]; s != nil {
			ta.DownSingular = s.(string)
		}

		if colsIntf, ok := t["columns"]; ok {
			ta.Columns = make(map[string]string)

			iterateMapOrSlice(colsIntf, func(name string, colIntf interface{}) {
				var alias string
				switch col := colIntf.(type) {
				case map[string]interface{}, map[interface{}]interface{}:
					cmap := cast.ToStringMap(colIntf)
					alias = cmap["alias"].(string)
				case string:
					alias = col
				}
				ta.Columns[name] = alias
			})
		}

		relationshipsIntf, ok := t["relationships"]
		if ok {
			iterateMapOrSlice(relationshipsIntf, func(name string, rIntf interface{}) {
				if ta.Relationships == nil {
					ta.Relationships = make(map[string]RelationshipAlias)
				}

				var ra RelationshipAlias
				rel := cast.ToStringMap(rIntf)

				if s := rel["local"]; s != nil {
					ra.Local = s.(string)
				}
				if s := rel["foreign"]; s != nil {
					ra.Foreign = s.(string)
				}

				ta.Relationships[name] = ra
			})
		}

		a.Tables[name] = ta
	})

	return a
}

func iterateMapOrSlice(mapOrSlice interface{}, fn func(name string, obj interface{})) {
	switch t := mapOrSlice.(type) {
	case map[string]interface{}, map[interface{}]interface{}:
		tmap := cast.ToStringMap(mapOrSlice)
		for name, table := range tmap {
			fn(name, table)
		}
	case []interface{}:
		for _, intf := range t {
			obj := cast.ToStringMap(intf)
			name := obj["name"].(string)
			fn(name, intf)
		}
	}
}

// ConvertTypeReplace is necessary because viper
func ConvertTypeReplace(i interface{}) []TypeReplace {
	if i == nil {
		return nil
	}

	intfArray := i.([]interface{})
	var replaces []TypeReplace
	for _, r := range intfArray {
		replaceIntf := cast.ToStringMap(r)
		replace := TypeReplace{}

		if replaceIntf["match"] == nil || replaceIntf["replace"] == nil {
			panic("replace types must specify both match and replace")
		}

		replace.Match = columnFromInterface(replaceIntf["match"])
		replace.Replace = columnFromInterface(replaceIntf["replace"])

		if imps := replaceIntf["imports"]; imps != nil {
			var err error
			replace.Imports, err = importers.SetFromInterface(imps)
			if err != nil {
				panic(err)
			}
		}

		replaces = append(replaces, replace)
	}

	return replaces
}

func columnFromInterface(i interface{}) (col drivers.Column) {
	m := cast.ToStringMap(i)
	if s := m["name"]; s != nil {
		col.Name = s.(string)
	}
	if s := m["type"]; s != nil {
		col.Type = s.(string)
	}
	if s := m["db_type"]; s != nil {
		col.DBType = s.(string)
	}
	if s := m["udt_name"]; s != nil {
		col.UDTName = s.(string)
	}
	if s := m["full_db_type"]; s != nil {
		col.FullDBType = s.(string)
	}
	if s := m["arr_type"]; s != nil {
		col.ArrType = new(string)
		*col.ArrType = s.(string)
	}
	if b := m["auto_generated"]; b != nil {
		col.AutoGenerated = b.(bool)
	}
	if b := m["nullable"]; b != nil {
		col.Nullable = b.(bool)
	}

	return col
}
<commit_msg>Fix yaml types replace conversion error<commit_after>package boilingcore

import (
	"path/filepath"
	"strings"

	"github.com/spf13/cast"
	"github.com/volatiletech/sqlboiler/drivers"
	"github.com/volatiletech/sqlboiler/importers"
)

// Config for the running of the commands
type Config struct {
	DriverName   string         `toml:"driver_name,omitempty" json:"driver_name,omitempty"`
	DriverConfig drivers.Config `toml:"driver_config,omitempty" json:"driver_config,omitempty"`

	PkgName          string   `toml:"pkg_name,omitempty" json:"pkg_name,omitempty"`
	OutFolder        string   `toml:"out_folder,omitempty" json:"out_folder,omitempty"`
	TemplateDirs     []string `toml:"template_dirs,omitempty" json:"template_dirs,omitempty"`
	Tags             []string `toml:"tags,omitempty" json:"tags,omitempty"`
	Replacements     []string `toml:"replacements,omitempty" json:"replacements,omitempty"`
	Debug            bool     `toml:"debug,omitempty" json:"debug,omitempty"`
	AddGlobal        bool     `toml:"add_global,omitempty" json:"add_global,omitempty"`
	AddPanic         bool     `toml:"add_panic,omitempty" json:"add_panic,omitempty"`
	NoContext        bool     `toml:"no_context,omitempty" json:"no_context,omitempty"`
	NoTests          bool     `toml:"no_tests,omitempty" json:"no_tests,omitempty"`
	NoHooks          bool     `toml:"no_hooks,omitempty" json:"no_hooks,omitempty"`
	NoAutoTimestamps bool     `toml:"no_auto_timestamps,omitempty" json:"no_auto_timestamps,omitempty"`
	NoRowsAffected   bool     `toml:"no_rows_affected,omitempty" json:"no_rows_affected,omitempty"`
	Wipe             bool     `toml:"wipe,omitempty" json:"wipe,omitempty"`
	StructTagCasing  string   `toml:"struct_tag_casing,omitempty" json:"struct_tag_casing,omitempty"`

	Imports importers.Collection `toml:"imports,omitempty" json:"imports,omitempty"`

	Aliases      Aliases       `toml:"aliases,omitempty" json:"aliases,omitempty"`
	TypeReplaces []TypeReplace `toml:"type_replaces,omitempty" json:"type_replaces,omitempty"`
}

// TypeReplace replaces a column type with something else
type TypeReplace struct {
	Match   drivers.Column `toml:"match,omitempty" json:"match,omitempty"`
	Replace drivers.Column `toml:"replace,omitempty" json:"replace,omitempty"`
	Imports importers.Set  `toml:"imports,omitempty" json:"imports,omitempty"`
}

// OutputDirDepth returns depth of output directory
func (c *Config) OutputDirDepth() int {
	d := filepath.ToSlash(filepath.Clean(c.OutFolder))
	if d == "." {
		return 0
	}

	return strings.Count(d, "/") + 1
}

// ConvertAliases is necessary because viper
//
// It also supports two different syntaxes, because of viper:
//
//   [aliases.tables.table_name]
//   fields... = "values"
//     [aliases.tables.columns]
//     colname = "alias"
//     [aliases.tables.relationships.fkey_name]
//     local   = "x"
//     foreign = "y"
//
// Or alternatively (when toml key names or viper's
// lowercasing of key names gets in the way):
//
//   [[aliases.tables]]
//   name = "table_name"
//   fields... = "values"
//     [[aliases.tables.columns]]
//     name  = "colname"
//     alias = "alias"
//     [[aliases.tables.relationships]]
//     name    = "fkey_name"
//     local   = "x"
//     foreign = "y"
func ConvertAliases(i interface{}) (a Aliases) {
	if i == nil {
		return a
	}

	topLevel := cast.ToStringMap(i)

	tablesIntf := topLevel["tables"]

	iterateMapOrSlice(tablesIntf, func(name string, tIntf interface{}) {
		if a.Tables == nil {
			a.Tables = make(map[string]TableAlias)
		}

		t := cast.ToStringMap(tIntf)

		var ta TableAlias

		if s := t["up_plural"]; s != nil {
			ta.UpPlural = s.(string)
		}
		if s := t["up_singular"]; s != nil {
			ta.UpSingular = s.(string)
		}
		if s := t["down_plural"]; s != nil {
			ta.DownPlural = s.(string)
		}
		if s := t["down_singular"]; s != nil {
			ta.DownSingular = s.(string)
		}

		if colsIntf, ok := t["columns"]; ok {
			ta.Columns = make(map[string]string)

			iterateMapOrSlice(colsIntf, func(name string, colIntf interface{}) {
				var alias string
				switch col := colIntf.(type) {
				case map[string]interface{}, map[interface{}]interface{}:
					cmap := cast.ToStringMap(colIntf)
					alias = cmap["alias"].(string)
				case string:
					alias = col
				}
				ta.Columns[name] = alias
			})
		}

		relationshipsIntf, ok := t["relationships"]
		if ok {
			iterateMapOrSlice(relationshipsIntf, func(name string, rIntf interface{}) {
				if ta.Relationships == nil {
					ta.Relationships = make(map[string]RelationshipAlias)
				}

				var ra RelationshipAlias
				rel := cast.ToStringMap(rIntf)

				if s := rel["local"]; s != nil {
					ra.Local = s.(string)
				}
				if s := rel["foreign"]; s != nil {
					ra.Foreign = s.(string)
				}

				ta.Relationships[name] = ra
			})
		}

		a.Tables[name] = ta
	})

	return a
}

func iterateMapOrSlice(mapOrSlice interface{}, fn func(name string, obj interface{})) {
	switch t := mapOrSlice.(type) {
	case map[string]interface{}, map[interface{}]interface{}:
		tmap := cast.ToStringMap(mapOrSlice)
		for name, table := range tmap {
			fn(name, table)
		}
	case []interface{}:
		for _, intf := range t {
			obj := cast.ToStringMap(intf)
			name := obj["name"].(string)
			fn(name, intf)
		}
	}
}

// ConvertTypeReplace is necessary because viper
func ConvertTypeReplace(i interface{}) []TypeReplace {
	if i == nil {
		return nil
	}

	intfArray := i.([]interface{})
	var replaces []TypeReplace
	for _, r := range intfArray {
		replaceIntf := cast.ToStringMap(r)
		replace := TypeReplace{}

		if replaceIntf["match"] == nil || replaceIntf["replace"] == nil {
			panic("replace types must specify both match and replace")
		}

		replace.Match = columnFromInterface(replaceIntf["match"])
		replace.Replace = columnFromInterface(replaceIntf["replace"])

		if imps := replaceIntf["imports"]; imps != nil {
			imps = cast.ToStringMap(imps)
			var err error
			replace.Imports, err = importers.SetFromInterface(imps)
			if err != nil {
				panic(err)
			}
		}

		replaces = append(replaces, replace)
	}

	return replaces
}

func columnFromInterface(i interface{}) (col drivers.Column) {
	m := cast.ToStringMap(i)
	if s := m["name"]; s != nil {
		col.Name = s.(string)
	}
	if s := m["type"]; s != nil {
		col.Type = s.(string)
	}
	if s := m["db_type"]; s != nil {
		col.DBType = s.(string)
	}
	if s := m["udt_name"]; s != nil {
		col.UDTName = s.(string)
	}
	if s := m["full_db_type"]; s != nil {
		col.FullDBType = s.(string)
	}
	if s := m["arr_type"]; s != nil {
		col.ArrType = new(string)
		*col.ArrType = s.(string)
	}
	if b := m["auto_generated"]; b != nil {
		col.AutoGenerated = b.(bool)
	}
	if b := m["nullable"]; b != nil {
		col.Nullable = b.(bool)
	}

	return col
}
<|endoftext|>


<commit_before>package booklitcmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"go/build"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/sirupsen/logrus"
	"github.com/vito/booklit"
	"github.com/vito/booklit/baselit"
	"github.com/vito/booklit/load"
	"github.com/vito/booklit/render"
)

type Command struct {
	Version func() `short:"v" long:"version" description:"Print the version of Boooklit and exit."`

	In  string `long:"in"  short:"i" required:"true" description:"Input .lit file."`
	Out string `long:"out" short:"o" required:"true" description:"Output directory in which to render."`

	ServerPort int `long:"serve" short:"s" description:"Start an HTTP server on the given port."`

	Plugins []string `long:"plugin" short:"p" description:"Package to import, providing a plugin."`

	Debug bool `long:"debug" short:"d" description:"Log at debug level."`

	AllowBrokenReferences bool `long:"allow-broken-references" description:"Replace broken references with a bogus tag."`

	HTMLEngine struct {
		Templates string `long:"templates" description:"Directory containing .tmpl files to load."`
	} `group:"HTML Rendering Engine" namespace:"html"`
}

func (cmd *Command) Execute(args []string) error {
	if cmd.Debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	isReexec := os.Getenv("BOOKLIT_REEXEC") != ""

	if cmd.ServerPort != 0 && !isReexec {
		return cmd.Serve()
	} else {
		paths, err := cmd.Build(isReexec)
		if isReexec {
			err = json.NewEncoder(os.Stdout).Encode(reexecOutput{
				Paths: paths,
			})
			if err != nil {
				return err
			}
		}

		return err
	}
}

func (cmd *Command) Serve() error {
	http.Handle("/", &Server{
		Command:    cmd,
		FileServer: http.FileServer(http.Dir(cmd.Out)),
	})

	return http.ListenAndServe(fmt.Sprintf(":%d", cmd.ServerPort), nil)
}

type reexecOutput struct {
	Paths []string
}

func (cmd *Command) Build(isReexec bool) ([]string, error) {
	if len(cmd.Plugins) > 0 && !isReexec {
		return cmd.reexec()
	}

	processor := &load.Processor{
		AllowBrokenReferences: cmd.AllowBrokenReferences,

		PluginFactories: []booklit.PluginFactory{
			baselit.NewPlugin,
		},
	}

	section, err := processor.LoadFile(cmd.In)
	if err != nil {
		return nil, err
	}

	err = os.MkdirAll(cmd.Out, 0755)
	if err != nil {
		return nil, err
	}

	engine := render.NewHTMLRenderingEngine()

	if cmd.HTMLEngine.Templates != "" {
		err := engine.LoadTemplates(cmd.HTMLEngine.Templates)
		if err != nil {
			return nil, err
		}
	}

	writer := render.Writer{
		Engine:      engine,
		Destination: cmd.Out,
	}

	err = writer.WriteSection(section)
	if err != nil {
		return nil, err
	}

	return cmd.pathsToWatch(section)
}

func (cmd *Command) pathsToWatch(section *booklit.Section) ([]string, error) {
	paths := cmd.sectionPaths(section)

	for _, plug := range cmd.Plugins {
		pkg, err := build.Import(plug, ".", 0)
		if err != nil {
			return nil, err
		}

		for _, file := range pkg.GoFiles {
			paths = append(paths, filepath.Join(pkg.Dir, file))
		}

		paths = append(paths, pkg.Dir)
	}

	templatesDir := cmd.HTMLEngine.Templates
	if templatesDir != "" {
		files, err := filepath.Glob(filepath.Join(templatesDir, "*.tmpl"))
		if err != nil {
			return nil, err
		}

		paths = append(paths, files...)

		paths = append(paths, templatesDir)
	}

	return paths, nil
}

func (cmd *Command) sectionPaths(section *booklit.Section) []string {
	pathsUniq := map[string]struct{}{section.Path: struct{}{}}

	for _, child := range section.Children {
		for _, path := range cmd.sectionPaths(child) {
			pathsUniq[path] = struct{}{}
		}
	}

	paths := []string{}
	for path, _ := range pathsUniq {
		paths = append(paths, path)
	}

	return paths
}

func (cmd *Command) reexec() ([]string, error) {
	tmpdir, err := ioutil.TempDir("", "booklit-reexec")
	if err != nil {
		return nil, err
	}

	defer func() {
		_ = os.RemoveAll(tmpdir)
	}()

	src := filepath.Join(tmpdir, "main.go")
	bin := filepath.Join(tmpdir, "booklit")

	goSrc := "package main\n"
	goSrc += "import \"github.com/vito/booklit/booklitcmd\"\n"
	for _, p := range cmd.Plugins {
		goSrc += "import _ \"" + p + "\"\n"
	}
	goSrc += "func main() {\n"
	goSrc += "	booklitcmd.Main()\n"
	goSrc += "}\n"

	err = ioutil.WriteFile(src, []byte(goSrc), 0644)
	if err != nil {
		return nil, err
	}

	build := exec.Command("go", "build", "-o", bin, src)
	build.Stdout = os.Stdout
	build.Stderr = os.Stderr
	err = build.Run()
	if err != nil {
		return nil, err
	}

	buf := new(bytes.Buffer)
	run := exec.Command(bin, os.Args[1:]...)
	run.Env = append(os.Environ(), "BOOKLIT_REEXEC=1")
	run.Stdout = buf
	run.Stderr = os.Stderr
	err = run.Run()
	if err != nil {
		return nil, err
	}

	var res reexecOutput
	err = json.Unmarshal(buf.Bytes(), &res)
	if err != nil {
		return nil, err
	}

	return res.Paths, nil
}
<commit_msg>fix err swallowing when reexecing<commit_after>package booklitcmd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"go/build"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/sirupsen/logrus"
	"github.com/vito/booklit"
	"github.com/vito/booklit/baselit"
	"github.com/vito/booklit/load"
	"github.com/vito/booklit/render"
)

type Command struct {
	Version func() `short:"v" long:"version" description:"Print the version of Boooklit and exit."`

	In  string `long:"in"  short:"i" required:"true" description:"Input .lit file."`
	Out string `long:"out" short:"o" required:"true" description:"Output directory in which to render."`

	ServerPort int `long:"serve" short:"s" description:"Start an HTTP server on the given port."`

	Plugins []string `long:"plugin" short:"p" description:"Package to import, providing a plugin."`

	Debug bool `long:"debug" short:"d" description:"Log at debug level."`

	AllowBrokenReferences bool `long:"allow-broken-references" description:"Replace broken references with a bogus tag."`

	HTMLEngine struct {
		Templates string `long:"templates" description:"Directory containing .tmpl files to load."`
	} `group:"HTML Rendering Engine" namespace:"html"`
}

func (cmd *Command) Execute(args []string) error {
	if cmd.Debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	isReexec := os.Getenv("BOOKLIT_REEXEC") != ""

	if cmd.ServerPort != 0 && !isReexec {
		return cmd.Serve()
	} else {
		paths, err := cmd.Build(isReexec)
		if err != nil {
			return err
		}

		if isReexec {
			err := json.NewEncoder(os.Stdout).Encode(reexecOutput{
				Paths: paths,
			})
			if err != nil {
				return err
			}
		}

		return nil
	}
}

func (cmd *Command) Serve() error {
	http.Handle("/", &Server{
		Command:    cmd,
		FileServer: http.FileServer(http.Dir(cmd.Out)),
	})

	return http.ListenAndServe(fmt.Sprintf(":%d", cmd.ServerPort), nil)
}

type reexecOutput struct {
	Paths []string
}

func (cmd *Command) Build(isReexec bool) ([]string, error) {
	if len(cmd.Plugins) > 0 && !isReexec {
		return cmd.reexec()
	}

	processor := &load.Processor{
		AllowBrokenReferences: cmd.AllowBrokenReferences,

		PluginFactories: []booklit.PluginFactory{
			baselit.NewPlugin,
		},
	}

	section, err := processor.LoadFile(cmd.In)
	if err != nil {
		return nil, err
	}

	err = os.MkdirAll(cmd.Out, 0755)
	if err != nil {
		return nil, err
	}

	engine := render.NewHTMLRenderingEngine()

	if cmd.HTMLEngine.Templates != "" {
		err := engine.LoadTemplates(cmd.HTMLEngine.Templates)
		if err != nil {
			return nil, err
		}
	}

	writer := render.Writer{
		Engine:      engine,
		Destination: cmd.Out,
	}

	err = writer.WriteSection(section)
	if err != nil {
		return nil, err
	}

	return cmd.pathsToWatch(section)
}

func (cmd *Command) pathsToWatch(section *booklit.Section) ([]string, error) {
	paths := cmd.sectionPaths(section)

	for _, plug := range cmd.Plugins {
		pkg, err := build.Import(plug, ".", 0)
		if err != nil {
			return nil, err
		}

		for _, file := range pkg.GoFiles {
			paths = append(paths, filepath.Join(pkg.Dir, file))
		}

		paths = append(paths, pkg.Dir)
	}

	templatesDir := cmd.HTMLEngine.Templates
	if templatesDir != "" {
		files, err := filepath.Glob(filepath.Join(templatesDir, "*.tmpl"))
		if err != nil {
			return nil, err
		}

		paths = append(paths, files...)

		paths = append(paths, templatesDir)
	}

	return paths, nil
}

func (cmd *Command) sectionPaths(section *booklit.Section) []string {
	pathsUniq := map[string]struct{}{section.Path: struct{}{}}

	for _, child := range section.Children {
		for _, path := range cmd.sectionPaths(child) {
			pathsUniq[path] = struct{}{}
		}
	}

	paths := []string{}
	for path, _ := range pathsUniq {
		paths = append(paths, path)
	}

	return paths
}

func (cmd *Command) reexec() ([]string, error) {
	tmpdir, err := ioutil.TempDir("", "booklit-reexec")
	if err != nil {
		return nil, err
	}

	defer func() {
		_ = os.RemoveAll(tmpdir)
	}()

	src := filepath.Join(tmpdir, "main.go")
	bin := filepath.Join(tmpdir, "booklit")

	goSrc := "package main\n"
	goSrc += "import \"github.com/vito/booklit/booklitcmd\"\n"
	for _, p := range cmd.Plugins {
		goSrc += "import _ \"" + p + "\"\n"
	}
	goSrc += "func main() {\n"
	goSrc += "	booklitcmd.Main()\n"
	goSrc += "}\n"

	err = ioutil.WriteFile(src, []byte(goSrc), 0644)
	if err != nil {
		return nil, err
	}

	build := exec.Command("go", "build", "-o", bin, src)
	build.Stdout = os.Stdout
	build.Stderr = os.Stderr
	err = build.Run()
	if err != nil {
		return nil, err
	}

	buf := new(bytes.Buffer)
	run := exec.Command(bin, os.Args[1:]...)
	run.Env = append(os.Environ(), "BOOKLIT_REEXEC=1")
	run.Stdout = buf
	run.Stderr = os.Stderr
	err = run.Run()
	if err != nil {
		return nil, err
	}

	var res reexecOutput
	err = json.Unmarshal(buf.Bytes(), &res)
	if err != nil {
		return nil, err
	}

	return res.Paths, nil
}
<|endoftext|>


<commit_before>package core

import "github.com/coel-lang/coel/src/lib/systemt"

type functionType struct {
	signature Signature
	function  func(...*Thunk) Value
}

func (f functionType) call(args Arguments) Value {
	ts, err := f.signature.Bind(args)

	if err != nil {
		return err
	}

	return f.function(ts...)
}

// NewLazyFunction creates a function whose arguments are evaluated lazily.
func NewLazyFunction(s Signature, f func(...*Thunk) Value) *Thunk {
	return Normal(functionType{
		signature: s,
		function:  f,
	})
}

// NewStrictFunction creates a function whose arguments are evaluated strictly.
func NewStrictFunction(s Signature, f func(...*Thunk) Value) *Thunk {
	return NewLazyFunction(s, func(ts ...*Thunk) Value {
		for _, t := range ts {
			tt := t
			systemt.Daemonize(func() { tt.Eval() })
		}

		return f(ts...)
	})
}

// NewEffectFunction creates a effect function which returns an effect value.
func NewEffectFunction(s Signature, f func(...*Thunk) Value) *Thunk {
	return Normal(functionType{
		s,
		func(ts ...*Thunk) Value { return newEffect(Normal(f(ts...))) },
	})
}

func (f functionType) string() Value {
	return StringType("<function>")
}
<commit_msg>Refactor function.go<commit_after>package core

import "github.com/coel-lang/coel/src/lib/systemt"

type functionType struct {
	signature Signature
	function  func(...*Thunk) Value
}

func (f functionType) call(args Arguments) Value {
	ts, err := f.signature.Bind(args)

	if err != nil {
		return err
	}

	return f.function(ts...)
}

// NewLazyFunction creates a function whose arguments are evaluated lazily.
func NewLazyFunction(s Signature, f func(...*Thunk) Value) *Thunk {
	return Normal(functionType{s, f})
}

// NewStrictFunction creates a function whose arguments are evaluated strictly.
func NewStrictFunction(s Signature, f func(...*Thunk) Value) *Thunk {
	return NewLazyFunction(s, func(ts ...*Thunk) Value {
		for _, t := range ts {
			tt := t
			systemt.Daemonize(func() { tt.Eval() })
		}

		return f(ts...)
	})
}

// NewEffectFunction creates a effect function which returns an effect value.
func NewEffectFunction(s Signature, f func(...*Thunk) Value) *Thunk {
	return Normal(functionType{s, func(ts ...*Thunk) Value { return newEffect(Normal(f(ts...))) }})
}

func (f functionType) string() Value {
	return StringType("<function>")
}
<|endoftext|>


<commit_before>package bridge

import (
	"bytes"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync"
	"time"
	"github.com/viant/toolbox"
	"fmt"
	"path"
	"os"
	"encoding/json"
	"encoding/base64"
	"unicode/utf8"
	"unicode"
)

//HttpBridgeConfig represent http bridge config
type HttpBridgeEndpointConfig struct {
	Port           string
	ReadTimeoutMs  int
	WriteTimeoutMs int
	MaxHeaderBytes int
}

//HttpBridgeProxyRoute represent http proxy route
type HttpBridgeProxyRoute struct {
	Pattern          string
	TargetURL        *url.URL
	ResponseModifier func(*http.Response) error
	Listener         func(request *http.Request, response *http.Response)
}

//HttpBridgeProxyConfig represent proxy config
type HttpBridgeProxyConfig struct {
	MaxIdleConnections    int
	RequestTimeoutMs      int
	KeepAliveTimeMs       int
	TLSHandshakeTimeoutMs int
	BufferPoolSize        int
	BufferSize            int
}

//HttpBridgeConfig represents HttpBridgeConfig config
type HttpBridgeConfig struct {
	Endpoint *HttpBridgeEndpointConfig
	Proxy    *HttpBridgeProxyConfig
	Routes   []*HttpBridgeProxyRoute
}

//ProxyHandlerFactory proxy handler factory
type HttpBridgeProxyHandlerFactory func(proxyConfig *HttpBridgeProxyConfig, route *HttpBridgeProxyRoute) (http.Handler, error)

//HttpBridge represents http bridge
type HttpBridge struct {
	Config   *HttpBridgeConfig
	Server   *http.Server
	Handlers map[string]http.Handler
}

//ListenAndServe start http endpoint
func (r *HttpBridge) ListenAndServe() error {
	return r.Server.ListenAndServe()
}

//ListenAndServe start http endpoint on secure port
func (r *HttpBridge) ListenAndServeTLS(certFile, keyFile string) error {
	return r.Server.ListenAndServeTLS(certFile, keyFile)
}

//NewHttpBridge creates a new instance of NewHttpBridge
func NewHttpBridge(config *HttpBridgeConfig, factory HttpBridgeProxyHandlerFactory) (*HttpBridge, error) {
	mux := http.NewServeMux()
	var handlers = make(map[string]http.Handler)
	for _, route := range config.Routes {
		handler, err := factory(config.Proxy, route)
		if err != nil {
			return nil, err
		}
		mux.Handle(route.Pattern, handler)
		handlers[route.Pattern] = handler
	}
	server := &http.Server{
		Addr:           ":" + config.Endpoint.Port,
		Handler:        mux,
		ReadTimeout:    time.Millisecond * time.Duration(config.Endpoint.ReadTimeoutMs),
		WriteTimeout:   time.Millisecond * time.Duration(config.Endpoint.WriteTimeoutMs),
		MaxHeaderBytes: config.Endpoint.MaxHeaderBytes,
	}
	return &HttpBridge{
		Server:   server,
		Config:   config,
		Handlers: handlers,
	}, nil
}

//NewProxyHandler creates a new proxy handler
func NewProxyHandler(proxyConfig *HttpBridgeProxyConfig, route *HttpBridgeProxyRoute) (http.Handler, error) {
	roundTripper := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   time.Duration(proxyConfig.RequestTimeoutMs) * time.Millisecond,
			KeepAlive: time.Duration(proxyConfig.KeepAliveTimeMs) * time.Millisecond,
		}).Dial,
		TLSHandshakeTimeout: time.Duration(proxyConfig.TLSHandshakeTimeoutMs) * time.Millisecond,
		MaxIdleConnsPerHost: proxyConfig.MaxIdleConnections,
	}
	director := func(request *http.Request) {
		request.URL.Scheme = route.TargetURL.Scheme
		request.URL.Host = route.TargetURL.Host
	}
	reverseProxy := &httputil.ReverseProxy{
		Transport:      roundTripper,
		BufferPool:     toolbox.NewBytesBufferPool(proxyConfig.BufferPoolSize, proxyConfig.BufferSize),
		ModifyResponse: route.ResponseModifier,
		Director:       director,
	}
	return reverseProxy, nil
}

//HttpTrip represents recorded round trip.
type HttpTrip struct {
	responseWriter     http.ResponseWriter
	Request            *http.Request
	responseBody       *bytes.Buffer
	responseStatusCode int
}

func (w *HttpTrip) Response() *http.Response {
	return &http.Response{
		Request:    w.Request,
		StatusCode: w.responseStatusCode,
		Header:     w.responseWriter.Header(),
		Body:       ioutil.NopCloser(bytes.NewReader(w.responseBody.Bytes())),
	}
}

func (w *HttpTrip) Write(b []byte) (int, error) {
	w.responseBody.Write(b)
	return w.responseWriter.Write(b)
}

func (w *HttpTrip) Header() http.Header {
	return w.responseWriter.Header()
}

func (w *HttpTrip) WriteHeader(status int) {
	w.responseStatusCode = status
	w.responseWriter.WriteHeader(status)
}

func (w *HttpTrip) Flush() {
	if flusher, ok := w.responseWriter.(http.Flusher); ok {
		flusher.Flush()
	}
}

func (w *HttpTrip) CloseNotify() <-chan bool {
	if closer, ok := w.responseWriter.(http.CloseNotifier); ok {
		return closer.CloseNotify()
	}
	return make(chan bool, 1)
}

//ListeningTripHandler represents endpoint recording handler
type ListeningTripHandler struct {
	handler         http.Handler
	pool            httputil.BufferPool
	listener        func(request *http.Request, response *http.Response)
	roundTripsMutex *sync.RWMutex
}

func (h *ListeningTripHandler) Notify(roundTrip *HttpTrip) {
	if h.listener != nil {
		h.listener(roundTrip.Request, roundTrip.Response())
	}
}

//drainBody reads all of b to memory and then returns two equivalent (modified version from  httputil)
func (h ListeningTripHandler) drainBody(reader io.ReadCloser) (io.ReadCloser, io.ReadCloser, error) {
	if reader == http.NoBody {
		return http.NoBody, http.NoBody, nil
	}
	var buf = new(bytes.Buffer)
	toolbox.CopyWithBufferPool(reader, buf, h.pool)
	return ioutil.NopCloser(bytes.NewReader(buf.Bytes())), ioutil.NopCloser(bytes.NewReader(buf.Bytes())), nil
}

func (h ListeningTripHandler) ServeHTTP(responseWriter http.ResponseWriter, request *http.Request) {
	var err error
	var originalRequest = request.WithContext(request.Context())
	if request.ContentLength > 0 {
		request.Body, originalRequest.Body, err = h.drainBody(request.Body)
		if err != nil {
			log.Printf("Faled to serve request :%v due to %v\n", request, err)
			return
		}
	}
	var recordedRoundTrip = &HttpTrip{
		responseWriter: responseWriter,
		Request:        originalRequest,
		responseBody:   new(bytes.Buffer),
	}
	responseWriter = http.ResponseWriter(recordedRoundTrip)
	defer h.Notify(recordedRoundTrip)
	h.handler.ServeHTTP(responseWriter, request)
}

func NewListeningHandler(handler http.Handler, bufferPoolSize, bufferSize int, listener func(request *http.Request, response *http.Response)) *ListeningTripHandler {
	var result = &ListeningTripHandler{
		handler:         handler,
		listener:        listener,
		pool:            toolbox.NewBytesBufferPool(bufferPoolSize, bufferSize),
		roundTripsMutex: &sync.RWMutex{},
	}
	return result
}

func NewProxyRecordingHandler(proxyConfig *HttpBridgeProxyConfig, route *HttpBridgeProxyRoute) (http.Handler, error) {
	handler, err := NewProxyHandler(proxyConfig, route)
	if err != nil {
		return nil, err
	}
	response := NewListeningHandler(handler, proxyConfig.BufferPoolSize, proxyConfig.BufferSize, route.Listener)
	return response, nil
}

func AsListeningTripHandler(handler http.Handler) *ListeningTripHandler {
	if result, ok := handler.(*ListeningTripHandler); ok {
		return result
	}
	return nil
}

type HttpRequest struct {
	Method string `json:",omitempty"`
	URL    string `json:",omitempty"`
	Header http.Header `json:",omitempty"`
	Body   string `json:",omitempty"`
}

type HttpResponse struct {
	Code   int
	Header http.Header `json:",omitempty"`
	Body   string `json:",omitempty"`
}

func ReaderAsText(reader io.Reader) string {

	if reader == nil {
		return ""
	}
	body, err := ioutil.ReadAll(reader)
	if err != nil {
		log.Fatal(err)
	}
	if isBinary(body) {
		buf := new(bytes.Buffer)
		encoder := base64.NewEncoder(base64.StdEncoding, buf)
		encoder.Write(body)
		encoder.Close()
		return fmt.Sprintf("base64:%v", string(buf.Bytes()))

	} else {
		return fmt.Sprintf("text:%v", string(body))
	}
}

func isBinary(input []byte) bool {
	for i, w := 0, 0; i < len(input); i += w {
		runeValue, width := utf8.DecodeRune(input[i:])
		if unicode.IsControl(runeValue) {
			return true
		}
		w = width
	}
	return false
}

func writeData(filename string, source interface{}, printStrOut bool) error  {
	toolbox.FileExists(filename)
	os.Remove(filename)

	logfile, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		log.Fatalf("error opening file: %v, %v", err, filename)
	}
	defer logfile.Close()

	buf, err := json.MarshalIndent(source, "", "\t")
	if err != nil {
		return err
	}
	_, err = logfile.Write(buf)

	if printStrOut {
		fmt.Printf("%v: %v\n", filename, string(buf))
	}

	return err
}

//HttpFileRecorder returns http route listener that will record request response to the passed in directory
func HttpFileRecorder(directory string, printStdOut bool) func(request *http.Request, response *http.Response) {
	tripCounter := 0

	err := toolbox.CreateDirIfNotExist(directory)
	if err != nil {
		fmt.Printf("failed to create directory%v %v\n, ", err, directory)
	}

	return func(request *http.Request, response *http.Response) {

		var body string

		if request.Body != nil {
			body = ReaderAsText(request.Body)
		}

		httpRequest := &HttpRequest{
		Method: request.Method,
		URL:    request.URL.String(),
		Header: request.Header,
		Body:   body,
		}

		err = writeData(path.Join(directory, fmt.Sprintf("%T-%v.json", *httpRequest, tripCounter)), httpRequest, printStdOut)
		if err != nil {
			fmt.Printf("failed to write request %v %v\n, ", err, request)
		}

		body = ReaderAsText(response.Body)
		request.Body = nil
		httpResponse := &HttpResponse{
		Code:   response.StatusCode,
		Header: response.Header,
		Body:   body,
		}

		err = writeData(path.Join(directory, fmt.Sprintf("%T-%v.json", *httpResponse, tripCounter)), httpResponse, printStdOut)
		if err != nil {
			fmt.Printf("failed to write response %v %v\n, ", err, response)
		}

		tripCounter++
	}

}
<commit_msg>added optional removal<commit_after>package bridge

import (
	"bytes"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync"
	"time"
	"github.com/viant/toolbox"
	"fmt"
	"path"
	"os"
	"encoding/json"
	"encoding/base64"
	"unicode/utf8"
	"unicode"
)

//HttpBridgeConfig represent http bridge config
type HttpBridgeEndpointConfig struct {
	Port           string
	ReadTimeoutMs  int
	WriteTimeoutMs int
	MaxHeaderBytes int
}

//HttpBridgeProxyRoute represent http proxy route
type HttpBridgeProxyRoute struct {
	Pattern          string
	TargetURL        *url.URL
	ResponseModifier func(*http.Response) error
	Listener         func(request *http.Request, response *http.Response)
}

//HttpBridgeProxyConfig represent proxy config
type HttpBridgeProxyConfig struct {
	MaxIdleConnections    int
	RequestTimeoutMs      int
	KeepAliveTimeMs       int
	TLSHandshakeTimeoutMs int
	BufferPoolSize        int
	BufferSize            int
}

//HttpBridgeConfig represents HttpBridgeConfig config
type HttpBridgeConfig struct {
	Endpoint *HttpBridgeEndpointConfig
	Proxy    *HttpBridgeProxyConfig
	Routes   []*HttpBridgeProxyRoute
}

//ProxyHandlerFactory proxy handler factory
type HttpBridgeProxyHandlerFactory func(proxyConfig *HttpBridgeProxyConfig, route *HttpBridgeProxyRoute) (http.Handler, error)

//HttpBridge represents http bridge
type HttpBridge struct {
	Config   *HttpBridgeConfig
	Server   *http.Server
	Handlers map[string]http.Handler
}

//ListenAndServe start http endpoint
func (r *HttpBridge) ListenAndServe() error {
	return r.Server.ListenAndServe()
}

//ListenAndServe start http endpoint on secure port
func (r *HttpBridge) ListenAndServeTLS(certFile, keyFile string) error {
	return r.Server.ListenAndServeTLS(certFile, keyFile)
}

//NewHttpBridge creates a new instance of NewHttpBridge
func NewHttpBridge(config *HttpBridgeConfig, factory HttpBridgeProxyHandlerFactory) (*HttpBridge, error) {
	mux := http.NewServeMux()
	var handlers = make(map[string]http.Handler)
	for _, route := range config.Routes {
		handler, err := factory(config.Proxy, route)
		if err != nil {
			return nil, err
		}
		mux.Handle(route.Pattern, handler)
		handlers[route.Pattern] = handler
	}
	server := &http.Server{
		Addr:           ":" + config.Endpoint.Port,
		Handler:        mux,
		ReadTimeout:    time.Millisecond * time.Duration(config.Endpoint.ReadTimeoutMs),
		WriteTimeout:   time.Millisecond * time.Duration(config.Endpoint.WriteTimeoutMs),
		MaxHeaderBytes: config.Endpoint.MaxHeaderBytes,
	}
	return &HttpBridge{
		Server:   server,
		Config:   config,
		Handlers: handlers,
	}, nil
}

//NewProxyHandler creates a new proxy handler
func NewProxyHandler(proxyConfig *HttpBridgeProxyConfig, route *HttpBridgeProxyRoute) (http.Handler, error) {
	roundTripper := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   time.Duration(proxyConfig.RequestTimeoutMs) * time.Millisecond,
			KeepAlive: time.Duration(proxyConfig.KeepAliveTimeMs) * time.Millisecond,
		}).Dial,
		TLSHandshakeTimeout: time.Duration(proxyConfig.TLSHandshakeTimeoutMs) * time.Millisecond,
		MaxIdleConnsPerHost: proxyConfig.MaxIdleConnections,
	}
	director := func(request *http.Request) {
		request.URL.Scheme = route.TargetURL.Scheme
		request.URL.Host = route.TargetURL.Host
	}
	reverseProxy := &httputil.ReverseProxy{
		Transport:      roundTripper,
		BufferPool:     toolbox.NewBytesBufferPool(proxyConfig.BufferPoolSize, proxyConfig.BufferSize),
		ModifyResponse: route.ResponseModifier,
		Director:       director,
	}
	return reverseProxy, nil
}

//HTTPTrip represents recorded round trip.
type HttpTrip struct {
	responseWriter     http.ResponseWriter
	Request            *http.Request
	responseBody       *bytes.Buffer
	responseStatusCode int
}

func (w *HttpTrip) Response() *http.Response {
	return &http.Response{
		Request:    w.Request,
		StatusCode: w.responseStatusCode,
		Header:     w.responseWriter.Header(),
		Body:       ioutil.NopCloser(bytes.NewReader(w.responseBody.Bytes())),
	}
}

func (w *HttpTrip) Write(b []byte) (int, error) {
	w.responseBody.Write(b)
	return w.responseWriter.Write(b)
}

func (w *HttpTrip) Header() http.Header {
	return w.responseWriter.Header()
}

func (w *HttpTrip) WriteHeader(status int) {
	w.responseStatusCode = status
	w.responseWriter.WriteHeader(status)
}

func (w *HttpTrip) Flush() {
	if flusher, ok := w.responseWriter.(http.Flusher); ok {
		flusher.Flush()
	}
}

func (w *HttpTrip) CloseNotify() <-chan bool {
	if closer, ok := w.responseWriter.(http.CloseNotifier); ok {
		return closer.CloseNotify()
	}
	return make(chan bool, 1)
}

//ListeningTripHandler represents endpoint recording handler
type ListeningTripHandler struct {
	handler         http.Handler
	pool            httputil.BufferPool
	listener        func(request *http.Request, response *http.Response)
	roundTripsMutex *sync.RWMutex
}

func (h *ListeningTripHandler) Notify(roundTrip *HttpTrip) {
	if h.listener != nil {
		h.listener(roundTrip.Request, roundTrip.Response())
	}
}

//drainBody reads all of b to memory and then returns two equivalent (modified version from  httputil)
func (h ListeningTripHandler) drainBody(reader io.ReadCloser) (io.ReadCloser, io.ReadCloser, error) {
	if reader == http.NoBody {
		return http.NoBody, http.NoBody, nil
	}
	var buf = new(bytes.Buffer)
	toolbox.CopyWithBufferPool(reader, buf, h.pool)
	return ioutil.NopCloser(bytes.NewReader(buf.Bytes())), ioutil.NopCloser(bytes.NewReader(buf.Bytes())), nil
}

func (h ListeningTripHandler) ServeHTTP(responseWriter http.ResponseWriter, request *http.Request) {
	var err error
	var originalRequest = request.WithContext(request.Context())
	if request.ContentLength > 0 {
		request.Body, originalRequest.Body, err = h.drainBody(request.Body)
		if err != nil {
			log.Printf("Faled to serve request :%v due to %v\n", request, err)
			return
		}
	}
	var recordedRoundTrip = &HttpTrip{
		responseWriter: responseWriter,
		Request:        originalRequest,
		responseBody:   new(bytes.Buffer),
	}
	responseWriter = http.ResponseWriter(recordedRoundTrip)
	defer h.Notify(recordedRoundTrip)
	h.handler.ServeHTTP(responseWriter, request)
}

func NewListeningHandler(handler http.Handler, bufferPoolSize, bufferSize int, listener func(request *http.Request, response *http.Response)) *ListeningTripHandler {
	var result = &ListeningTripHandler{
		handler:         handler,
		listener:        listener,
		pool:            toolbox.NewBytesBufferPool(bufferPoolSize, bufferSize),
		roundTripsMutex: &sync.RWMutex{},
	}
	return result
}

func NewProxyRecordingHandler(proxyConfig *HttpBridgeProxyConfig, route *HttpBridgeProxyRoute) (http.Handler, error) {
	handler, err := NewProxyHandler(proxyConfig, route)
	if err != nil {
		return nil, err
	}
	response := NewListeningHandler(handler, proxyConfig.BufferPoolSize, proxyConfig.BufferSize, route.Listener)
	return response, nil
}

func AsListeningTripHandler(handler http.Handler) *ListeningTripHandler {
	if result, ok := handler.(*ListeningTripHandler); ok {
		return result
	}
	return nil
}

type HttpRequest struct {
	Method string `json:",omitempty"`
	URL    string `json:",omitempty"`
	Header http.Header `json:",omitempty"`
	Body   string `json:",omitempty"`
}

type HttpResponse struct {
	Code   int
	Header http.Header `json:",omitempty"`
	Body   string `json:",omitempty"`
}

func ReaderAsText(reader io.Reader) string {

	if reader == nil {
		return ""
	}
	body, err := ioutil.ReadAll(reader)
	if err != nil {
		log.Fatal(err)
	}
	if isBinary(body) {
		buf := new(bytes.Buffer)
		encoder := base64.NewEncoder(base64.StdEncoding, buf)
		encoder.Write(body)
		encoder.Close()
		return fmt.Sprintf("base64:%v", string(buf.Bytes()))

	} else {
		return fmt.Sprintf("text:%v", string(body))
	}
}

func isBinary(input []byte) bool {
	for i, w := 0, 0; i < len(input); i += w {
		runeValue, width := utf8.DecodeRune(input[i:])
		if unicode.IsControl(runeValue) {
			return true
		}
		w = width
	}
	return false
}

func writeData(filename string, source interface{}, printStrOut bool) error  {
	if toolbox.FileExists(filename) {
		os.Remove(filename)
	}

	logfile, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		log.Fatalf("error opening file: %v, %v", err, filename)
	}
	defer logfile.Close()

	buf, err := json.MarshalIndent(source, "", "\t")
	if err != nil {
		return err
	}
	_, err = logfile.Write(buf)

	if printStrOut {
		fmt.Printf("%v: %v\n", filename, string(buf))
	}

	return err
}

//HttpFileRecorder returns http route listener that will record request response to the passed in directory
func HttpFileRecorder(directory string, printStdOut bool) func(request *http.Request, response *http.Response) {
	tripCounter := 0

	err := toolbox.CreateDirIfNotExist(directory)
	if err != nil {
		fmt.Printf("failed to create directory%v %v\n, ", err, directory)
	}

	return func(request *http.Request, response *http.Response) {

		var body string

		if request.Body != nil {
			body = ReaderAsText(request.Body)
		}

		httpRequest := &HttpRequest{
		Method: request.Method,
		URL:    request.URL.String(),
		Header: request.Header,
		Body:   body,
		}

		err = writeData(path.Join(directory, fmt.Sprintf("%T-%v.json", *httpRequest, tripCounter)), httpRequest, printStdOut)
		if err != nil {
			fmt.Printf("failed to write request %v %v\n, ", err, request)
		}

		body = ReaderAsText(response.Body)
		request.Body = nil
		httpResponse := &HttpResponse{
		Code:   response.StatusCode,
		Header: response.Header,
		Body:   body,
		}

		err = writeData(path.Join(directory, fmt.Sprintf("%T-%v.json", *httpResponse, tripCounter)), httpResponse, printStdOut)
		if err != nil {
			fmt.Printf("failed to write response %v %v\n, ", err, response)
		}

		tripCounter++
	}

}
<|endoftext|>