<commit_before>package ext

import opentracing "github.com/opentracing/opentracing-go"

// These constants define common tag names recommended for better portability across
// tracing systems and languages/platforms.
//
// The tag names are defined as typed strings, so that in addition to the usual use
//
//     span.setTag(TagName, value)
//
// they also support value type validation via this additional syntax:
//
//    TagName.Set(span, value)
//
var (
	//////////////////////////////////////////////////////////////////////
	// SpanKind (client/server)
	//////////////////////////////////////////////////////////////////////

	// SpanKind hints at relationship between spans, e.g. client/server
	SpanKind = spanKindTag("span.kind")

	// SpanKindRPCClient marks a span representing the client-side of an RPC
	// or other remote call
	SpanKindRPCClient = SpanKindEnum("client")

	// SpanKindRPCServer marks a span representing the server-side of an RPC
	// or other remote call
	SpanKindRPCServer = SpanKindEnum("server")

	//////////////////////////////////////////////////////////////////////
	// Sampling hint
	//////////////////////////////////////////////////////////////////////

	// SamplingPriority determines the priority of sampling this Span.
	SamplingPriority = uint16Tag("sampling.priority")

	//////////////////////////////////////////////////////////////////////
	// Peer tags. These tags can be emitted by either client-side of
	// server-side to describe the other side/service in a peer-to-peer
	// communications, like an RPC call.
	//////////////////////////////////////////////////////////////////////

	// PeerService records the service name of the peer
	PeerService = stringTag("peer.service")

	// PeerHostname records the host name of the peer
	PeerHostname = stringTag("peer.hostname")

	// PeerHostIPv4 records IP v4 host address of the peer
	PeerHostIPv4 = uint32Tag("peer.ipv4")

	// PeerHostIPv6 records IP v6 host address of the peer
	PeerHostIPv6 = stringTag("peer.ipv6")

	// PeerPort records port number of the peer
	PeerPort = uint16Tag("peer.port")

	//////////////////////////////////////////////////////////////////////
	// HTTP Tags
	//////////////////////////////////////////////////////////////////////

	// HTTPUrl should be the URL of the request being handled in this segment
	// of the trace, in standard URI format. The protocol is optional.
	HTTPUrl = stringTag("http.url")

	// HTTPMethod is the HTTP method of the request, and is case-insensitive.
	HTTPMethod = stringTag("http.method")

	// HTTPStatusCode is the numeric HTTP status code (200, 404, etc) of the
	// HTTP response.
	HTTPStatusCode = uint16Tag("htttp.status_code")
)

// ---

// SpanKindEnum represents common span types
type SpanKindEnum string

type spanKindTag string

// Add adds a string tag to the `span`
func (tag spanKindTag) Set(span opentracing.Span, value SpanKindEnum) {
	span.SetTag(string(tag), value)
}

// ---

type stringTag string

// Add adds a string tag to the `span`
func (tag stringTag) Set(span opentracing.Span, value string) {
	span.SetTag(string(tag), value)
}

// ---

type uint32Tag string

// Add adds a uint32 tag to the `span`
func (tag uint32Tag) Set(span opentracing.Span, value uint32) {
	span.SetTag(string(tag), value)
}

// ---

type uint16Tag string

// Add adds a uint16 tag to the `span`
func (tag uint16Tag) Set(span opentracing.Span, value uint16) {
	span.SetTag(string(tag), value)
}
<commit_msg>Fix typo<commit_after>package ext

import opentracing "github.com/opentracing/opentracing-go"

// These constants define common tag names recommended for better portability across
// tracing systems and languages/platforms.
//
// The tag names are defined as typed strings, so that in addition to the usual use
//
//     span.setTag(TagName, value)
//
// they also support value type validation via this additional syntax:
//
//    TagName.Set(span, value)
//
var (
	//////////////////////////////////////////////////////////////////////
	// SpanKind (client/server)
	//////////////////////////////////////////////////////////////////////

	// SpanKind hints at relationship between spans, e.g. client/server
	SpanKind = spanKindTag("span.kind")

	// SpanKindRPCClient marks a span representing the client-side of an RPC
	// or other remote call
	SpanKindRPCClient = SpanKindEnum("client")

	// SpanKindRPCServer marks a span representing the server-side of an RPC
	// or other remote call
	SpanKindRPCServer = SpanKindEnum("server")

	//////////////////////////////////////////////////////////////////////
	// Sampling hint
	//////////////////////////////////////////////////////////////////////

	// SamplingPriority determines the priority of sampling this Span.
	SamplingPriority = uint16Tag("sampling.priority")

	//////////////////////////////////////////////////////////////////////
	// Peer tags. These tags can be emitted by either client-side of
	// server-side to describe the other side/service in a peer-to-peer
	// communications, like an RPC call.
	//////////////////////////////////////////////////////////////////////

	// PeerService records the service name of the peer
	PeerService = stringTag("peer.service")

	// PeerHostname records the host name of the peer
	PeerHostname = stringTag("peer.hostname")

	// PeerHostIPv4 records IP v4 host address of the peer
	PeerHostIPv4 = uint32Tag("peer.ipv4")

	// PeerHostIPv6 records IP v6 host address of the peer
	PeerHostIPv6 = stringTag("peer.ipv6")

	// PeerPort records port number of the peer
	PeerPort = uint16Tag("peer.port")

	//////////////////////////////////////////////////////////////////////
	// HTTP Tags
	//////////////////////////////////////////////////////////////////////

	// HTTPUrl should be the URL of the request being handled in this segment
	// of the trace, in standard URI format. The protocol is optional.
	HTTPUrl = stringTag("http.url")

	// HTTPMethod is the HTTP method of the request, and is case-insensitive.
	HTTPMethod = stringTag("http.method")

	// HTTPStatusCode is the numeric HTTP status code (200, 404, etc) of the
	// HTTP response.
	HTTPStatusCode = uint16Tag("http.status_code")
)

// ---

// SpanKindEnum represents common span types
type SpanKindEnum string

type spanKindTag string

// Add adds a string tag to the `span`
func (tag spanKindTag) Set(span opentracing.Span, value SpanKindEnum) {
	span.SetTag(string(tag), value)
}

// ---

type stringTag string

// Add adds a string tag to the `span`
func (tag stringTag) Set(span opentracing.Span, value string) {
	span.SetTag(string(tag), value)
}

// ---

type uint32Tag string

// Add adds a uint32 tag to the `span`
func (tag uint32Tag) Set(span opentracing.Span, value uint32) {
	span.SetTag(string(tag), value)
}

// ---

type uint16Tag string

// Add adds a uint16 tag to the `span`
func (tag uint16Tag) Set(span opentracing.Span, value uint16) {
	span.SetTag(string(tag), value)
}
<|endoftext|>


<commit_before>package helm_test

import (
	"os"
	"strings"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/werf/werf/integration/utils"
)

var _ = Describe("deploy and rollback chart", func() {
	var releaseName string
	var releaseNamespace string

	BeforeEach(func() {
		releaseName = utils.ProjectName()
		releaseNamespace = releaseName
	})

	AfterEach(func() {
		utils.RunSucceedCommand(
			testDirPath,
			werfBinPath,
			"helm", "uninstall", releaseName, "--namespace", releaseNamespace,
		)
	})

	When("deploy local chart", func() {
		BeforeEach(func() {
			utils.CopyIn(utils.FixturePath("chart_1"), testDirPath)
		})

		It("should deploy chart in working directory", func() {
			utils.RunSucceedCommand(
				testDirPath,
				werfBinPath,
				"helm", "install", releaseName, ".", "--namespace", releaseNamespace,
			)
		})

		When("first release has been deployed", func() {
			BeforeEach(func() {
				utils.RunSucceedCommand(
					testDirPath,
					werfBinPath,
					"helm", "install", releaseName, ".", "--namespace", releaseNamespace,
				)
			})

			It("should get release templates and values", func() {
				output := utils.SucceedCommandOutputString(
					testDirPath,
					werfBinPath,
					"helm", "get", "all", releaseName, "--namespace", releaseNamespace,
				)

				expectedSubStrings := []string{
					"REVISION: 1",
					"COMPUTED VALUES",
					"HOOKS",
					"MANIFEST",
					"# Source: chart/templates/serviceaccount.yaml",
					"# Source: chart/templates/service.yaml",
					"# Source: chart/templates/deployment.yaml",
				}

				for _, expectedSubString := range expectedSubStrings {
					Ω(output).Should(ContainSubstring(expectedSubString))
				}
			})

			When("second release has been deployed", func() {
				BeforeEach(func() {
					Ω(os.RemoveAll(testDirPath)).ShouldNot(HaveOccurred())
					utils.CopyIn(utils.FixturePath("chart_2"), testDirPath)

					utils.RunSucceedCommand(
						testDirPath,
						werfBinPath,
						"helm", "upgrade", releaseName, ".", "--namespace", releaseNamespace,
					)
				})

				It("should get release templates and values", func() {
					output := utils.SucceedCommandOutputString(
						testDirPath,
						werfBinPath,
						"helm", "get", "all", releaseName, "--namespace", releaseNamespace,
					)

					expectedSubStrings := []string{
						"REVISION: 2",
						"COMPUTED VALUES",
						"HOOKS",
						"MANIFEST",
						"# Source: chart/templates/service.yaml",
					}

					notExpectedSubStrings := []string{
						"# Source: chart/templates/serviceaccount.yaml",
						"# Source: chart/templates/deployment.yaml",
					}

					for _, expectedSubString := range expectedSubStrings {
						Ω(output).Should(ContainSubstring(expectedSubString))
					}

					for _, notExpectedSubString := range notExpectedSubStrings {
						Ω(output).ShouldNot(ContainSubstring(notExpectedSubString))
					}
				})

				It("should list release", func() {
					output := utils.SucceedCommandOutputString(
						testDirPath,
						werfBinPath,
						"helm", "list", "--namespace", releaseNamespace,
					)

					Ω(output).Should(ContainSubstring(releaseName))
				})

				It("should get release history", func() {
					output := utils.SucceedCommandOutputString(
						testDirPath,
						werfBinPath,
						"helm", "history", releaseName, "--namespace", releaseNamespace,
					)

					Ω(strings.Count(output, "superseded")).Should(BeEquivalentTo(1))
					Ω(strings.Count(output, "deployed")).Should(BeEquivalentTo(1))
				})

				It("should rollback release", func() {
					utils.RunSucceedCommand(
						testDirPath,
						werfBinPath,
						"helm", "rollback", releaseName, "1", "--namespace", releaseNamespace,
					)

					output := utils.SucceedCommandOutputString(
						testDirPath,
						werfBinPath,
						"helm", "get", "all", releaseName, "--namespace", releaseNamespace,
					)

					Ω(output).Should(ContainSubstring("REVISION: 3"))
				})
			})
		})
	})

	When("deploy by chart reference", func() {
		BeforeEach(func() {
			stubs.SetEnv("WERF_HELM_HOME", testDirPath)

			utils.RunSucceedCommand(
				testDirPath,
				werfBinPath,
				"helm", "repo", "add", "stable", "https://charts.helm.sh/stable",
			)
		})

		It("should deploy chart by chart reference", func() {
			utils.RunSucceedCommand(
				testDirPath,
				werfBinPath,
				"helm", "install", releaseName, "stable/nginx-ingress", "--namespace", releaseNamespace,
			)
		})
	})
})
<commit_msg>[tests] Helm Suite: fix test helm home<commit_after>package helm_test

import (
	"os"
	"strings"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/werf/werf/integration/utils"
)

