Enhance commit retrieval with branch & tag prefix support

pkg/querier/vcs/client/github.go

@@ -12,6 +12,7 @@ import (
 	"golang.org/x/oauth2"
 
 	vcsv1 "github.com/grafana/pyroscope/api/gen/proto/go/vcs/v1"
+	"github.com/grafana/pyroscope/pkg/util/connectgrpc"
 )
 
 // GithubClient returns a github client.
@@ -28,6 +29,11 @@ type githubClient struct {
 func (gh *githubClient) GetCommit(ctx context.Context, owner, repo, ref string) (*vcsv1.GetCommitResponse, error) {
 	commit, _, err := gh.client.Repositories.GetCommit(ctx, owner, repo, ref, nil)
 	if err != nil {
+		var githubErr *github.ErrorResponse
+		if errors.As(err, &githubErr) {
+			code := connectgrpc.HTTPToCode(int32(githubErr.Response.StatusCode))
+			return nil, connect.NewError(code, err)
+		}
 		return nil, err
 	}
 

pkg/querier/vcs/service.go

@@ -27,6 +27,10 @@ type Service struct {
 	httpClient *http.Client
 }
 
+type gitHubCommitGetter interface {
+	GetCommit(context.Context, string, string, string) (*vcsv1.GetCommitResponse, error)
+}
+
 func New(logger log.Logger, reg prometheus.Registerer) *Service {
 	httpClient := client.InstrumentedHTTPClient(logger, reg)
 
@@ -189,13 +193,41 @@ func (q *Service) GetCommit(ctx context.Context, req *connect.Request[vcsv1.GetC
 		return nil, err
 	}
 
-	commit, err := ghClient.GetCommit(ctx, gitURL.GetOwnerName(), gitURL.GetRepoName(), req.Msg.Ref)
+	owner := gitURL.GetOwnerName()
+	repo := gitURL.GetRepoName()
+	ref := req.Msg.GetRef()
+
+	commit, err := q.tryGetCommit(ctx, ghClient, owner, repo, ref)
 	if err != nil {
 		return nil, err
 	}
+
 	return connect.NewResponse(commit), nil
 }
 
+// tryGetCommit attempts to retrieve a commit using different ref formats (commit hash, branch, tag).
+// It tries each format in order and returns the first successful result.
+func (q *Service) tryGetCommit(ctx context.Context, client gitHubCommitGetter, owner, repo, ref string) (*vcsv1.GetCommitResponse, error) {
+	refFormats := []string{
+		ref,            // Try as a commit hash
+		"heads/" + ref, // Try as a branch
+		"tags/" + ref,  // Try as a tag
+	}
+
+	var lastErr error
+	for _, format := range refFormats {
+		commit, err := client.GetCommit(ctx, owner, repo, format)
+		if err == nil {
+			return commit, nil
+		}
+
+		lastErr = err
+		q.logger.Log("err", lastErr, "msg", "Failed to get commit", "ref", format)
+	}
+
+	return nil, lastErr
+}
+
 func rejectExpiredToken(token *oauth2.Token) error {
 	if time.Now().After(token.Expiry) {
 		return connect.NewError(connect.CodeUnauthenticated, fmt.Errorf("token is expired"))

pkg/querier/vcs/service_test.go

@@ -0,0 +1,108 @@
+package vcs
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"testing"
+
+	"github.com/go-kit/log"
+	"github.com/google/go-github/v58/github"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	vcsv1 "github.com/grafana/pyroscope/api/gen/proto/go/vcs/v1"
+)
+
+type gitHubCommitGetterMock struct {
+	mock.Mock
+}
+
+func (m *gitHubCommitGetterMock) GetCommit(ctx context.Context, owner, repo, ref string) (*vcsv1.GetCommitResponse, error) {
+	args := m.Called(ctx, owner, repo, ref)
+	if args.Get(0) == nil {
+		return nil, args.Error(1)
+	}
+
+	return args.Get(0).(*vcsv1.GetCommitResponse), args.Error(1)
+}
+
+func TestTryGetCommit(t *testing.T) {
+	svc := Service{logger: log.NewNopLogger()}
+
+	tests := []struct {
+		name       string
+		setupMock  func(*gitHubCommitGetterMock)
+		ref        string
+		wantCommit *vcsv1.GetCommitResponse
+		wantErr    bool
+	}{
+		{
+			name: "Direct commit hash",
+			setupMock: func(m *gitHubCommitGetterMock) {
+				m.On("GetCommit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+					Return(&vcsv1.GetCommitResponse{Sha: "abc123"}, nil)
+			},
+			ref:        "",
+			wantCommit: &vcsv1.GetCommitResponse{Sha: "abc123"},
+			wantErr:    false,
+		},
+		{
+			name: "Branch reference with 'heads/' prefix",
+			setupMock: func(m *gitHubCommitGetterMock) {
+				m.On("GetCommit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+					Return(nil, assert.AnError).Times(1)
+				m.On("GetCommit", mock.Anything, mock.Anything, mock.Anything, "heads/main").
+					Return(&vcsv1.GetCommitResponse{Sha: "def456"}, nil).Times(1)
+			},
+			ref:        "main",
+			wantCommit: &vcsv1.GetCommitResponse{Sha: "def456"},
+			wantErr:    false,
+		},
+		{
+			name: "Tag reference with 'tags/' prefix",
+			setupMock: func(m *gitHubCommitGetterMock) {
+				m.On("GetCommit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+					Return(nil, assert.AnError).Times(2)
+				m.On("GetCommit", mock.Anything, mock.Anything, mock.Anything, "tags/v1").
+					Return(&vcsv1.GetCommitResponse{Sha: "def456"}, nil).Times(1)
+			},
+			ref:        "v1",
+			wantCommit: &vcsv1.GetCommitResponse{Sha: "def456"},
+			wantErr:    false,
+		},
+		{
+			name: "GitHub API returns not found error",
+			setupMock: func(m *gitHubCommitGetterMock) {
+				notFoundErr := &github.ErrorResponse{
+					Response: &http.Response{StatusCode: http.StatusNotFound},
+				}
+				m.On("GetCommit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+					Return(nil, notFoundErr).Times(3)
+			},
+			ref:        "nonexistent",
+			wantCommit: nil,
+			wantErr:    true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockGetter := new(gitHubCommitGetterMock)
+			tt.setupMock(mockGetter)
+
+			gotCommit, err := svc.tryGetCommit(context.Background(), mockGetter, "owner", "repo", tt.ref)
+
+			if tt.wantErr {
+				assert.Error(t, err)
+				var githubErr *github.ErrorResponse
+				assert.True(t, errors.As(err, &githubErr), "Expected a GitHub error")
+			} else {
+				assert.NoError(t, err)
+			}
+
+			assert.Equal(t, tt.wantCommit, gotCommit)
+			mockGetter.AssertExpectations(t)
+		})
+	}
+}