feat: rewrite alias/upstream computation in Go #4395

michaelkedar · 2025-11-25T06:01:32Z

This is pretty much a direct translation of the existing python code (which makes some of it unidiomatic go), changed to use the new Vulnerability instead of Bug.

One difference is the alias and upstream computations now run at the same time, and the writes to the Vulnerability protos now happen together. I had to modify the recoverer to be able to handle multiple updating multiple fields when the GCS generation is mismatched to support this properly.

Also ported over the datastore emulator manager, and made some interfaces for the google cloud clients (with the added benefit of introducing write retries which resolves #4355)
This has allowed my to bring over the existing tests.

I've renamed the alias-computation cron to relations - Cloud Deploy is going to create a new workload for this, and the old one is going to have to be deleted manually.

In a follow-up, I'm going to add another goroutine to calculate the related field, since these aren't being kept up to date at the moment.

michaelkedar · 2025-11-26T04:58:12Z

/gemini review

gemini-code-assist

Code Review

This is a major rewrite of the alias and upstream computation logic from Python to Go, and the changes are extensive but well-executed. The introduction of interfaces for cloud clients (CloudStorage, Publisher) and mock implementations for testing is a significant improvement in code quality and testability. The new relations Go command is well-structured, and the test coverage is comprehensive. The addition of a datastore emulator for testing is also a great enhancement to the development workflow. I've identified a minor resource leak and a point of code complexity that could be simplified, but overall this is a high-quality contribution.

go/cmd/relations/relations.go

gemini-code-assist · 2025-11-26T05:00:47Z

go/cmd/relations/upstream.go

+		for _, upstream := range group.UpstreamIDs {
+			if _, ok := visited[upstream]; !ok && !slices.Contains(toVisit, upstream) {
+				toVisit = append(toVisit, upstream)
+			} else {
+				if u, ok := upstreamMap[vulnID]; !ok {
+					upstreamMap[vulnID] = []string{upstream}
+				} else if !slices.Contains(u, upstream) {
+					upstreamMap[vulnID] = append(u, upstream)
+				}
+			}
+		}
+		// Add the immediate upstreams of the vuln to the map
+		upstreamMap[vulnID] = group.UpstreamIDs
+		for _, upstream := range group.UpstreamIDs {
+			if _, ok := visited[upstream]; !ok && !slices.Contains(toVisit, upstream) {
+				toVisit = append(toVisit, upstream)
+			}
+		}


This block of code for populating toVisit and upstreamMap appears to be more complex than necessary and contains some redundant logic. The first loop from line 135 seems to be partially overwritten by the assignment on line 147 and the second loop on line 148. This can be simplified for clarity and to remove the redundancy.

if len(group.UpstreamIDs) == 0 { continue } // Add the immediate upstreams of the vuln to the map upstreamMap[vulnID] = group.UpstreamIDs for _, upstream := range group.UpstreamIDs { if _, ok := visited[upstream]; !ok && !slices.Contains(toVisit, upstream) { toVisit = append(toVisit, upstream) } }

FWIW, this logic is basically directly translated from the python implementation:

osv.dev/gcp/workers/alias/upstream_computation.py

Lines 171 to 181 in 3016163

for upstream in upstreams:

if upstream not in visited and upstream not in to_visit:

to_visit.add(upstream)

else:

if bug_id not in upstream_map:

upstream_map[bug_id] = set([upstream])

else:

upstream_map[bug_id].add(upstream)

# Add the immediate upstreams of the bug to the dict

upstream_map[bug_id] = upstreams

to_visit.update(upstreams - visited)

michaelkedar · 2025-11-26T23:43:08Z

/gemini review

gemini-code-assist

Code Review

This is a significant and well-executed pull request that rewrites the alias and upstream computation logic from Python to Go. The new Go implementation is well-structured, and the introduction of interfaces for cloud clients and comprehensive test utilities is a fantastic improvement for testability and maintainability. The core logic appears to be a faithful translation of the original Python code. I've identified a couple of critical compilation errors related to the misuse of sync.WaitGroup that need to be addressed. Otherwise, the changes are excellent.

go/cmd/relations/relations.go

go/cmd/relations/update.go

another-rex · 2025-11-27T00:05:53Z

gcp/workers/recoverer/recoverer.py

-          vuln_proto.aliases[:] = aliases
-          if aliases_modified > modified:
-            modified = aliases_modified
+      fields = field.split(',')


Can you describe in the PR description what this change is for?

Added:

One difference is the alias and upstream computations now run at the same time, and the writes to the Vulnerability protos now happen together. I had to modify the recoverer to be able to handle multiple updating multiple fields when the GCS generation is mismatched to support this properly.

another-rex · 2025-11-27T00:09:18Z

go/cmd/exporter/exporter.go

 	//    generate the final files, sending the data to be written to `routerToWriteCh`.
 	// 5. A pool of `writer` workers receive the file data and write it to the output.
-	gcsObjToDownloaderCh := make(chan *storage.ObjectHandle)
+	gcsObjToDownloaderCh := make(chan string)


nit: rename this variable since it's taking in a path now
nit2: Alias string to a gcsPath type and put that as the channel type.

renamed to gcsPathToDownloaderCh
re: nit2 - I don't really see the benefit of doing that - it's fairly standard to use strings as paths in go (e.g. in the path package)

It's more to clarify whether you are meant to pass in a path, a name, or a vuln ID...etc. This makes it a bit more self documenting. I think you don't even need to cast when passing in a string for aliases

another-rex · 2025-11-27T00:10:12Z

go/cmd/exporter/exporter.go

-		if errors.Is(err, iterator.Done) {
-			break
-		}
+	for path, err := range vulnClient.Objects(ctx, gcsProtoPrefix) {


This is so clean!

go/osv/clients/gcs_client.go

go/cmd/relations/alias.go

another-rex · 2025-11-27T05:07:52Z

go/cmd/relations/upstream.go

+	query := datastore.NewQuery("Vulnerability").FilterField("upstream_raw", ">", "")
+	it := cl.Run(ctx, query)
+
+	vulns := make(map[string][]string)


Add a comment here about what vulns is suppose to contain.

I renamed this to rawUpstreams which is more descriptive

go/cmd/relations/upstream_test.go

another-rex · 2025-11-27T05:10:03Z

go/cmd/relations/upstream.go

+	"google.golang.org/api/iterator"
+)
+
+func computeUpstream(targetBugUpstream []string, bugs map[string][]string) []string {


Can you add a comment describing what this function does, and what the two args are. it's unclear why targetBugUpstream is a list.

I added back in all the original docstrings from python to the go code, which should make it slightly clearer.

I also changed the function signature to be computeUpstream(vulnID []string, rawUpstreams map[string][]string) []string , which is simpler and clearer.

go/cmd/relations/upstream.go

go/osv/clients/gcs_client.go

michaelkedar added 7 commits November 20, 2025 11:40

alias WIP

beddd4b

alias wip 2

35bade2

close the clients

a809032

minor changes

3e9c35c

upstream computation

384836d

lint

67b768d

recovery

207359a

gemini-code-assist bot reviewed Nov 26, 2025

View reviewed changes

michaelkedar added 2 commits November 26, 2025 16:11

fix leaks

0c90403

lont

1da9942

michaelkedar force-pushed the relationships branch from 8bd87f7 to 1da9942 Compare November 26, 2025 05:58

michaelkedar added 4 commits November 26, 2025 17:23

add test runners

7e38f84

fix test path

4c59077

setup the venv

dd5be0a

deploy go alias/upstream

91fdbac

gemini-code-assist bot reviewed Nov 26, 2025

View reviewed changes

go/cmd/relations/relations.go Show resolved Hide resolved

go/cmd/relations/update.go Show resolved Hide resolved

michaelkedar changed the title ~~feat: rewrite alias/upstream computation in Go (WIP)~~ feat: rewrite alias/upstream computation in Go Nov 26, 2025

Merge branch 'master' into relationships

c81d518

michaelkedar marked this pull request as ready for review November 26, 2025 23:48

michaelkedar requested review from another-rex, cuixq and jess-lowe November 26, 2025 23:49

michaelkedar mentioned this pull request Nov 27, 2025

Improve the OSV.dev database situation #3850

Open

another-rex reviewed Nov 27, 2025

View reviewed changes

michaelkedar added 4 commits December 1, 2025 10:47

some review nits

3fdbdba

bug -> vuln

697f2e4

make it neater

95e8c54

gcs comment

0118be6

michaelkedar requested a review from another-rex December 1, 2025 02:27

another-rex approved these changes Dec 1, 2025

View reviewed changes

	for upstream in upstreams:
	if upstream not in visited and upstream not in to_visit:
	to_visit.add(upstream)
	else:
	if bug_id not in upstream_map:
	upstream_map[bug_id] = set([upstream])
	else:
	upstream_map[bug_id].add(upstream)
	# Add the immediate upstreams of the bug to the dict
	upstream_map[bug_id] = upstreams
	to_visit.update(upstreams - visited)

feat: rewrite alias/upstream computation in Go #4395

Are you sure you want to change the base?

feat: rewrite alias/upstream computation in Go #4395

Uh oh!

Conversation

michaelkedar commented Nov 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

michaelkedar commented Nov 26, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

gemini-code-assist bot Nov 26, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

michaelkedar commented Nov 26, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

michaelkedar commented Nov 25, 2025 •

edited

Loading