var _ = Describe("deploy and rollback chart", func() {
	var releaseName string
	var releaseNamespace string

	BeforeEach(func() {
		releaseName = utils.ProjectName()
		releaseNamespace = releaseName
	})

	AfterEach(func() {
		utils.RunSucceedCommand(
			testDirPath,
			werfBinPath,
			"helm", "uninstall", releaseName, "--namespace", releaseNamespace,
		)
	})

	When("deploy local chart", func() {
		BeforeEach(func() {
			utils.CopyIn(utils.FixturePath("chart_1"), testDirPath)
		})

		It("should deploy chart in working directory", func() {
			utils.RunSucceedCommand(
				testDirPath,
				werfBinPath,
				"helm", "install", releaseName, ".", "--namespace", releaseNamespace,
			)
		})

		When("first release has been deployed", func() {
			BeforeEach(func() {
				utils.RunSucceedCommand(
					testDirPath,
					werfBinPath,
					"helm", "install", releaseName, ".", "--namespace", releaseNamespace,
				)
			})

			It("should get release templates and values", func() {
				output := utils.SucceedCommandOutputString(
					testDirPath,
					werfBinPath,
					"helm", "get", "all", releaseName, "--namespace", releaseNamespace,
				)

				expectedSubStrings := []string{
					"REVISION: 1",
					"COMPUTED VALUES",
					"HOOKS",
					"MANIFEST",
					"# Source: chart/templates/serviceaccount.yaml",
					"# Source: chart/templates/service.yaml",
					"# Source: chart/templates/deployment.yaml",
				}

				for _, expectedSubString := range expectedSubStrings {
					Ω(output).Should(ContainSubstring(expectedSubString))
				}
			})

			When("second release has been deployed", func() {
				BeforeEach(func() {
					Ω(os.RemoveAll(testDirPath)).ShouldNot(HaveOccurred())
					utils.CopyIn(utils.FixturePath("chart_2"), testDirPath)

					utils.RunSucceedCommand(
						testDirPath,
						werfBinPath,
						"helm", "upgrade", releaseName, ".", "--namespace", releaseNamespace,
					)
				})

				It("should get release templates and values", func() {
					output := utils.SucceedCommandOutputString(
						testDirPath,
						werfBinPath,
						"helm", "get", "all", releaseName, "--namespace", releaseNamespace,
					)

					expectedSubStrings := []string{
						"REVISION: 2",
						"COMPUTED VALUES",
						"HOOKS",
						"MANIFEST",
						"# Source: chart/templates/service.yaml",
					}

					notExpectedSubStrings := []string{
						"# Source: chart/templates/serviceaccount.yaml",
						"# Source: chart/templates/deployment.yaml",
					}

					for _, expectedSubString := range expectedSubStrings {
						Ω(output).Should(ContainSubstring(expectedSubString))
					}

					for _, notExpectedSubString := range notExpectedSubStrings {
						Ω(output).ShouldNot(ContainSubstring(notExpectedSubString))
					}
				})

				It("should list release", func() {
					output := utils.SucceedCommandOutputString(
						testDirPath,
						werfBinPath,
						"helm", "list", "--namespace", releaseNamespace,
					)

					Ω(output).Should(ContainSubstring(releaseName))
				})

				It("should get release history", func() {
					output := utils.SucceedCommandOutputString(
						testDirPath,
						werfBinPath,
						"helm", "history", releaseName, "--namespace", releaseNamespace,
					)

					Ω(strings.Count(output, "superseded")).Should(BeEquivalentTo(1))
					Ω(strings.Count(output, "deployed")).Should(BeEquivalentTo(1))
				})

				It("should rollback release", func() {
					utils.RunSucceedCommand(
						testDirPath,
						werfBinPath,
						"helm", "rollback", releaseName, "1", "--namespace", releaseNamespace,
					)

					output := utils.SucceedCommandOutputString(
						testDirPath,
						werfBinPath,
						"helm", "get", "all", releaseName, "--namespace", releaseNamespace,
					)

					Ω(output).Should(ContainSubstring("REVISION: 3"))
				})
			})
		})
	})

	When("deploy by chart reference", func() {
		BeforeEach(func() {
			stubs.SetEnv("XDG_DATA_HOME", testDirPath)
			stubs.SetEnv("XDG_CACHE_HOME", testDirPath)
			stubs.SetEnv("XDG_CONFIG_HOME", testDirPath)

			utils.RunSucceedCommand(
				testDirPath,
				werfBinPath,
				"helm", "repo", "add", "stable", "https://charts.helm.sh/stable",
			)
		})

		It("should deploy chart by chart reference", func() {
			utils.RunSucceedCommand(
				testDirPath,
				werfBinPath,
				"helm", "install", releaseName, "stable/nginx-ingress", "--namespace", releaseNamespace,
			)
		})
	})
})
<|endoftext|>


<commit_before>/*
Package framed provides an implementation of io.ReadWriteCloser that reads and
writes whole frames only.

Frames are length-prefixed.  The first two bytes are an unsigned 16 bit int
stored in little-endian byte order indicating the length of the content.  The
remaining bytes are the actual content of the frame.

The use of a uint16 means that the maximum possible frame size (MAX_FRAME_SIZE)
is 65535.
*/
package framed

import (
	"encoding/binary"
	"fmt"
	"io"
	"sync"
)

const (
	// MAX_FRAME_SIZE is the maximum possible size of a frame (not including the
	// length prefix)
	MAX_FRAME_SIZE = 65535
)

var endianness = binary.LittleEndian

/*
A framed.Reader enhances an io.ReadCloser to read data in contiguous frames.
It implements the io.Reader interface, but nlike typical Readers it only returns
whole frames.  Unlike typical Writers, it will not allow frames to be
fragmented.

Although the underlying stream may be safe to use from multiple
goroutines, a framed.Reader is not.
*/
type Reader struct {
	Stream io.Reader // the raw underlying connection
}

/*
A framed.Writer enhances an io.WriteCLoser to write data in contiguous frames.
It implements the io.Writer interface, but unlike typical Writers, it includes
information that allows a corresponding framed.Reader to read whole frames
without them being fragmented.

A framed.Writer also supports a method that writes multiple buffers to the
underlying stream as a single frame.
*/
type Writer struct {
	Stream io.Writer // the raw underlying connection
	mutex  sync.Mutex
}

func NewReader(r io.Reader) *Reader {
	return &Reader{Stream: r}
}

func NewWriter(w io.Writer) *Writer {
	return &Writer{Stream: w}
}

/*
Read implements the function from io.Reader.  Unlike io.Reader.Read,
frame.Read only returns full frames of data (assuming that the data was written
by a framed.Writer).
*/
func (framed *Reader) Read(buffer []byte) (n int, err error) {
	var nb uint16
	err = binary.Read(framed.Stream, endianness, &nb)
	if err != nil {
		return
	}

	n = int(nb)

	bufferSize := len(buffer)
	if n > bufferSize {
		return 0, fmt.Errorf("Buffer of size %d is too small to hold frame of size %d", bufferSize, n)
	}
	// Read into buffer
	var totalRead int
	for totalRead = 0; totalRead < n; {
		var nr int
		nr, err = framed.Stream.Read(buffer[totalRead:n])
		if err != nil {
			return
		}
		totalRead += nr
	}
	if totalRead != n {
		err = fmt.Errorf("%d bytes read did not match %d bytes expected", totalRead, n)
	}
	return
}

/*
Write implements the Write method from io.Writer.  It prepends a frame length
header that allows the framed.Reader on the other end to read the whole frame.
*/
func (framed *Writer) Write(frame []byte) (n int, err error) {
	framed.mutex.Lock()
	defer framed.mutex.Unlock()

	n = len(frame)

	// Write the length header
	if err = binary.Write(framed.Stream, endianness, uint16(n)); err != nil {
		return
	}

	// Write the data
	var written int
	if written, err = framed.Stream.Write(frame); err != nil {
		return
	}
	if written != n {
		err = fmt.Errorf("%d bytes written, expected to write %d", written, n)
	}
	return
}

func (framed *Writer) WritePieces(pieces ...[]byte) (n int, err error) {
	framed.mutex.Lock()
	defer framed.mutex.Unlock()

	for _, piece := range pieces {
		n = n + len(piece)
	}

	// Write the length header
	if err = binary.Write(framed.Stream, endianness, uint16(n)); err != nil {
		return
	}

	// Write the data
	var written int
	for _, piece := range pieces {
		var nw int
		if nw, err = framed.Stream.Write(piece); err != nil {
			return
		}
		written = written + nw
	}
	if written != n {
		err = fmt.Errorf("%d bytes written, expected to write %d", written, n)
	}
	return
}
<commit_msg>Using io.ReadFull in place of custom logic<commit_after>/*
Package framed provides an implementation of io.ReadWriteCloser that reads and
writes whole frames only.

Frames are length-prefixed.  The first two bytes are an unsigned 16 bit int
stored in little-endian byte order indicating the length of the content.  The
remaining bytes are the actual content of the frame.

The use of a uint16 means that the maximum possible frame size (MAX_FRAME_SIZE)
is 65535.
*/
package framed

import (
	"encoding/binary"
	"fmt"
	"io"
	"sync"
)

const (
	// MAX_FRAME_SIZE is the maximum possible size of a frame (not including the
	// length prefix)
	MAX_FRAME_SIZE = 65535
)

var endianness = binary.LittleEndian

/*
A framed.Reader enhances an io.ReadCloser to read data in contiguous frames.
It implements the io.Reader interface, but nlike typical Readers it only returns
whole frames.  Unlike typical Writers, it will not allow frames to be
fragmented.

Although the underlying stream may be safe to use from multiple
goroutines, a framed.Reader is not.
*/
type Reader struct {
	Stream io.Reader // the raw underlying connection
}

/*
A framed.Writer enhances an io.WriteCLoser to write data in contiguous frames.
It implements the io.Writer interface, but unlike typical Writers, it includes
information that allows a corresponding framed.Reader to read whole frames
without them being fragmented.

A framed.Writer also supports a method that writes multiple buffers to the
underlying stream as a single frame.
*/
type Writer struct {
	Stream io.Writer // the raw underlying connection
	mutex  sync.Mutex
}

func NewReader(r io.Reader) *Reader {
	return &Reader{Stream: r}
}

func NewWriter(w io.Writer) *Writer {
	return &Writer{Stream: w}
}

/*
Read implements the function from io.Reader.  Unlike io.Reader.Read,
frame.Read only returns full frames of data (assuming that the data was written
by a framed.Writer).
*/
func (framed *Reader) Read(buffer []byte) (n int, err error) {
	var nb uint16
	err = binary.Read(framed.Stream, endianness, &nb)
	if err != nil {
		return
	}

	n = int(nb)

	bufferSize := len(buffer)
	if n > bufferSize {
		return 0, fmt.Errorf("Buffer of size %d is too small to hold frame of size %d", bufferSize, n)
	}
	// Read into buffer
	n, err = io.ReadFull(framed.Stream, buffer[:n])
	return
}

/*
Write implements the Write method from io.Writer.  It prepends a frame length
header that allows the framed.Reader on the other end to read the whole frame.
*/
func (framed *Writer) Write(frame []byte) (n int, err error) {
	framed.mutex.Lock()
	defer framed.mutex.Unlock()

	n = len(frame)

	// Write the length header
	if err = binary.Write(framed.Stream, endianness, uint16(n)); err != nil {
		return
	}

	// Write the data
	var written int
	if written, err = framed.Stream.Write(frame); err != nil {
		return
	}
	if written != n {
		err = fmt.Errorf("%d bytes written, expected to write %d", written, n)
	}
	return
}

func (framed *Writer) WritePieces(pieces ...[]byte) (n int, err error) {
	framed.mutex.Lock()
	defer framed.mutex.Unlock()

	for _, piece := range pieces {
		n = n + len(piece)
	}

	// Write the length header
	if err = binary.Write(framed.Stream, endianness, uint16(n)); err != nil {
		return
	}

	// Write the data
	var written int
	for _, piece := range pieces {
		var nw int
		if nw, err = framed.Stream.Write(piece); err != nil {
			return
		}
		written = written + nw
	}
	if written != n {
		err = fmt.Errorf("%d bytes written, expected to write %d", written, n)
	}
	return
}
<|endoftext|>


<commit_before>//go:build e2e
// +build e2e

package main

import (
	"context"
	"fmt"
	"os"
	"testing"
	"time"

	tfe "github.com/hashicorp/go-tfe"
	"github.com/hashicorp/go-uuid"
)

const (
	expectConsoleTimeout = 15 * time.Second
)

type tfCommand struct {
	command           []string
	expectedCmdOutput string
	expectedErr       string
	expectError       bool
	userInput         []string
	postInputOutput   []string
}

type operationSets struct {
	commands []tfCommand
	prep     func(t *testing.T, orgName, dir string)
}

type testCases map[string]struct {
	operations  []operationSets
	validations func(t *testing.T, orgName string)
}

