Repository: ogen-go/ogen
Pull request #734 (2022-12-29 19:23:16)

fix(sgcollector): do not put entire query result into the memory

codecov[bot]: # [Codecov](https://codecov.io/gh/ogen-go/ogen/pull/734?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go) Report
Base: **74.31**% // Head: **74.01**% // Decreases project coverage by **`-0.30%`** :warning:
> Coverage data is based on head [(`bf99528`)](https://codecov.io/gh/ogen-go/ogen/pull/734?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go) compared to base [(`381d15e`)](https://codecov.io/gh/ogen-go/ogen/commit/381d15ef02a602261e24b0a2dd53d36dccb7bfab?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go).
> Patch coverage: 10.31% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #734      +/-   ##
==========================================
- Coverage   74.31%   74.01%   -0.31%     
==========================================
  Files         182      182              
  Lines       12961    13034      +73     
==========================================
+ Hits         9632     9647      +15     
- Misses       2835     2893      +58     
  Partials      494      494              
```


| [Impacted Files](https://codecov.io/gh/ogen-go/ogen/pull/734?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go) | Coverage Δ | |
|---|---|---|
| [tools/sgcollector/main.go](https://codecov.io/gh/ogen-go/ogen/pull/734/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go#diff-dG9vbHMvc2djb2xsZWN0b3IvbWFpbi5nbw==) | `0.00% <0.00%> (ø)` | |
| [tools/sgcollector/sourcegraph.go](https://codecov.io/gh/ogen-go/ogen/pull/734/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go#diff-dG9vbHMvc2djb2xsZWN0b3Ivc291cmNlZ3JhcGguZ28=) | `0.00% <0.00%> (ø)` | |
| [tools/sgcollector/worker.go](https://codecov.io/gh/ogen-go/ogen/pull/734/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go#diff-dG9vbHMvc2djb2xsZWN0b3Ivd29ya2VyLmdv) | `9.92% <38.23%> (+9.92%)` | :arrow_up: |
| [gen/gen\_contents.go](https://codecov.io/gh/ogen-go/ogen/pull/734/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go#diff-Z2VuL2dlbl9jb250ZW50cy5nbw==) | `75.00% <0.00%> (+0.73%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/ogen-go/ogen/pull/734?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ogen-go).


Patch:
diff --git a/tools/sgcollector/main.go b/tools/sgcollector/main.go
index 5bc1a335..16177489 100644
--- a/tools/sgcollector/main.go
+++ b/tools/sgcollector/main.go
@@ -118,26 +118,23 @@ func run(ctx context.Context) error {
 		defer close(links)
 
 		for _, q := range queries {
-			results, err := search(ctx, q)
-			if err != nil {
-				return err
-			}
-			logger.Info("Query complete",
-				zap.String("query", q),
-				zap.Int("total", len(results.Matches)),
-				zap.Int("match_count", results.MatchCount),
-			)
+			logger := logger.With(zap.String("query", q))
 
-			for i, m := range results.Matches {
+			logger.Info("Start query")
+			counter := 0
+			if err := search(ctx, q, func(match FileMatch) error {
 				select {
-				case <-ctx.Done():
-					return nil
-				case links <- m:
+				case links <- match:
 					total++
+					counter++
+					return nil
+				case <-ctx.Done():
+					return ctx.Err()
 				}
-				// Zero element to let GC collect FileMatch's fields.
-				results.Matches[i] = FileMatch{}
+			}); err != nil {
+				return err
 			}
+			logger.Info("Query complete", zap.Int("total", counter))
 		}
 		return nil
 	})
@@ -160,6 +157,7 @@ func run(ctx context.Context) error {
 						return nil
 					}
 
+					logger.Debug("Processing link", zap.Inline(m))
 					err := worker(ctx, m, reporters)
 					if err != nil {
 						logger.Error("Error",
diff --git a/tools/sgcollector/sourcegraph.go b/tools/sgcollector/sourcegraph.go
index 8d6c8ba3..b18debda 100644
--- a/tools/sgcollector/sourcegraph.go
+++ b/tools/sgcollector/sourcegraph.go
@@ -4,9 +4,12 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"io"
 	"net/http"
+	"os"
 
 	"github.com/go-faster/errors"
+	"github.com/go-faster/jx"
 	"go.uber.org/zap/zapcore"
 )
 
@@ -41,26 +44,6 @@ type (
 		Description string `json:"description"`
 	}
 
-	SearchResult struct {
-		Matches             []FileMatch `json:"results"`
-		LimitHit            bool        `json:"limitHit"`
-		MatchCount          int         `json:"matchCount"`
-		ElapsedMilliseconds int         `json:"elapsedMilliseconds"`
-		Alert               Alert       `json:"alert"`
-	}
-
-	SearchResults struct {
-		Results SearchResult `json:"results"`
-	}
-
-	Data struct {
-		Search SearchResults `json:"search"`
-	}
-
-	Response struct {
-		Data Data `json:"data"`
-	}
-
 	QueryVariables struct {
 		Query string `json:"query"`
 	}
@@ -130,23 +113,111 @@ fragment SearchResultsAlertFields on SearchResults {
 }
 `
 
-func search(ctx context.Context, query string) (SearchResult, error) {
-	r, err := querySourcegraph(ctx, GraphQLQuery{
+func search(ctx context.Context, query string, cb func(FileMatch) error) (rerr error) {
+	unmarshal := func(d *jx.Decoder, v any) error {
+		raw, err := d.Raw()
+		if err != nil {
+			return errors.Wrap(err, "scan")
+		}
+
+		if err := json.Unmarshal(raw, v); err != nil {
+			return errors.Wrap(err, "unmarshal")
+		}
+
+		return nil
+	}
+
+	err := querySourcegraph(ctx, GraphQLQuery{
 		Query: graphQLQuery,
 		Variables: QueryVariables{
 			Query: query,
 		},
+	}, func(d *jx.Decoder, key []byte) error {
+		switch string(key) {
+		case "results":
+			return d.Arr(func(d *jx.Decoder) error {
+				var fm FileMatch
+				if err := unmarshal(d, &fm); err != nil {
+					return errors.Wrap(err, "parse FileMatch")
+				}
+
+				if err := cb(fm); err != nil {
+					return errors.Wrap(err, "callback")
+				}
+				return nil
+			})
+		case "alert":
+			var a Alert
+			if err := unmarshal(d, &a); err != nil {
+				return errors.Wrap(err, "parse Alert")
+			}
+
+			if a.Title != "" {
+				alert := a.Title
+				if a.Description != "" {
+					alert = a.Description
+				}
+				return errors.Errorf("alert: %s", alert)
+			}
+
+			return nil
+		default:
+			return d.Skip()
+		}
 	})
 	if err != nil {
-		return r, errors.Wrapf(err, "query: %q", query)
+		return errors.Wrapf(err, "query: %q", query)
 	}
-	return r, nil
+	return nil
 }
 
-func querySourcegraph(ctx context.Context, q GraphQLQuery) (SearchResult, error) {
+func querySourcegraph(ctx context.Context, q GraphQLQuery, cb func(d *jx.Decoder, key []byte) error) error {
+	// Handling of the response may take a while, so we save the response to avoid timeouts.
+	f, err := os.CreateTemp("", "sgcollector*")
+	if err != nil {
+		return errors.Wrap(err, "create temp file")
+	}
+	defer func() {
+		_ = f.Close()
+	}()
+
+	if err := sendSourcegraph(ctx, q, f); err != nil {
+		return err
+	}
+
+	if _, err := f.Seek(0, io.SeekStart); err != nil {
+		return errors.Wrap(err, "seek")
+	}
+
+	d := jx.Decode(f, 4096)
+	return d.ObjBytes(func(d *jx.Decoder, key []byte) error {
+		switch string(key) {
+		case "data":
+			return d.ObjBytes(func(d *jx.Decoder, key []byte) error {
+				switch string(key) {
+				case "search":
+					return d.ObjBytes(func(d *jx.Decoder, key []byte) error {
+						switch string(key) {
+						case "results":
+							return d.ObjBytes(cb)
+						default:
+							return d.Skip()
+						}
+					})
+				default:
+					return d.Skip()
+				}
+			})
+		default:
+			return d.Skip()
+		}
+	})
+}
+
+func sendSourcegraph(ctx context.Context, q GraphQLQuery, out io.Writer) error {
 	body, err := json.Marshal(q)
 	if err != nil {
-		return SearchResult{}, errors.Wrap(err, "encode")
+		return errors.Wrap(err, "encode query")
 	}
 
 	req, err := http.NewRequestWithContext(ctx,
@@ -154,34 +225,23 @@ func querySourcegraph(ctx context.Context, q GraphQLQuery) (SearchResult, error)
 		bytes.NewReader(body),
 	)
 	if err != nil {
-		return SearchResult{}, errors.Wrap(err, "create request")
+		return errors.Wrap(err, "create request")
 	}
 
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
-		return SearchResult{}, errors.Wrap(err, "do request")
+		return errors.Wrap(err, "do request")
 	}
 	defer func() {
 		_ = resp.Body.Close()
 	}()
 
 	if resp.StatusCode >= 400 {
-		return SearchResult{}, errors.Errorf("http error: %s", resp.Status)
+		return errors.Errorf("http error: %s", resp.Status)
 	}
 
-	var r Response
-	if err := json.NewDecoder(resp.Body).Decode(&r); err != nil {
-		return SearchResult{}, errors.Wrap(err, "parse")
+	if _, err := io.Copy(out, resp.Body); err != nil {
+		return errors.Wrap(err, "copy response")
 	}
-
-	result := r.Data.Search.Results
-	if a := result.Alert; a.Title != "" {
-		alert := a.Title
-		if a.Description != "" {
-			alert = a.Description
-		}
-		return SearchResult{}, errors.Errorf("alert: %s", alert)
-	}
-
-	return result, nil
+	return nil
 }
diff --git a/tools/sgcollector/worker.go b/tools/sgcollector/worker.go
index 780dda06..189d694b 100644
--- a/tools/sgcollector/worker.go
+++ b/tools/sgcollector/worker.go
@@ -1,40 +1,26 @@
 package main
 
 import (
-	"bytes"
 	"context"
 	"crypto/sha256"
 	"fmt"
 	"net/http"
+	"net/url"
+	"regexp"
 	"strings"
 	"time"
 
-	"github.com/ghodss/yaml"
 	"github.com/go-faster/errors"
 	"github.com/go-faster/jx"
+	yaml "github.com/go-faster/yamlx"
 	"golang.org/x/exp/slices"
 
 	"github.com/ogen-go/ogen"
 	"github.com/ogen-go/ogen/gen"
+	"github.com/ogen-go/ogen/internal/location"
 )
 
-var (
-	errPanic  = errors.New("panic")
-	bomPrefix = []byte{0xEF, 0xBB, 0xBF}
-)
-
-func convertYAMLtoJSON(data []byte) (_ []byte, rErr error) {
-	defer func() {
-		if rr := recover(); rr != nil {
-			rErr = errors.Errorf("panic: %#v", rr)
-		}
-	}()
-	j, err := yaml.YAMLToJSON(data)
-	if err != nil {
-		return nil, err
-	}
-	return j, nil
-}
+var errPanic = errors.New("panic")
 
 func validateJSON(data []byte) error {
 	d := jx.GetDecoder()
@@ -43,8 +29,33 @@ func validateJSON(data []byte) error {
 	return d.Validate()
 }
 
+var gitPathRegex = regexp.MustCompile(`^/(?P<owner>[^/]+)/(?P<repo>[^/]+)/(-/blob|blob)/(?P<ref>[^/]+)/(?P<path>.*)$`)
+
+func getRootURL(m FileMatch) (*url.URL, bool) {
+	for _, external := range m.File.ExternalURLs {
+		u, err := url.Parse(external.URL)
+		if err != nil {
+			continue
+		}
+		if !gitPathRegex.MatchString(u.Path) {
+			continue
+		}
+		switch external.ServiceKind {
+		case "GITHUB":
+			u.Host = "raw.githubusercontent.com"
+			u.Path = gitPathRegex.ReplaceAllString(u.Path, "/$owner/$repo/$ref/$path")
+		case "GITLAB":
+			u.Path = gitPathRegex.ReplaceAllString(u.Path, "/$owner/$repo/raw/$ref/$path")
+		default:
+			continue
+		}
+		return u, true
+	}
+	return nil, false
+}
+
 func worker(ctx context.Context, m FileMatch, r *Reporters) (rErr error) {
-	data := bytes.TrimPrefix([]byte(m.File.Content), bomPrefix)
+	data := []byte(m.File.Content)
 	isYAML := strings.HasSuffix(m.File.Name, ".yml") || strings.HasSuffix(m.File.Name, ".yaml")
 	hash := sha256.Sum256(data)
 
@@ -61,7 +72,15 @@ func worker(ctx context.Context, m FileMatch, r *Reporters) (rErr error) {
 		}
 	}()
 
-	pse := generate(data, isYAML)
+	rootURL := &url.URL{
+		Scheme: "jsonschema",
+		Host:   "dummy",
+	}
+	if u, ok := getRootURL(m); ok {
+		rootURL = u
+	}
+
+	pse := generate(data, isYAML, m.File.Path, rootURL)
 	if pse != nil {
 		if err := r.report(ctx, pse.stage, Report{
 			File:           m,
@@ -104,17 +123,16 @@ func workerHTTPClient() *http.Client {
 	}
 }
 
-func generate(data []byte, isYAML bool) *GenerateError {
+func generate(data []byte, isYAML bool, fileName string, rootURL *url.URL) *GenerateError {
 	if isYAML {
-		j, err := convertYAMLtoJSON(data)
-		if err != nil {
+		var n *yaml.Node
+		if err := yaml.Unmarshal(data, &n); err != nil {
 			return &GenerateError{stage: InvalidYAML, err: err}
 		}
-		data = j
-	}
-
-	if err := validateJSON(data); err != nil {
-		return &GenerateError{stage: InvalidJSON, err: err}
+	} else {
+		if err := validateJSON(data); err != nil {
+			return &GenerateError{stage: InvalidJSON, err: err}
+		}
 	}
 
 	spec, err := ogen.Parse(data)
@@ -129,6 +147,7 @@ func generate(data []byte, isYAML bool) *GenerateError {
 	g, err := gen.NewGenerator(spec, gen.Options{
 		InferSchemaType: true,
 		AllowRemote:     true,
+		RootURL:         rootURL,
 		Remote: gen.RemoteOptions{
 			HTTPClient: workerHTTPClient(),
 			ReadFile: func(string) ([]byte, error) {
@@ -147,6 +166,7 @@ func generate(data []byte, isYAML bool) *GenerateError {
 			}
 			notImpl = append(notImpl, name)
 		},
+		File: location.NewFile(fileName, rootURL.String(), data),
 	})
 
 	slices.Sort(notImpl)


Test patch:
diff --git a/tools/sgcollector/worker_test.go b/tools/sgcollector/worker_test.go
new file mode 100644
index 00000000..f23b6881
--- /dev/null
+++ b/tools/sgcollector/worker_test.go
@@ -0,0 +1,85 @@
+package main
+
+import (
+	"fmt"
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_getRootURL(t *testing.T) {
+	github := func(u string) FileMatch {
+		return FileMatch{
+			File: File{
+				ExternalURLs: []ExternalURL{
+					{ServiceKind: "GITHUB", URL: u},
+				},
+			},
+		}
+	}
+	gitlab := func(u string) FileMatch {
+		return FileMatch{
+			File: File{
+				ExternalURLs: []ExternalURL{
+					{ServiceKind: "GITLAB", URL: u},
+				},
+			},
+		}
+	}
+	mustURL := func(s string) *url.URL {
+		u, err := url.Parse(s)
+		require.NoError(t, err)
+		return u
+	}
+
+	tests := []struct {
+		m     FileMatch
+		want  *url.URL
+		want1 bool
+	}{
+		// GitHub links.
+		{
+			github("https://github.com/owner/repo/blob/abc-the-commit/dir1/dir2/file.yml"),
+			mustURL("https://raw.githubusercontent.com/owner/repo/abc-the-commit/dir1/dir2/file.yml"),
+			true,
+		},
+		{
+			github("https://github.com/owner/repo/blob/abc-the-commit/file.yml"),
+			mustURL("https://raw.githubusercontent.com/owner/repo/abc-the-commit/file.yml"),
+			true,
+		},
+		// Bad GitHub links.
+		{
+			github("https://github.com/owner/blob/abc-the-commit/file.yml"),
+			nil,
+			false,
+		},
+
+		// GitLab links.
+		{
+			gitlab("https://gitlab.com/owner/repo/blob/abc-the-commit/file.yml"),
+			mustURL("https://gitlab.com/owner/repo/raw/abc-the-commit/file.yml"),
+			true,
+		},
+		{
+			gitlab("https://gitlab.com/owner/repo/-/blob/abc-the-commit/file.yml"),
+			mustURL("https://gitlab.com/owner/repo/raw/abc-the-commit/file.yml"),
+			true,
+		},
+	}
+	for i, tt := range tests {
+		tt := tt
+		t.Run(fmt.Sprintf("Test%d", i+1), func(t *testing.T) {
+			a := require.New(t)
+
+			got, ok := getRootURL(tt.m)
+			if !tt.want1 {
+				a.False(ok)
+				return
+			}
+			a.True(ok)
+			a.Equal(tt.want, got)
+		})
+	}
+}