func createOrganization(t *testing.T) (*tfe.Organization, func()) {
	ctx := context.Background()
	org, err := tfeClient.Organizations.Create(ctx, tfe.OrganizationCreateOptions{
		Name:  tfe.String("tst-" + randomString(t)),
		Email: tfe.String(fmt.Sprintf("%s@tfe.local", randomString(t))),
	})
	if err != nil {
		t.Fatal(err)
	}

	_, err = tfeClient.Admin.Organizations.Update(ctx, org.Name, tfe.AdminOrganizationUpdateOptions{
		AccessBetaTools: tfe.Bool(true),
	})
	if err != nil {
		t.Fatal(err)
	}

	return org, func() {
		if err := tfeClient.Organizations.Delete(ctx, org.Name); err != nil {
			t.Errorf("Error destroying organization! WARNING: Dangling resources\n"+
				"may exist! The full error is shown below.\n\n"+
				"Organization: %s\nError: %s", org.Name, err)
		}
	}
}

func createWorkspace(t *testing.T, orgName string, wOpts tfe.WorkspaceCreateOptions) *tfe.Workspace {
	ctx := context.Background()
	w, err := tfeClient.Workspaces.Create(ctx, orgName, wOpts)
	if err != nil {
		t.Fatal(err)
	}

	return w
}

func getWorkspace(workspaces []*tfe.Workspace, workspace string) (*tfe.Workspace, bool) {
	for _, ws := range workspaces {
		if ws.Name == workspace {
			return ws, false
		}
	}
	return nil, true
}

func randomString(t *testing.T) string {
	v, err := uuid.GenerateUUID()
	if err != nil {
		t.Fatal(err)
	}
	return v
}

func terraformConfigLocalBackend() string {
	return fmt.Sprintf(`
terraform {
  backend "local" {
  }
}

output "val" {
  value = "${terraform.workspace}"
}
`)
}

func terraformConfigRemoteBackendName(org, name string) string {
	return fmt.Sprintf(`
terraform {
  backend "remote" {
    hostname = "%s"
    organization = "%s"

    workspaces {
      name = "%s"
    }
  }
}

output "val" {
  value = "${terraform.workspace}"
}
`, tfeHostname, org, name)
}

func terraformConfigRemoteBackendPrefix(org, prefix string) string {
	return fmt.Sprintf(`
terraform {
  backend "remote" {
    hostname = "%s"
    organization = "%s"

    workspaces {
      prefix = "%s"
    }
  }
}

output "val" {
  value = "${terraform.workspace}"
}
`, tfeHostname, org, prefix)
}

func terraformConfigCloudBackendTags(org, tag string) string {
	return fmt.Sprintf(`
terraform {
  cloud {
    hostname = "%s"
    organization = "%s"

    workspaces {
      tags = ["%s"]
    }
  }
}

output "tag_val" {
  value = "%s"
}
`, tfeHostname, org, tag, tag)
}

func terraformConfigCloudBackendName(org, name string) string {
	return fmt.Sprintf(`
terraform {
  cloud {
    hostname = "%s"
    organization = "%s"

    workspaces {
      name = "%s"
    }
  }
}

output "val" {
  value = "${terraform.workspace}"
}
`, tfeHostname, org, name)
}

func writeMainTF(t *testing.T, block string, dir string) {
	f, err := os.Create(fmt.Sprintf("%s/main.tf", dir))
	if err != nil {
		t.Fatal(err)
	}
	_, err = f.WriteString(block)
	if err != nil {
		t.Fatal(err)
	}
	f.Close()
}

// Ensure that TFC/E has a particular terraform version.
func hasTerraformVersion(t *testing.T, version string) bool {
	opts := tfe.AdminTerraformVersionsListOptions{
		ListOptions: tfe.ListOptions{
			PageNumber: 1,
			PageSize:   100,
		},
	}
	hasVersion := false

findTfVersion:
	for {
		// TODO: update go-tfe Read() to retrieve a terraform version by name.
		// Currently you can only retrieve by ID.
		tfVersionList, err := tfeClient.Admin.TerraformVersions.List(context.Background(), opts)
		if err != nil {
			t.Fatalf("Could not retrieve list of terraform versions: %v", err)
		}
		for _, item := range tfVersionList.Items {
			if item.Version == version {
				hasVersion = true
				break findTfVersion
			}
		}

		// Exit the loop when we've seen all pages.
		if tfVersionList.CurrentPage >= tfVersionList.TotalPages {
			break
		}

		// Update the page number to get the next page.
		opts.PageNumber = tfVersionList.NextPage
	}

	return hasVersion
}
<commit_msg>cloud/e2e: Disable cost estimation<commit_after>//go:build e2e
// +build e2e

package main

import (
	"context"
	"fmt"
	"os"
	"testing"
	"time"

	tfe "github.com/hashicorp/go-tfe"
	"github.com/hashicorp/go-uuid"
)

const (
	expectConsoleTimeout = 15 * time.Second
)

type tfCommand struct {
	command           []string
	expectedCmdOutput string
	expectedErr       string
	expectError       bool
	userInput         []string
	postInputOutput   []string
}

type operationSets struct {
	commands []tfCommand
	prep     func(t *testing.T, orgName, dir string)
}

type testCases map[string]struct {
	operations  []operationSets
	validations func(t *testing.T, orgName string)
}

func createOrganization(t *testing.T) (*tfe.Organization, func()) {
	ctx := context.Background()
	org, err := tfeClient.Organizations.Create(ctx, tfe.OrganizationCreateOptions{
		Name:                  tfe.String("tst-" + randomString(t)),
		Email:                 tfe.String(fmt.Sprintf("%s@tfe.local", randomString(t))),
		CostEstimationEnabled: tfe.Bool(false),
	})
	if err != nil {
		t.Fatal(err)
	}

	_, err = tfeClient.Admin.Organizations.Update(ctx, org.Name, tfe.AdminOrganizationUpdateOptions{
		AccessBetaTools: tfe.Bool(true),
	})
	if err != nil {
		t.Fatal(err)
	}

	return org, func() {
		if err := tfeClient.Organizations.Delete(ctx, org.Name); err != nil {
			t.Errorf("Error destroying organization! WARNING: Dangling resources\n"+
				"may exist! The full error is shown below.\n\n"+
				"Organization: %s\nError: %s", org.Name, err)
		}
	}
}

func createWorkspace(t *testing.T, orgName string, wOpts tfe.WorkspaceCreateOptions) *tfe.Workspace {
	ctx := context.Background()
	w, err := tfeClient.Workspaces.Create(ctx, orgName, wOpts)
	if err != nil {
		t.Fatal(err)
	}

	return w
}

func getWorkspace(workspaces []*tfe.Workspace, workspace string) (*tfe.Workspace, bool) {
	for _, ws := range workspaces {
		if ws.Name == workspace {
			return ws, false
		}
	}
	return nil, true
}

func randomString(t *testing.T) string {
	v, err := uuid.GenerateUUID()
	if err != nil {
		t.Fatal(err)
	}
	return v
}

func terraformConfigLocalBackend() string {
	return fmt.Sprintf(`
terraform {
  backend "local" {
  }
}

output "val" {
  value = "${terraform.workspace}"
}
`)
}

func terraformConfigRemoteBackendName(org, name string) string {
	return fmt.Sprintf(`
terraform {
  backend "remote" {
    hostname = "%s"
    organization = "%s"

    workspaces {
      name = "%s"
    }
  }
}

output "val" {
  value = "${terraform.workspace}"
}
`, tfeHostname, org, name)
}

func terraformConfigRemoteBackendPrefix(org, prefix string) string {
	return fmt.Sprintf(`
terraform {
  backend "remote" {
    hostname = "%s"
    organization = "%s"

    workspaces {
      prefix = "%s"
    }
  }
}

output "val" {
  value = "${terraform.workspace}"
}
`, tfeHostname, org, prefix)
}

func terraformConfigCloudBackendTags(org, tag string) string {
	return fmt.Sprintf(`
terraform {
  cloud {
    hostname = "%s"
    organization = "%s"

    workspaces {
      tags = ["%s"]
    }
  }
}

output "tag_val" {
  value = "%s"
}
`, tfeHostname, org, tag, tag)
}

func terraformConfigCloudBackendName(org, name string) string {
	return fmt.Sprintf(`
terraform {
  cloud {
    hostname = "%s"
    organization = "%s"

    workspaces {
      name = "%s"
    }
  }
}

output "val" {
  value = "${terraform.workspace}"
}
`, tfeHostname, org, name)
}

func writeMainTF(t *testing.T, block string, dir string) {
	f, err := os.Create(fmt.Sprintf("%s/main.tf", dir))
	if err != nil {
		t.Fatal(err)
	}
	_, err = f.WriteString(block)
	if err != nil {
		t.Fatal(err)
	}
	f.Close()
}

// Ensure that TFC/E has a particular terraform version.
func hasTerraformVersion(t *testing.T, version string) bool {
	opts := tfe.AdminTerraformVersionsListOptions{
		ListOptions: tfe.ListOptions{
			PageNumber: 1,
			PageSize:   100,
		},
	}
	hasVersion := false

findTfVersion:
	for {
		// TODO: update go-tfe Read() to retrieve a terraform version by name.
		// Currently you can only retrieve by ID.
		tfVersionList, err := tfeClient.Admin.TerraformVersions.List(context.Background(), opts)
		if err != nil {
			t.Fatalf("Could not retrieve list of terraform versions: %v", err)
		}
		for _, item := range tfVersionList.Items {
			if item.Version == version {
				hasVersion = true
				break findTfVersion
			}
		}

		// Exit the loop when we've seen all pages.
		if tfVersionList.CurrentPage >= tfVersionList.TotalPages {
			break
		}

		// Update the page number to get the next page.
		opts.PageNumber = tfVersionList.NextPage
	}

	return hasVersion
}
<|endoftext|>


<commit_before>package registry

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"sync"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/hcp-sdk-go/clients/cloud-packer-service/preview/2021-04-30/models"
	registryimage "github.com/hashicorp/packer-plugin-sdk/packer/registry/image"
	"github.com/hashicorp/packer/internal/registry/env"
	"google.golang.org/grpc/codes"
)

// Bucket represents a single Image bucket on the HCP Packer registry.
type Bucket struct {
	Slug        string
	Description string
	Destination string
	Labels      map[string]string
	Iteration   *Iteration
	client      *Client
}

// NewBucketWithIteration initializes a simple Bucket that can be used publishing Packer build
// images to the HCP Packer registry.
func NewBucketWithIteration(opts IterationOptions) (*Bucket, error) {
	b := Bucket{}

	i, err := NewIteration(opts)
	if err != nil {
		return nil, err
	}

	b.Iteration = i
	return &b, nil
}

func (b *Bucket) Validate() error {
	if b.Slug == "" {
		return fmt.Errorf("no Packer bucket name defined; either the environment variable %q is undefined or the HCL configuration has no build name", env.HCPPackerBucket)
	}
	return nil
}

// connect initializes a client connection to a remote HCP Packer Registry service on HCP.
// Upon a successful connection the initialized client is persisted on the Bucket b for later usage.
func (b *Bucket) connect() error {
	if b.client != nil {
		return nil
	}

	registryClient, err := NewClient()
	if err != nil {
		return errors.New("Failed to create client connection to artifact registry: " + err.Error())
	}
	b.client = registryClient
	return nil
}

// Initialize registers the Bucket b with the configured HCP Packer Registry.
// Upon initialization a Bucket will be upserted to, and new iteration will be created for the build if the configured
// fingerprint has no associated iterations. Lastly, the initialization process with register the builds that need to be
// completed before an iteration can be marked as DONE.
//
// b.Initialize() must be called before any data can be published to the configured HCP Packer Registry.
// TODO ensure initialize can only be called once
func (b *Bucket) Initialize(ctx context.Context) error {

	if err := b.connect(); err != nil {
		return err
	}

	b.Destination = fmt.Sprintf("%s/%s", b.client.OrganizationID, b.client.ProjectID)

	bucketInput := &models.HashicorpCloudPackerCreateBucketRequest{
		BucketSlug:  b.Slug,
		Description: b.Description,
		Labels:      b.Labels,
	}

	err := UpsertBucket(ctx, b.client, bucketInput)
	if err != nil {
		return fmt.Errorf("failed to initialize bucket %q: %w", b.Slug, err)
	}

	return b.initializeIteration(ctx)
}

func (b *Bucket) RegisterBuildForComponent(sourceName string) {
	if b == nil {
		return
	}

	if _, ok := b.Iteration.builds.Load(sourceName); ok {
		return
	}
	b.Iteration.expectedBuilds = append(b.Iteration.expectedBuilds, sourceName)
}

// CreateInitialBuildForIteration will create a build entry on the HCP Packer Registry for the named componentType.
// This initial creation is needed so that Packer can properly track when an iteration is complete.
func (b *Bucket) CreateInitialBuildForIteration(ctx context.Context, componentType string) error {

	status := models.HashicorpCloudPackerBuildStatusUNSET
	buildInput := &models.HashicorpCloudPackerCreateBuildRequest{
		BucketSlug:  b.Slug,
		Fingerprint: b.Iteration.Fingerprint,
		IterationID: b.Iteration.ID,
		Build: &models.HashicorpCloudPackerBuildCreateBody{
			ComponentType: componentType,
			PackerRunUUID: b.Iteration.RunUUID,
			Status:        status,
		},
	}

	id, err := CreateBuild(ctx, b.client, buildInput)
	if err != nil {
		return err
	}

	build := &Build{
		ID:            id,
		ComponentType: componentType,
		RunUUID:       b.Iteration.RunUUID,
		Status:        status,
		Labels:        make(map[string]string),
		Images:        make(map[string]registryimage.Image),
	}

	log.Println("[TRACE] creating initial build for component", componentType)
	b.Iteration.builds.Store(componentType, build)

	return nil
}

// UpdateBuildStatus updates the status of a build entry on the HCP Packer registry with its current local status.
func (b *Bucket) UpdateBuildStatus(ctx context.Context, name string, status models.HashicorpCloudPackerBuildStatus) error {
	if status == models.HashicorpCloudPackerBuildStatusDONE {
		return b.markBuildComplete(ctx, name)
	}

	// Lets check if we have something already for this build
	build, ok := b.Iteration.builds.Load(name)
	if !ok {
		return fmt.Errorf("no build for the component %q associated to the iteration %q", name, b.Iteration.ID)
	}

	buildToUpdate, ok := build.(*Build)
	if !ok {
		return fmt.Errorf("the build for the component %q does not appear to be a valid registry Build", name)
	}

	if buildToUpdate.ID == "" {
		return fmt.Errorf("the build for the component %q does not have a valid id", name)
	}

	buildInput := &models.HashicorpCloudPackerUpdateBuildRequest{
		BuildID: buildToUpdate.ID,
		Updates: &models.HashicorpCloudPackerBuildUpdates{
			PackerRunUUID: buildToUpdate.RunUUID,
			Labels:        buildToUpdate.Labels,
			Status:        status,
		},
	}

	_, err := UpdateBuild(ctx, b.client, buildInput)
	if err != nil {
		return err
	}
	buildToUpdate.Status = status
	b.Iteration.builds.Store(name, buildToUpdate)
	return nil
}

// markBuildComplete should be called to set a build on the HCP Packer registry to DONE.
// Upon a successful call markBuildComplete will publish all images created by the named build,
// and set the registry build to done. A build with no images can not be set to DONE.
func (b *Bucket) markBuildComplete(ctx context.Context, name string) error {
	build, ok := b.Iteration.builds.Load(name)
	if !ok {
		return fmt.Errorf("no build for the component %q associated to the iteration %q", name, b.Iteration.ID)
	}

	buildToUpdate, ok := build.(*Build)
	if !ok {
		return fmt.Errorf("the build for the component %q does not appear to be a valid registry Build", name)
	}

	if buildToUpdate.ID == "" {
		return fmt.Errorf("the build for the component %q does not have a valid id", name)
	}

	status := models.HashicorpCloudPackerBuildStatusDONE

	if buildToUpdate.Status == status {
		// let's no mess with anything that is already done
		return nil
	}

	buildInput := &models.HashicorpCloudPackerUpdateBuildRequest{
		BuildID: buildToUpdate.ID,
		Updates: &models.HashicorpCloudPackerBuildUpdates{
			PackerRunUUID: buildToUpdate.RunUUID,
			Labels:        buildToUpdate.Labels,
			Status:        status,
		},
	}

	if len(buildToUpdate.Images) == 0 {
		return fmt.Errorf("setting a build to DONE with no published images is not currently supported.")
	}

	var providerName string
	images := make([]*models.HashicorpCloudPackerImageCreateBody, 0, len(buildToUpdate.Images))
	for _, image := range buildToUpdate.Images {
		if providerName == "" {
			providerName = image.ProviderName
		}
		images = append(images, &models.HashicorpCloudPackerImageCreateBody{ImageID: image.ImageID, Region: image.ProviderRegion})
	}

	buildInput.Updates.CloudProvider = providerName
	buildInput.Updates.Images = images

	_, err := UpdateBuild(ctx, b.client, buildInput)
	if err != nil {
		return err
	}

	buildToUpdate.Status = status
	b.Iteration.builds.Store(name, buildToUpdate)
	return nil
}

// UpdateImageForBuild appends one or more images artifacts to the build referred to by componentType.
func (b *Bucket) UpdateImageForBuild(componentType string, images ...registryimage.Image) error {
	return b.Iteration.AddImageToBuild(componentType, images...)
}

// UpdateLabelsForBuild merges the contents of data to the labels associated with the build referred to by componentType.
func (b *Bucket) UpdateLabelsForBuild(componentType string, data map[string]string) error {
	return b.Iteration.AddLabelsToBuild(componentType, data)
}

// Load defaults from environment variables
func (b *Bucket) LoadDefaultSettingsFromEnv() {
	// Configure HCP Packer Registry destination
	if b.Slug == "" {
		b.Slug = os.Getenv(env.HCPPackerBucket)
	}

	// Set some iteration values. For Packer RunUUID should always be set.
	// Creating an iteration differently? Let's not overwrite a UUID that might be set.
	if b.Iteration.RunUUID == "" {
		b.Iteration.RunUUID = os.Getenv("PACKER_RUN_UUID")
	}

}

// createIteration creates an empty iteration for a give bucket on the HCP Packer registry.
// The iteration can then be stored locally and used for tracking build status and images for a running
// Packer build.
func (b *Bucket) createIteration() (*models.HashicorpCloudPackerIteration, error) {
	iterationInput := &models.HashicorpCloudPackerCreateIterationRequest{
		BucketSlug:  b.Slug,
		Fingerprint: b.Iteration.Fingerprint,
	}

	iterationResp, err := CreateIteration(context.TODO(), b.client, iterationInput)
	if err != nil {
		return nil, fmt.Errorf("failed to create Iteration for Bucket %s with error: %w", b.Slug, err)
	}

	if iterationResp == nil {
		return nil, fmt.Errorf("failed to create Iteration for Bucket %s with error: %w", b.Slug, err)
	}

	log.Println("[TRACE] a valid iteration for build was created with the Id", iterationResp.ID)
	return iterationResp, nil
}

// initializeIteration populates the bucket iteration with the details needed for tracking builds for a Packer run.
// If an existing Packer registry iteration exists for the said iteration fingerprint, calling initialize on iteration
// that doesn't yet exist will call createIteration to create the entry on the HCP packer registry for the given bucket.
// All build details will be created (if they don't exists) and added to b.Iteration.builds for tracking during runtime.
func (b *Bucket) initializeIteration(ctx context.Context) error {

	// load existing iteration using fingerprint.
	iterationResp, err := GetIteration(ctx, b.client, b.Slug, b.Iteration.Fingerprint)
	if checkErrorCode(err, codes.Aborted) {
		//probably means Iteration doesn't exist need a way to check the error
		iterationResp, err = b.createIteration()
	}

	if err != nil {
		return fmt.Errorf("failed to initialize iteration for fingerprint %s: %s", b.Iteration.Fingerprint, err)
	}

	if iterationResp == nil {
		return fmt.Errorf("failed to initialize iteration details for Bucket %s with error: %w", b.Slug, err)
	}

	log.Println("[TRACE] a valid iteration was retrieved with the id", iterationResp.ID)
	b.Iteration.ID = iterationResp.ID

	// If the iteration is completed and there are no new builds to add, Packer
	// should exit and inform the user that artifacts already exists for the
	// fingerprint associated with the iteration.
	if iterationResp.Complete {
		return fmt.Errorf("This iteration associated to the fingerprint %s is complete. "+
			"If you wish to add a new build to this image a new iteration must be created by changing the build fingerprint.", b.Iteration.Fingerprint)
	}

	// list all this iteration's builds so we can figure out which ones
	// we want to run against. TODO: pagination?
	existingBuilds, err := ListBuilds(ctx, b.client, b.Slug, iterationResp.ID)
	if err != nil {
		return fmt.Errorf("error listing builds for this existing iteration: %s", err)
	}

	var toCreate []string
	for _, expected := range b.Iteration.expectedBuilds {
		var found bool
		for _, existing := range existingBuilds {
			if existing.ComponentType == expected {
				found = true
				build := &Build{
					ID:            existing.ID,
					ComponentType: existing.ComponentType,
					RunUUID:       b.Iteration.RunUUID,
					Status:        existing.Status,
					Labels:        existing.Labels,
				}
				b.Iteration.builds.Store(existing.ComponentType, build)

				// TODO validate that this is safe. For builds that are DONE do we want to keep track of completed things
				// potential issue on updating the status of a build that is already DONE. Is this possible?
				for _, image := range existing.Images {

					err := b.UpdateImageForBuild(existing.ComponentType, registryimage.Image{
						ImageID:        image.ImageID,
						ProviderRegion: image.Region,
					})

					if err != nil {
						log.Printf("[TRACE] unable to load existing images for %q: %v", existing.ComponentType, err)
					}
				}

				log.Printf("[TRACE] a build of component type %s already exists; skipping the create call", expected)
				break
			}
		}

		if !found {
			missingbuild := expected
			toCreate = append(toCreate, missingbuild)
		}
	}

	var errs *multierror.Error
	var wg sync.WaitGroup

	for _, buildName := range toCreate {
		wg.Add(1)
		go func(name string) {
			defer wg.Done()
			log.Printf("[TRACE] registering build with iteration for %q.", name)
			// Need a way to handle skipping builds that were already created.
			// TODO when we load an existing iteration we will probably have a build Id so we should skip.
			// we also need to bubble up the errors here.
			err := b.CreateInitialBuildForIteration(ctx, name)
			if checkErrorCode(err, codes.AlreadyExists) {
				// Check whether build is complete, and if so, skip it.
				log.Printf("[TRACE] build %s already exists in Packer registry, continuing...", name)
				return
			}

			errs = multierror.Append(errs, err)
		}(buildName)
	}
	wg.Wait()

	return errs.ErrorOrNil()
}

// IsExpectingBuildForComponent returns true if the component referenced by buildName is part of the iteration
// and is not marked as DONE on the HCP Packer registry.
func (b *Bucket) IsExpectingBuildForComponent(buildName string) bool {

	v, ok := b.Iteration.builds.Load(buildName)
	if !ok {
		return false
	}

	build := v.(*Build)
	hasBuildID := build.ID != ""
	hasImages := len(build.Images) == 0
	isNotDone := build.Status != models.HashicorpCloudPackerBuildStatusDONE

	return hasBuildID && hasImages && isNotDone
}
<commit_msg>add source image id to build from PARtifact<commit_after>package registry

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"sync"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/hcp-sdk-go/clients/cloud-packer-service/preview/2021-04-30/models"
	registryimage "github.com/hashicorp/packer-plugin-sdk/packer/registry/image"
	"github.com/hashicorp/packer/internal/registry/env"
	"google.golang.org/grpc/codes"
)

// Bucket represents a single Image bucket on the HCP Packer registry.
type Bucket struct {
	Slug        string
	Description string
	Destination string
	Labels      map[string]string
	Iteration   *Iteration
	client      *Client
}

// NewBucketWithIteration initializes a simple Bucket that can be used publishing Packer build
// images to the HCP Packer registry.
func NewBucketWithIteration(opts IterationOptions) (*Bucket, error) {
	b := Bucket{}

	i, err := NewIteration(opts)
	if err != nil {
		return nil, err
	}

	b.Iteration = i
	return &b, nil
}

func (b *Bucket) Validate() error {
	if b.Slug == "" {
		return fmt.Errorf("no Packer bucket name defined; either the environment variable %q is undefined or the HCL configuration has no build name", env.HCPPackerBucket)
	}
	return nil
}

// connect initializes a client connection to a remote HCP Packer Registry service on HCP.
// Upon a successful connection the initialized client is persisted on the Bucket b for later usage.
func (b *Bucket) connect() error {
	if b.client != nil {
		return nil
	}

	registryClient, err := NewClient()
	if err != nil {
		return errors.New("Failed to create client connection to artifact registry: " + err.Error())
	}
	b.client = registryClient
	return nil
}

// Initialize registers the Bucket b with the configured HCP Packer Registry.
// Upon initialization a Bucket will be upserted to, and new iteration will be created for the build if the configured
// fingerprint has no associated iterations. Lastly, the initialization process with register the builds that need to be
// completed before an iteration can be marked as DONE.
//
// b.Initialize() must be called before any data can be published to the configured HCP Packer Registry.
// TODO ensure initialize can only be called once
func (b *Bucket) Initialize(ctx context.Context) error {

	if err := b.connect(); err != nil {
		return err
	}

	b.Destination = fmt.Sprintf("%s/%s", b.client.OrganizationID, b.client.ProjectID)

	bucketInput := &models.HashicorpCloudPackerCreateBucketRequest{
		BucketSlug:  b.Slug,
		Description: b.Description,
		Labels:      b.Labels,
	}

	err := UpsertBucket(ctx, b.client, bucketInput)
	if err != nil {
		return fmt.Errorf("failed to initialize bucket %q: %w", b.Slug, err)
	}

	return b.initializeIteration(ctx)
}

func (b *Bucket) RegisterBuildForComponent(sourceName string) {
	if b == nil {
		return
	}

	if _, ok := b.Iteration.builds.Load(sourceName); ok {
		return
	}
	b.Iteration.expectedBuilds = append(b.Iteration.expectedBuilds, sourceName)
}

// CreateInitialBuildForIteration will create a build entry on the HCP Packer Registry for the named componentType.
// This initial creation is needed so that Packer can properly track when an iteration is complete.
func (b *Bucket) CreateInitialBuildForIteration(ctx context.Context, componentType string) error {
	status := models.HashicorpCloudPackerBuildStatusUNSET
	buildInput := &models.HashicorpCloudPackerCreateBuildRequest{
		BucketSlug:  b.Slug,
		Fingerprint: b.Iteration.Fingerprint,
		IterationID: b.Iteration.ID,
		Build: &models.HashicorpCloudPackerBuildCreateBody{
			ComponentType: componentType,
			PackerRunUUID: b.Iteration.RunUUID,
			Status:        status,
		},
	}

	id, err := CreateBuild(ctx, b.client, buildInput)
	if err != nil {
		return err
	}

	build := &Build{
		ID:            id,
		ComponentType: componentType,
		RunUUID:       b.Iteration.RunUUID,
		Status:        status,
		Labels:        make(map[string]string),
		Images:        make(map[string]registryimage.Image),
	}

	log.Println("[TRACE] creating initial build for component", componentType)
	b.Iteration.builds.Store(componentType, build)

	return nil
}

// UpdateBuildStatus updates the status of a build entry on the HCP Packer registry with its current local status.
func (b *Bucket) UpdateBuildStatus(ctx context.Context, name string, status models.HashicorpCloudPackerBuildStatus) error {
	if status == models.HashicorpCloudPackerBuildStatusDONE {
		return b.markBuildComplete(ctx, name)
	}

	// Lets check if we have something already for this build
	build, ok := b.Iteration.builds.Load(name)
	if !ok {
		return fmt.Errorf("no build for the component %q associated to the iteration %q", name, b.Iteration.ID)
	}

	buildToUpdate, ok := build.(*Build)
	if !ok {
		return fmt.Errorf("the build for the component %q does not appear to be a valid registry Build", name)
	}

	if buildToUpdate.ID == "" {
		return fmt.Errorf("the build for the component %q does not have a valid id", name)
	}

	buildInput := &models.HashicorpCloudPackerUpdateBuildRequest{
		BuildID: buildToUpdate.ID,
		Updates: &models.HashicorpCloudPackerBuildUpdates{
			PackerRunUUID: buildToUpdate.RunUUID,
			Labels:        buildToUpdate.Labels,
			Status:        status,
		},
	}

	_, err := UpdateBuild(ctx, b.client, buildInput)
	if err != nil {
		return err
	}
	buildToUpdate.Status = status
	b.Iteration.builds.Store(name, buildToUpdate)
	return nil
}

// markBuildComplete should be called to set a build on the HCP Packer registry to DONE.
// Upon a successful call markBuildComplete will publish all images created by the named build,
// and set the registry build to done. A build with no images can not be set to DONE.
func (b *Bucket) markBuildComplete(ctx context.Context, name string) error {
	build, ok := b.Iteration.builds.Load(name)
	if !ok {
		return fmt.Errorf("no build for the component %q associated to the iteration %q", name, b.Iteration.ID)
	}

	buildToUpdate, ok := build.(*Build)
	if !ok {
		return fmt.Errorf("the build for the component %q does not appear to be a valid registry Build", name)
	}

	if buildToUpdate.ID == "" {
		return fmt.Errorf("the build for the component %q does not have a valid id", name)
	}

	status := models.HashicorpCloudPackerBuildStatusDONE

	if buildToUpdate.Status == status {
		// let's no mess with anything that is already done
		return nil
	}

	buildInput := &models.HashicorpCloudPackerUpdateBuildRequest{
		BuildID: buildToUpdate.ID,
		Updates: &models.HashicorpCloudPackerBuildUpdates{
			PackerRunUUID: buildToUpdate.RunUUID,
			Labels:        buildToUpdate.Labels,
			Status:        status,
		},
	}

	if len(buildToUpdate.Images) == 0 {
		return fmt.Errorf("setting a build to DONE with no published images is not currently supported.")
	}

	var providerName, sourceID string
	images := make([]*models.HashicorpCloudPackerImageCreateBody, 0, len(buildToUpdate.Images))
	for _, image := range buildToUpdate.Images {
		// These values will always be the same for all images in a single build,
		// so we can just set it inside the loop without consequence
		if providerName == "" {
			providerName = image.ProviderName
		}
		if image.SourceImageID != "" {
			sourceID = image.SourceImageID
		}

		images = append(images, &models.HashicorpCloudPackerImageCreateBody{ImageID: image.ImageID, Region: image.ProviderRegion})
	}

	buildInput.Updates.CloudProvider = providerName
	buildInput.Updates.SourceImageID = sourceID
	buildInput.Updates.Images = images

	_, err := UpdateBuild(ctx, b.client, buildInput)
	if err != nil {
		return err
	}

	buildToUpdate.Status = status
	b.Iteration.builds.Store(name, buildToUpdate)
	return nil
}

// UpdateImageForBuild appends one or more images artifacts to the build referred to by componentType.
func (b *Bucket) UpdateImageForBuild(componentType string, images ...registryimage.Image) error {
	return b.Iteration.AddImageToBuild(componentType, images...)
}

// UpdateLabelsForBuild merges the contents of data to the labels associated with the build referred to by componentType.
func (b *Bucket) UpdateLabelsForBuild(componentType string, data map[string]string) error {
	return b.Iteration.AddLabelsToBuild(componentType, data)
}

// Load defaults from environment variables
func (b *Bucket) LoadDefaultSettingsFromEnv() {
	// Configure HCP Packer Registry destination
	if b.Slug == "" {
		b.Slug = os.Getenv(env.HCPPackerBucket)
	}

	// Set some iteration values. For Packer RunUUID should always be set.
	// Creating an iteration differently? Let's not overwrite a UUID that might be set.
	if b.Iteration.RunUUID == "" {
		b.Iteration.RunUUID = os.Getenv("PACKER_RUN_UUID")
	}

}

// createIteration creates an empty iteration for a give bucket on the HCP Packer registry.
// The iteration can then be stored locally and used for tracking build status and images for a running
// Packer build.
func (b *Bucket) createIteration() (*models.HashicorpCloudPackerIteration, error) {
	iterationInput := &models.HashicorpCloudPackerCreateIterationRequest{
		BucketSlug:  b.Slug,
		Fingerprint: b.Iteration.Fingerprint,
	}

	iterationResp, err := CreateIteration(context.TODO(), b.client, iterationInput)
	if err != nil {
		return nil, fmt.Errorf("failed to create Iteration for Bucket %s with error: %w", b.Slug, err)
	}

	if iterationResp == nil {
		return nil, fmt.Errorf("failed to create Iteration for Bucket %s with error: %w", b.Slug, err)
	}

	log.Println("[TRACE] a valid iteration for build was created with the Id", iterationResp.ID)
	return iterationResp, nil
}

// initializeIteration populates the bucket iteration with the details needed for tracking builds for a Packer run.
// If an existing Packer registry iteration exists for the said iteration fingerprint, calling initialize on iteration
// that doesn't yet exist will call createIteration to create the entry on the HCP packer registry for the given bucket.
// All build details will be created (if they don't exists) and added to b.Iteration.builds for tracking during runtime.
func (b *Bucket) initializeIteration(ctx context.Context) error {

	// load existing iteration using fingerprint.
	iterationResp, err := GetIteration(ctx, b.client, b.Slug, b.Iteration.Fingerprint)
	if checkErrorCode(err, codes.Aborted) {
		//probably means Iteration doesn't exist need a way to check the error
		iterationResp, err = b.createIteration()
	}

	if err != nil {
		return fmt.Errorf("failed to initialize iteration for fingerprint %s: %s", b.Iteration.Fingerprint, err)
	}

	if iterationResp == nil {
		return fmt.Errorf("failed to initialize iteration details for Bucket %s with error: %w", b.Slug, err)
	}

	log.Println("[TRACE] a valid iteration was retrieved with the id", iterationResp.ID)
	b.Iteration.ID = iterationResp.ID

	// If the iteration is completed and there are no new builds to add, Packer
	// should exit and inform the user that artifacts already exists for the
	// fingerprint associated with the iteration.
	if iterationResp.Complete {
		return fmt.Errorf("This iteration associated to the fingerprint %s is complete. "+
			"If you wish to add a new build to this image a new iteration must be created by changing the build fingerprint.", b.Iteration.Fingerprint)
	}

	// list all this iteration's builds so we can figure out which ones
	// we want to run against. TODO: pagination?
	existingBuilds, err := ListBuilds(ctx, b.client, b.Slug, iterationResp.ID)
	if err != nil {
		return fmt.Errorf("error listing builds for this existing iteration: %s", err)
	}

	var toCreate []string
	for _, expected := range b.Iteration.expectedBuilds {
		var found bool
		for _, existing := range existingBuilds {
			if existing.ComponentType == expected {
				found = true
				build := &Build{
					ID:            existing.ID,
					ComponentType: existing.ComponentType,
					RunUUID:       b.Iteration.RunUUID,
					Status:        existing.Status,
					Labels:        existing.Labels,
				}
				b.Iteration.builds.Store(existing.ComponentType, build)

				// TODO validate that this is safe. For builds that are DONE do we want to keep track of completed things
				// potential issue on updating the status of a build that is already DONE. Is this possible?
				for _, image := range existing.Images {

					err := b.UpdateImageForBuild(existing.ComponentType, registryimage.Image{
						ImageID:        image.ImageID,
						ProviderRegion: image.Region,
					})

					if err != nil {
						log.Printf("[TRACE] unable to load existing images for %q: %v", existing.ComponentType, err)
					}
				}

				log.Printf("[TRACE] a build of component type %s already exists; skipping the create call", expected)
				break
			}
		}

		if !found {
			missingbuild := expected
			toCreate = append(toCreate, missingbuild)
		}
	}

	var errs *multierror.Error
	var wg sync.WaitGroup

	for _, buildName := range toCreate {
		wg.Add(1)
		go func(name string) {
			defer wg.Done()
			log.Printf("[TRACE] registering build with iteration for %q.", name)
			// Need a way to handle skipping builds that were already created.
			// TODO when we load an existing iteration we will probably have a build Id so we should skip.
			// we also need to bubble up the errors here.
			err := b.CreateInitialBuildForIteration(ctx, name)
			if checkErrorCode(err, codes.AlreadyExists) {
				// Check whether build is complete, and if so, skip it.
				log.Printf("[TRACE] build %s already exists in Packer registry, continuing...", name)
				return
			}

			errs = multierror.Append(errs, err)
		}(buildName)
	}
	wg.Wait()

	return errs.ErrorOrNil()
}

// IsExpectingBuildForComponent returns true if the component referenced by buildName is part of the iteration
// and is not marked as DONE on the HCP Packer registry.
func (b *Bucket) IsExpectingBuildForComponent(buildName string) bool {

	v, ok := b.Iteration.builds.Load(buildName)
	if !ok {
		return false
	}

	build := v.(*Build)
	hasBuildID := build.ID != ""
	hasImages := len(build.Images) == 0
	isNotDone := build.Status != models.HashicorpCloudPackerBuildStatusDONE

	return hasBuildID && hasImages && isNotDone
}
<|endoftext|>


<commit_before>package game

import (
	"github.com/MichaelDiBernardo/srl/lib/math"
	"testing"
)

var lTestActor = &Spec{
	Family:  FamActor,
	Genus:   GenMonster,
	Species: "TestSpecies",
	Name:    "Hi",
	Traits:  &Traits{},
}

var lTestItem = &Spec{
	Family:  FamItem,
	Genus:   GenEquip,
	Species: "TestSpecies",
	Name:    "Hiiii",
	Traits:  &Traits{},
}

func TestOkPlaceActor(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, IdentLevel)
	obj := newObj(lTestActor)
	pos := math.Pt(1, 1)

	ok := l.Place(obj, pos)

	if !ok {
		t.Error(`Place((1, 1) was false, want true`)
	}

	if l.At(pos).Actor != obj {
		t.Error(`Place((1, 1)) did not set tile actor to obj`)
	}

	if l.At(pos) != obj.Tile {
		t.Error(`Place((1, 1)) did not set actor's tile to obj`)
	}
}

func TestSecondPlaceActorCleansUp(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, IdentLevel)
	obj := g.NewObj(lTestActor)
	startpos := math.Pt(1, 1)
	endpos := math.Pt(2, 2)

	l.Place(obj, startpos)
	l.Place(obj, endpos)

	if l.At(startpos).Actor != nil {
		t.Error(`Place((2, 2)) did not set (1, 1) tile actor to nil`)
	}
	if l.At(endpos).Actor != obj {
		t.Error(`Place((2, 2)) did not set tile actor to obj`)
	}
}

func TestBadActorPlaceOntoSolid(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, SquareLevel)
	obj := g.NewObj(lTestActor)
	pos := math.Pt(0, 0)

	ok := l.Place(obj, pos)

	if ok {
		t.Error(`Place( (0,0) ) onto FeatWall ok was true; want false`)
	}
}

func TestBadPlaceActorOntoOccupiedTile(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, IdentLevel)
	a1, a2 := g.NewObj(lTestActor), g.NewObj(lTestActor)
	pos := math.Pt(0, 0)

	l.Place(a1, pos)
	ok := l.Place(a2, pos)

	if ok {
		t.Error(`Place onto other actor: ok was true; want false`)
	}
}

func TestPlaceAddsActorToList(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, nil, IdentLevel)
	obj := g.NewObj(lTestActor)
	startpos := math.Pt(1, 1)

	l.Place(obj, startpos)

	if actual := len(l.actors); actual != 1 {
		t.Errorf(`Place(obj) put %d actors in list; want 1`, actual)
	}
	if actual := l.actors[0]; actual != obj {
		t.Errorf(`Place(obj) put %v in list; want %v`, actual, obj)
	}
}

func TestBadPlaceDoesNotAddActorToList(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, SquareLevel)
	obj := g.NewObj(lTestActor)
	startpos := math.Pt(0, 0)

	l.Place(obj, startpos)

	if actual := len(l.actors); actual != 0 {
		t.Errorf(`Place(obj) put %d actors in list; want 0`, actual)
	}
}

func TestPlaceSingleItem(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, SquareLevel)
	obj := g.NewObj(lTestItem)
	dest := math.Pt(1, 1)

	l.Place(obj, dest)

	items := l.At(dest).Items

	if size := items.Len(); size != 1 {
		t.Errorf(`Place(item) put %d items; want 1`)
	}

	if item := items.Top(); item != obj {
		t.Errorf(`Place(item) was %v, want %v`, item, obj)
	}
}

func TestOkRemoveActor(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, IdentLevel)
	obj := newObj(lTestActor)
	pos := math.Pt(1, 1)

	l.Place(obj, pos)
	l.Remove(obj)

	if obj.Level != nil {
		t.Error(`obj.Level was not nil`)
	}
	if obj.Tile != nil {
		t.Error(`obj.Tile was not nil`)
	}
	if l.At(pos).Actor != nil {
		t.Error(`Actor's previous tile had tile.Actor != nil`)
	}
	if len(l.actors) > 0 {
		t.Error(`l.actors was not empty after removal.`)
	}
}
<commit_msg>Thanks `go vet`<commit_after>package game

import (
	"github.com/MichaelDiBernardo/srl/lib/math"
	"testing"
)

var lTestActor = &Spec{
	Family:  FamActor,
	Genus:   GenMonster,
	Species: "TestSpecies",
	Name:    "Hi",
	Traits:  &Traits{},
}

var lTestItem = &Spec{
	Family:  FamItem,
	Genus:   GenEquip,
	Species: "TestSpecies",
	Name:    "Hiiii",
	Traits:  &Traits{},
}

func TestOkPlaceActor(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, IdentLevel)
	obj := newObj(lTestActor)
	pos := math.Pt(1, 1)

	ok := l.Place(obj, pos)

	if !ok {
		t.Error(`Place((1, 1) was false, want true`)
	}

	if l.At(pos).Actor != obj {
		t.Error(`Place((1, 1)) did not set tile actor to obj`)
	}

	if l.At(pos) != obj.Tile {
		t.Error(`Place((1, 1)) did not set actor's tile to obj`)
	}
}

func TestSecondPlaceActorCleansUp(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, IdentLevel)
	obj := g.NewObj(lTestActor)
	startpos := math.Pt(1, 1)
	endpos := math.Pt(2, 2)

	l.Place(obj, startpos)
	l.Place(obj, endpos)

	if l.At(startpos).Actor != nil {
		t.Error(`Place((2, 2)) did not set (1, 1) tile actor to nil`)
	}
	if l.At(endpos).Actor != obj {
		t.Error(`Place((2, 2)) did not set tile actor to obj`)
	}
}

func TestBadActorPlaceOntoSolid(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, SquareLevel)
	obj := g.NewObj(lTestActor)
	pos := math.Pt(0, 0)

	ok := l.Place(obj, pos)

	if ok {
		t.Error(`Place( (0,0) ) onto FeatWall ok was true; want false`)
	}
}

func TestBadPlaceActorOntoOccupiedTile(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, IdentLevel)
	a1, a2 := g.NewObj(lTestActor), g.NewObj(lTestActor)
	pos := math.Pt(0, 0)

	l.Place(a1, pos)
	ok := l.Place(a2, pos)

	if ok {
		t.Error(`Place onto other actor: ok was true; want false`)
	}
}

func TestPlaceAddsActorToList(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, nil, IdentLevel)
	obj := g.NewObj(lTestActor)
	startpos := math.Pt(1, 1)

	l.Place(obj, startpos)

	if actual := len(l.actors); actual != 1 {
		t.Errorf(`Place(obj) put %d actors in list; want 1`, actual)
	}
	if actual := l.actors[0]; actual != obj {
		t.Errorf(`Place(obj) put %v in list; want %v`, actual, obj)
	}
}

func TestBadPlaceDoesNotAddActorToList(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, SquareLevel)
	obj := g.NewObj(lTestActor)
	startpos := math.Pt(0, 0)

	l.Place(obj, startpos)

	if actual := len(l.actors); actual != 0 {
		t.Errorf(`Place(obj) put %d actors in list; want 0`, actual)
	}
}

func TestPlaceSingleItem(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, SquareLevel)
	obj := g.NewObj(lTestItem)
	dest := math.Pt(1, 1)

	l.Place(obj, dest)

	items := l.At(dest).Items

	if size := items.Len(); size != 1 {
		t.Errorf(`Place(item) put %d items; want 1`, size)
	}

	if item := items.Top(); item != obj {
		t.Errorf(`Place(item) was %v, want %v`, item, obj)
	}
}

func TestOkRemoveActor(t *testing.T) {
	g := NewGame()
	l := NewLevel(4, 4, g, IdentLevel)
	obj := newObj(lTestActor)
	pos := math.Pt(1, 1)

	l.Place(obj, pos)
	l.Remove(obj)

	if obj.Level != nil {
		t.Error(`obj.Level was not nil`)
	}
	if obj.Tile != nil {
		t.Error(`obj.Tile was not nil`)
	}
	if l.At(pos).Actor != nil {
		t.Error(`Actor's previous tile had tile.Actor != nil`)
	}
	if len(l.actors) > 0 {
		t.Error(`l.actors was not empty after removal.`)
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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
//
//     https://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 bigtable

import (
	"context"

	cbt "cloud.google.com/go/bigtable"
	"github.com/datacommonsorg/mixer/internal/util"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type chanData struct {
	// Identifier of the data when reading multiple BigTable rows
	// concurrently. It could be place dcid, statvar dcid or other "key".
	token string
	// Data read from Cloud Bigtable
	data interface{}
}

// readRowFn generates a function to be used as the callback function in Bigtable Read.
// This utilizes the Golang closure so the arguments can be scoped in the
// generated function.
func readRowFn(
	errCtx context.Context,
	btTable *cbt.Table,
	rowSetPart cbt.RowSet,
	getToken func(string) (string, error),
	action func([]byte) (interface{}, error),
	elemChan chan chanData,
) func() error {
	return func() error {
		if err := btTable.ReadRows(errCtx, rowSetPart,
			func(btRow cbt.Row) bool {
				if len(btRow[BtFamily]) == 0 {
					return true
				}
				raw := btRow[BtFamily][0].Value

				if getToken == nil {
					getToken = util.KeyToDcid
				}
				token, err := getToken(btRow.Key())
				if err != nil {
					return false
				}
				jsonRaw, err := util.UnzipAndDecode(string(raw))
				if err != nil {
					return false
				}
				elem, err := action(jsonRaw)
				if err != nil {
					return false
				}
				elemChan <- chanData{token, elem}
				return true
			}); err != nil {
			return err
		}
		return nil
	}
}

// Read reads BigTable rows from multiple Bigtable in parallel.
//
// Reading multiple rows is chunked as the size limit for RowSet is 500KB.
//
// Args:
// baseBt: The bigtable that holds the base cache
// rowSet: BigTable rowSet containing the row keys.
// action: A callback function that converts the raw bytes into appropriate
//		go struct based on the cache content.
// getToken: A function to get back the indexed token (like place dcid) from
//		bigtable row key.
//
func Read(
	ctx context.Context,
	btGroup *Group,
	rowSet cbt.RowSet,
	action func([]byte) (interface{}, error),
	getToken func(string) (string, error),
) ([]map[string]interface{}, error) {
	tables := btGroup.Tables()
	if len(tables) == 0 {
		return nil, status.Errorf(
			codes.NotFound, "Bigtable instance is not specified")
	}

	// Function start
	var rowSetSize int
	var rowList cbt.RowList
	var rowRangeList cbt.RowRangeList
	switch v := rowSet.(type) {
	case cbt.RowList:
		rowList = rowSet.(cbt.RowList)
		rowSetSize = len(rowList)
	case cbt.RowRangeList:
		rowRangeList = rowSet.(cbt.RowRangeList)
		rowSetSize = len(rowRangeList)
	default:
		return nil, status.Errorf(codes.Internal, "Unsupported RowSet type: %v", v)
	}
	if rowSetSize == 0 {
		return nil, nil
	}

	chans := make(map[int]chan chanData)
	for i := 0; i < len(tables); i++ {
		chans[i] = make(chan chanData, rowSetSize)
	}
	errs, errCtx := errgroup.WithContext(ctx)
	for i := 0; i <= rowSetSize/BtBatchQuerySize; i++ {
		left := i * BtBatchQuerySize
		right := (i + 1) * BtBatchQuerySize
		if right > rowSetSize {
			right = rowSetSize
		}
		var rowSetPart cbt.RowSet
		if len(rowList) > 0 {
			rowSetPart = rowList[left:right]
		} else {
			rowSetPart = rowRangeList[left:right]
		}
		// Read from all the given tables.
		if len(tables) > 0 {
			for j := 0; j < len(tables); j++ {
				j := j
				if tables[j] != nil {
					errs.Go(readRowFn(errCtx, tables[j], rowSetPart, getToken, action, chans[j]))
				}
			}
		}
	}
	err := errs.Wait()
	if err != nil {
		return nil, err
	}
	for i := 0; i < len(chans); i++ {
		close(chans[i])
	}

	result := []map[string]interface{}{}
	if tables != nil {
		for i := 0; i < len(tables); i++ {
			item := map[string]interface{}{}
			for elem := range chans[i] {
				item[elem.token] = elem.data
			}
			result = append(result, item)
		}
	}
	return result, nil
}
<commit_msg>Adapt Bigtable Reader to take distinct keys for each import group table (#838)<commit_after>// Copyright 2020 Google LLC
//
// 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
//
//     https://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 bigtable

import (
	"context"

	cbt "cloud.google.com/go/bigtable"
	"github.com/datacommonsorg/mixer/internal/util"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type chanData struct {
	// Identifier of the data when reading multiple BigTable rows
	// concurrently. It could be place dcid, statvar dcid or other "key".
	token string
	// Data read from Cloud Bigtable
	data interface{}
}

// readRowFn generates a function to be used as the callback function in Bigtable Read.
// This utilizes the Golang closure so the arguments can be scoped in the
// generated function.
func readRowFn(
	errCtx context.Context,
	btTable *cbt.Table,
	rowSetPart cbt.RowSet,
	getToken func(string) (string, error),
	action func([]byte) (interface{}, error),
	elemChan chan chanData,
) func() error {
	return func() error {
		if err := btTable.ReadRows(errCtx, rowSetPart,
			func(btRow cbt.Row) bool {
				if len(btRow[BtFamily]) == 0 {
					return true
				}
				raw := btRow[BtFamily][0].Value

				if getToken == nil {
					getToken = util.KeyToDcid
				}
				token, err := getToken(btRow.Key())
				if err != nil {
					return false
				}
				jsonRaw, err := util.UnzipAndDecode(string(raw))
				if err != nil {
					return false
				}
				elem, err := action(jsonRaw)
				if err != nil {
					return false
				}
				elemChan <- chanData{token, elem}
				return true
			}); err != nil {
			return err
		}
		return nil
	}
}

// Read reads BigTable rows from multiple Bigtable in parallel.
// Note all Bigtable read use the same set of rowList.
func Read(
	ctx context.Context,
	btGroup *Group,
	rowList cbt.RowList,
	action func([]byte) (interface{}, error),
	getToken func(string) (string, error),
) ([]map[string]interface{}, error) {
	rowListMap := map[int]cbt.RowList{}
	for i := 0; i < len(btGroup.Tables()); i++ {
		rowListMap[i] = rowList
	}
	return ReadWithGroupRowList(ctx, btGroup, rowListMap, action, getToken)
}

// ReadWithGroupRowList reads BigTable rows from multiple Bigtable in parallel.
// Reading is chunked as the size limit for RowSet is 500KB.
//
// Note the read could have different RowList for each import group Bigtable as
// needed by the pagination APIs.
func ReadWithGroupRowList(
	ctx context.Context,
	btGroup *Group,
	rowListMap map[int]cbt.RowList,
	action func([]byte) (interface{}, error),
	getToken func(string) (string, error),
) ([]map[string]interface{}, error) {
	tables := btGroup.Tables()
	if len(tables) == 0 {
		return nil, status.Errorf(codes.NotFound, "Bigtable instance is not specified")
	}
	if len(tables) != len(rowListMap) {
		return nil, status.Errorf(codes.Internal, "Number of tables and rowList don't match")
	}
	// Channels for each import group read.
	chans := make(map[int]chan chanData)
	for i := 0; i < len(tables); i++ {
		chans[i] = make(chan chanData, len(rowListMap[i]))
	}

	errs, errCtx := errgroup.WithContext(ctx)
	// Read from each import group tables. Note each table could have different
	// rowList in pagination APIs.
	for i := 0; i < len(tables); i++ {
		rowSet := rowListMap[i]
		rowSetSize := len(rowSet)
		for j := 0; j <= rowSetSize/BtBatchQuerySize; j++ {
			left := j * BtBatchQuerySize
			right := (j + 1) * BtBatchQuerySize
			if right > rowSetSize {
				right = rowSetSize
			}
			rowSetPart := rowSet[left:right]
			if tables[i] != nil {
				errs.Go(readRowFn(errCtx, tables[i], rowSetPart, getToken, action, chans[i]))
			}
		}
	}
	err := errs.Wait()
	if err != nil {
		return nil, err
	}
	for i := 0; i < len(chans); i++ {
		close(chans[i])
	}

	result := []map[string]interface{}{}
	if tables != nil {
		for i := 0; i < len(tables); i++ {
			item := map[string]interface{}{}
			for elem := range chans[i] {
				item[elem.token] = elem.data
			}
			result = append(result, item)
		}
	}
	return result, nil
}
<|endoftext|>


<commit_before>package resolvers

import (
	"context"
	"fmt"
	"log"
	"math"
	"time"

	"github.com/emwalker/digraph/loaders"
	"github.com/emwalker/digraph/models"
	"github.com/volatiletech/sqlboiler/boil"
	"github.com/volatiletech/sqlboiler/queries"
	"github.com/volatiletech/sqlboiler/queries/qm"
)

type topicResolver struct {
	*Resolver
}

func getTopicLoader(ctx context.Context) *loaders.TopicLoader {
	return ctx.Value(loaders.TopicLoaderKey).(*loaders.TopicLoader)
}

func topicConnection(rows []*models.Topic, err error) (models.TopicConnection, error) {
	if err != nil {
		return models.TopicConnection{}, err
	}

	var edges []*models.TopicEdge
	for _, topic := range rows {
		edges = append(edges, &models.TopicEdge{Node: models.TopicValue{topic, false}})
	}

	return models.TopicConnection{Edges: edges}, nil
}

func topicRepository(ctx context.Context, topic *models.TopicValue) (*models.Repository, error) {
	if topic.R != nil && topic.R.Repository != nil {
		return topic.R.Repository, nil
	}
	repo, err := fetchRepository(ctx, topic.RepositoryID)
	return &repo, err
}

func topicOrganization(ctx context.Context, topic *models.TopicValue) (*models.Organization, error) {
	if topic.R != nil && topic.R.Organization != nil {
		return topic.R.Organization, nil
	}
	org, err := fetchOrganization(ctx, topic.OrganizationID)
	return &org, err
}

func availableTopics(
	ctx context.Context, exec boil.ContextExecutor, actor *models.User, first *int,
) (models.TopicConnection, error) {
	var topics []*models.Topic
	err := queries.Raw(`
	select t.* from topics t
	inner join organizations o on o.id = t.organization_id
	inner join organization_members om on om.organization_id = o.id
	inner join users u on om.user_id = u.id
	where u.id = $1
	order by t.name
	limit $2
	`, actor.ID, limitFrom(first)).Bind(ctx, exec, &topics)

	if err != nil {
		return models.TopicConnection{}, err
	}

	return topicConnection(topics, err)
}

// AvailableParentTopics returns the topics that can be added to the link.
func (r *topicResolver) AvailableParentTopics(
	ctx context.Context, topic *models.TopicValue, first *int, after *string, last *int, before *string,
) (models.TopicConnection, error) {
	return availableTopics(ctx, r.DB, r.Actor, first)
}

// ChildTopics returns a set of topics.
func (r *topicResolver) ChildTopics(
	ctx context.Context, topic *models.TopicValue, searchString *string, first *int, after *string,
	last *int, before *string,
) (models.TopicConnection, error) {
	mods := []qm.QueryMod{
		qm.Load("ParentTopics"),
		qm.OrderBy("name"),
	}

	if searchString != nil && *searchString != "" {
		mods = append(mods, qm.Where("name ilike '%%' || ? || '%%'", *searchString))
	}

	topics, err := topic.ChildTopics(mods...).All(ctx, r.DB)
	return topicConnection(topics, err)
}

// CreatedAt returns the time of the topic's creation.
func (r *topicResolver) CreatedAt(_ context.Context, topic *models.TopicValue) (string, error) {
	return topic.CreatedAt.Format(time.RFC3339), nil
}

// Description returns a description of the topic.
func (r *topicResolver) Description(_ context.Context, topic *models.TopicValue) (*string, error) {
	return topic.Description.Ptr(), nil
}

func (r *topicResolver) ID(_ context.Context, topic *models.TopicValue) (string, error) {
	return topic.ID, nil
}

// Links returns a set of links.
func (r *topicResolver) Links(
	ctx context.Context, topic *models.TopicValue, searchString *string, first *int, after *string,
	last *int, before *string,
) (models.LinkConnection, error) {
	mods := []qm.QueryMod{
		qm.Load("ParentTopics"),
		qm.OrderBy("created_at desc"),
	}

	if searchString != nil && *searchString != "" {
		mods = append(mods, qm.Where("title ilike '%%' || ? || '%%'", searchString))
	}

	scope := topic.ChildLinks(mods...)
	return linkConnection(scope.All(ctx, r.DB))
}

func (r *topicResolver) Name(_ context.Context, topic *models.TopicValue) (string, error) {
	return topic.Name, nil
}

// NewlyAdded returns true if the topic was just added.
func (r *topicResolver) NewlyAdded(_ context.Context, topic *models.TopicValue) (bool, error) {
	return topic.NewlyAdded, nil
}

// Organization returns an organization.
func (r *topicResolver) Organization(
	ctx context.Context, topic *models.TopicValue,
) (models.Organization, error) {
	org, err := topicOrganization(ctx, topic)
	return *org, err
}

// ParentTopics returns a set of topics.
func (r *topicResolver) ParentTopics(
	ctx context.Context, topic *models.TopicValue, first *int, after *string, last *int, before *string,
) (models.TopicConnection, error) {
	if topic.R != nil && topic.R.ParentTopics != nil {
		return topicConnection(topic.R.ParentTopics, nil)
	}

	log.Printf("Fetching parent topics for topic %s", topic.ID)
	mods := []qm.QueryMod{
		qm.OrderBy("name"),
	}

	topics, err := topic.ParentTopics(mods...).All(ctx, r.DB)
	return topicConnection(topics, err)
}

// Repository returns the repostory of the topic.
func (r *topicResolver) Repository(
	ctx context.Context, topic *models.TopicValue,
) (models.Repository, error) {
	repo, err := topicRepository(ctx, topic)
	return *repo, err
}

// ResourcePath returns a path to the item.
func (r *topicResolver) ResourcePath(ctx context.Context, topic *models.TopicValue) (string, error) {
	repo, err := topicRepository(ctx, topic)
	if err != nil {
		return "", err
	}

	org, err := topicOrganization(ctx, topic)
	if err != nil {
		return "", err
	}

	if repo.System {
		return fmt.Sprintf("/%s/topics/%s", org.Login, topic.ID), nil
	}
	return fmt.Sprintf("/%s/%s/topics/%s", org.Login, repo.Name, topic.ID), nil
}

func (r *topicResolver) matchingDescendantTopics(
	ctx context.Context, topic *models.TopicValue, searchString string, limit int,
) ([]*models.TopicValue, error) {
	var rows []struct {
		ID string
	}

	err := queries.Raw(`
	with recursive child_topics as (
	  select parent_id, parent_id as child_id
	  from topic_topics where parent_id = $1
	union
	  select pt.child_id, ct.child_id
	  from topic_topics ct
	  inner join child_topics pt on pt.child_id = ct.parent_id
	)
	select t.id from topics t
	inner join child_topics ct on ct.child_id = t.id
	where t.name ilike '%%' || $2 || '%%'
	limit $3
	`, topic.ID, searchString, limit).Bind(ctx, r.DB, &rows)

	if err != nil {
		return nil, err
	}

	if len(rows) < 1 {
		return nil, nil
	}

	var ids []interface{}
	for _, row := range rows {
		ids = append(ids, row.ID)
	}

	topics, err := models.Topics(qm.WhereIn("id in ?", ids...)).All(ctx, r.DB)
	if err != nil {
		return nil, err
	}

	var topicValues []*models.TopicValue
	for _, t := range topics {
		topicValues = append(topicValues, &models.TopicValue{t, false})
	}

	return topicValues, nil
}

func (r *topicResolver) matchingDescendantLinks(
	ctx context.Context, topic *models.TopicValue, searchString string, limit int,
) ([]*models.Link, error) {
	var rows []struct {
		ID string
	}

	err := queries.Raw(`
	with recursive child_topics as (
	  select parent_id, parent_id as child_id
	  from topic_topics where parent_id = $1
	union
	  select pt.child_id, ct.child_id
	  from topic_topics ct
	  inner join child_topics pt on pt.child_id = ct.parent_id
	)
	select l.id from links l
	inner join link_topics lt on l.id = lt.child_id
	inner join child_topics ct on ct.child_id = lt.parent_id
	where l.title ilike '%%' || $2 || '%%'
	limit $3
	`, topic.ID, searchString, limit).Bind(ctx, r.DB, &rows)

	if err != nil {
		return nil, err
	}

	if len(rows) < 1 {
		return nil, nil
	}

	var ids []interface{}
	for _, row := range rows {
		ids = append(ids, row.ID)
	}

	return models.Links(
		qm.Load("ParentTopics"),
		qm.WhereIn("id in ?", ids...),
	).All(ctx, r.DB)
}

func (r *topicResolver) Search(
	ctx context.Context, topic *models.TopicValue, searchString string, first *int, after *string,
	last *int, before *string,
) (models.SearchResultItemConnection, error) {
	log.Printf("Searching topic %s for '%s'", topic.ID, searchString)

	var (
		err    error
		topics []*models.TopicValue
		links  []*models.Link
		edges  []*models.SearchResultItemEdge
	)

	var limit = math.MaxInt32
	if first != nil {
		limit = *first
	}

	if topics, err = r.matchingDescendantTopics(ctx, topic, searchString, limit); err != nil {
		return models.SearchResultItemConnection{}, err
	}

	limit -= len(topics)
	if links, err = r.matchingDescendantLinks(ctx, topic, searchString, limit); err != nil {
		return models.SearchResultItemConnection{}, err
	}

	for _, topic := range topics {
		edges = append(edges, &models.SearchResultItemEdge{Node: *topic})
	}

	for _, link := range links {
		edges = append(edges, &models.SearchResultItemEdge{Node: *link})
	}

	return models.SearchResultItemConnection{Edges: edges}, nil
}

// UpdatedAt returns the time of the most recent update.
func (r *topicResolver) UpdatedAt(_ context.Context, topic *models.TopicValue) (string, error) {
	return topic.UpdatedAt.Format(time.RFC3339), nil
}
<commit_msg>Fix a bug introduced in the previous commit<commit_after>package resolvers

import (
	"context"
	"fmt"
	"log"
	"math"
	"time"

	"github.com/emwalker/digraph/loaders"
	"github.com/emwalker/digraph/models"
	"github.com/volatiletech/sqlboiler/boil"
	"github.com/volatiletech/sqlboiler/queries"
	"github.com/volatiletech/sqlboiler/queries/qm"
)

type topicResolver struct {
	*Resolver
}

func getTopicLoader(ctx context.Context) *loaders.TopicLoader {
	return ctx.Value(loaders.TopicLoaderKey).(*loaders.TopicLoader)
}

func topicConnection(rows []*models.Topic, err error) (models.TopicConnection, error) {
	if err != nil {
		return models.TopicConnection{}, err
	}

	var edges []*models.TopicEdge
	for _, topic := range rows {
		edges = append(edges, &models.TopicEdge{Node: models.TopicValue{topic, false}})
	}

	return models.TopicConnection{Edges: edges}, nil
}

func topicRepository(ctx context.Context, topic *models.TopicValue) (*models.Repository, error) {
	if topic.R != nil && topic.R.Repository != nil {
		return topic.R.Repository, nil
	}
	repo, err := fetchRepository(ctx, topic.RepositoryID)
	return &repo, err
}

func topicOrganization(ctx context.Context, topic *models.TopicValue) (*models.Organization, error) {
	if topic.R != nil && topic.R.Organization != nil {
		return topic.R.Organization, nil
	}
	org, err := fetchOrganization(ctx, topic.OrganizationID)
	return &org, err
}

func availableTopics(
	ctx context.Context, exec boil.ContextExecutor, actor *models.User, first *int,
) (models.TopicConnection, error) {
	var topics []*models.Topic
	err := queries.Raw(`
	select t.* from topics t
	inner join organizations o on o.id = t.organization_id
	inner join organization_members om on om.organization_id = o.id
	inner join users u on om.user_id = u.id
	where u.id = $1
	order by t.name
	limit $2
	`, actor.ID, limitFrom(first)).Bind(ctx, exec, &topics)

	if err != nil {
		return models.TopicConnection{}, err
	}

	return topicConnection(topics, err)
}

// AvailableParentTopics returns the topics that can be added to the link.
func (r *topicResolver) AvailableParentTopics(
	ctx context.Context, topic *models.TopicValue, first *int, after *string, last *int, before *string,
) (models.TopicConnection, error) {
	return availableTopics(ctx, r.DB, r.Actor, first)
}

// ChildTopics returns a set of topics.
func (r *topicResolver) ChildTopics(
	ctx context.Context, topic *models.TopicValue, searchString *string, first *int, after *string,
	last *int, before *string,
) (models.TopicConnection, error) {
	mods := []qm.QueryMod{
		qm.Load("ParentTopics"),
		qm.OrderBy("name"),
	}

	if searchString != nil && *searchString != "" {
		mods = append(mods, qm.Where("name ilike '%%' || ? || '%%'", *searchString))
	}

	topics, err := topic.ChildTopics(mods...).All(ctx, r.DB)
	return topicConnection(topics, err)
}

// CreatedAt returns the time of the topic's creation.
func (r *topicResolver) CreatedAt(_ context.Context, topic *models.TopicValue) (string, error) {
	return topic.CreatedAt.Format(time.RFC3339), nil
}

// Description returns a description of the topic.
func (r *topicResolver) Description(_ context.Context, topic *models.TopicValue) (*string, error) {
	return topic.Description.Ptr(), nil
}

func (r *topicResolver) ID(_ context.Context, topic *models.TopicValue) (string, error) {
	return topic.ID, nil
}

// Links returns a set of links.
func (r *topicResolver) Links(
	ctx context.Context, topic *models.TopicValue, searchString *string, first *int, after *string,
	last *int, before *string,
) (models.LinkConnection, error) {
	mods := []qm.QueryMod{
		qm.Load("ParentTopics"),
		qm.OrderBy("created_at desc"),
	}

	if searchString != nil && *searchString != "" {
		mods = append(mods, qm.Where("title ilike '%%' || ? || '%%'", searchString))
	}

	scope := topic.ChildLinks(mods...)
	return linkConnection(scope.All(ctx, r.DB))
}

func (r *topicResolver) Name(_ context.Context, topic *models.TopicValue) (string, error) {
	return topic.Name, nil
}

// NewlyAdded returns true if the topic was just added.
func (r *topicResolver) NewlyAdded(_ context.Context, topic *models.TopicValue) (bool, error) {
	return topic.NewlyAdded, nil
}

// Organization returns an organization.
func (r *topicResolver) Organization(
	ctx context.Context, topic *models.TopicValue,
) (models.Organization, error) {
	org, err := topicOrganization(ctx, topic)
	return *org, err
}

// ParentTopics returns a set of topics.
func (r *topicResolver) ParentTopics(
	ctx context.Context, topic *models.TopicValue, first *int, after *string, last *int, before *string,
) (models.TopicConnection, error) {
	if topic.R != nil && topic.R.ParentTopics != nil {
		return topicConnection(topic.R.ParentTopics, nil)
	}

	log.Printf("Fetching parent topics for topic %s", topic.ID)
	mods := []qm.QueryMod{
		qm.OrderBy("name"),
	}

	topics, err := topic.ParentTopics(mods...).All(ctx, r.DB)
	return topicConnection(topics, err)
}

// Repository returns the repostory of the topic.
func (r *topicResolver) Repository(
	ctx context.Context, topic *models.TopicValue,
) (models.Repository, error) {
	repo, err := topicRepository(ctx, topic)
	return *repo, err
}

// ResourcePath returns a path to the item.
func (r *topicResolver) ResourcePath(ctx context.Context, topic *models.TopicValue) (string, error) {
	repo, err := topicRepository(ctx, topic)
	if err != nil {
		return "", err
	}

	org, err := topicOrganization(ctx, topic)
	if err != nil {
		return "", err
	}

	if repo.System {
		return fmt.Sprintf("/%s/topics/%s", org.Login, topic.ID), nil
	}
	return fmt.Sprintf("/%s/%s/topics/%s", org.Login, repo.Name, topic.ID), nil
}

func (r *topicResolver) matchingDescendantTopics(
	ctx context.Context, topic *models.TopicValue, searchString string, limit int,
) ([]*models.TopicValue, error) {
	var rows []struct {
		ID string
	}

	err := queries.Raw(`
	with recursive child_topics as (
	  select parent_id, parent_id as child_id
	  from topic_topics where parent_id = $1
	union
	  select pt.child_id, ct.child_id
	  from topic_topics ct
	  inner join child_topics pt on pt.child_id = ct.parent_id
	)
	select t.id from topics t
	inner join child_topics ct on ct.child_id = t.id
	where t.name ilike '%%' || $2 || '%%'
	limit $3
	`, topic.ID, searchString, limit).Bind(ctx, r.DB, &rows)

	if err != nil {
		return nil, err
	}

	if len(rows) < 1 {
		return nil, nil
	}

	var ids []interface{}
	for _, row := range rows {
		ids = append(ids, row.ID)
	}

	topics, err := models.Topics(qm.WhereIn("id in ?", ids...)).All(ctx, r.DB)
	if err != nil {
		return nil, err
	}

	var topicValues []*models.TopicValue
	for _, t := range topics {
		topicValues = append(topicValues, &models.TopicValue{t, false})
	}

	return topicValues, nil
}

func (r *topicResolver) matchingDescendantLinks(
	ctx context.Context, topic *models.TopicValue, searchString string, limit int,
) ([]*models.LinkValue, error) {
	var rows []struct {
		ID string
	}

	err := queries.Raw(`
	with recursive child_topics as (
	  select parent_id, parent_id as child_id
	  from topic_topics where parent_id = $1
	union
	  select pt.child_id, ct.child_id
	  from topic_topics ct
	  inner join child_topics pt on pt.child_id = ct.parent_id
	)
	select l.id from links l
	inner join link_topics lt on l.id = lt.child_id
	inner join child_topics ct on ct.child_id = lt.parent_id
	where l.title ilike '%%' || $2 || '%%'
	limit $3
	`, topic.ID, searchString, limit).Bind(ctx, r.DB, &rows)

	if err != nil {
		return nil, err
	}

	if len(rows) < 1 {
		return nil, nil
	}

	var ids []interface{}
	for _, row := range rows {
		ids = append(ids, row.ID)
	}

	links, err := models.Links(
		qm.Load("ParentTopics"),
		qm.WhereIn("id in ?", ids...),
	).All(ctx, r.DB)
	if err != nil {
		return nil, err
	}

	var linkValues []*models.LinkValue
	for _, l := range links {
		linkValues = append(linkValues, &models.LinkValue{l, false})
	}

	return linkValues, nil
}

func (r *topicResolver) Search(
	ctx context.Context, topic *models.TopicValue, searchString string, first *int, after *string,
	last *int, before *string,
) (models.SearchResultItemConnection, error) {
	log.Printf("Searching topic %s for '%s'", topic.ID, searchString)

	var (
		err    error
		topics []*models.TopicValue
		links  []*models.LinkValue
		edges  []*models.SearchResultItemEdge
	)

	var limit = math.MaxInt32
	if first != nil {
		limit = *first
	}

	if topics, err = r.matchingDescendantTopics(ctx, topic, searchString, limit); err != nil {
		return models.SearchResultItemConnection{}, err
	}

	limit -= len(topics)
	if links, err = r.matchingDescendantLinks(ctx, topic, searchString, limit); err != nil {
		return models.SearchResultItemConnection{}, err
	}

	for _, topic := range topics {
		edges = append(edges, &models.SearchResultItemEdge{Node: *topic})
	}

	for _, link := range links {
		edges = append(edges, &models.SearchResultItemEdge{Node: *link})
	}

	return models.SearchResultItemConnection{Edges: edges}, nil
}

// UpdatedAt returns the time of the most recent update.
func (r *topicResolver) UpdatedAt(_ context.Context, topic *models.TopicValue) (string, error) {
	return topic.UpdatedAt.Format(time.RFC3339), nil
}
<|endoftext|>