Skip to content

Commit

Permalink
refactor codegen parts (#1863)
Browse files Browse the repository at this point in the history 
* refactor codegen parts
  • Loading branch information
@motatoes
motatoes authored Dec 27, 2024
1 parent 7706260 commit 34d1db1
Show file tree
Hide file tree
Showing 2 changed files with 156 additions and 150 deletions.
302 changes: 154 additions & 148 deletions backend/controllers/github.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -767,154 +767,8 @@ func handleIssueCommentEvent(gh utils.GithubClientProvider, payload *github.Issu

// terraform code generator
if os.Getenv("DIGGER_GENERATION_ENABLED") == "1" {
if strings.HasPrefix(*payload.Comment.Body, "digger generate") {
projectName := ci.ParseProjectName(*payload.Comment.Body)
if projectName == "" {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: generate requires argument -p <project_name>: %v", err))
log.Printf("missing project in command: %v", *payload.Comment.Body)
return fmt.Errorf("generate requires argument -p <project_name>: %v", err)
}

project := config.GetProject(projectName)
if project == nil {
commentReporterManager.UpdateComment(fmt.Sprintf("could not find project %v in digger.yml", projectName))
log.Printf("could not find project %v in digger.yml", projectName)
return fmt.Errorf("could not find project %v in digger.yml", projectName)
}

commentReporterManager.UpdateComment(fmt.Sprintf(":white_check_mark: Successfully loaded project"))

generationEndpoint := os.Getenv("DIGGER_GENERATION_ENDPOINT")
if generationEndpoint == "" {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: server does not have generation endpoint configured, please verify"))
log.Printf("server does not have generation endpoint configured, please verify")
return fmt.Errorf("server does not have generation endpoint configured, please verify")
}
webhookSecret := os.Getenv("DIGGER_GENERATION_WEBHOOK_SECRET")

// Get all code content from the repository at a specific commit
getCodeFromCommit := func(ghService *dg_github.GithubService, repoOwner, repoName string, commitSha *string, projectDir string) (string, error) {
const MaxPatchSize = 1024 * 1024 // 1MB limit

// Get the commit's changes compared to default branch
comparison, _, err := ghService.Client.Repositories.CompareCommits(
context.Background(),
repoOwner,
repoName,
defaultBranch,
*commitSha,
nil,
)
if err != nil {
return "", fmt.Errorf("error comparing commits: %v", err)
}

var appCode strings.Builder
for _, file := range comparison.Files {
if file.Patch == nil {
continue // Skip files without patches
}
log.Printf("Processing patch for file: %s", *file.Filename)
if *file.Additions > 0 {
lines := strings.Split(*file.Patch, "\n")
for _, line := range lines {
if strings.HasPrefix(line, "+") && !strings.HasPrefix(line, "+++") {
appCode.WriteString(strings.TrimPrefix(line, "+"))
appCode.WriteString("\n")
}
}
}
appCode.WriteString("\n")
}

if appCode.Len() == 0 {
return "", fmt.Errorf("no code changes found in commit %s. Please ensure the PR contains added or modified code", *commitSha)
}

return appCode.String(), nil
}

appCode, err := getCodeFromCommit(ghService, repoOwner, repoName, commitSha, project.Dir)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to get code content: %v", err))
log.Printf("Error getting code content: %v", err)
return fmt.Errorf("error getting code content: %v", err)
}

commentReporterManager.UpdateComment(fmt.Sprintf(":white_check_mark: Successfully loaded code from commit"))

log.Printf("the app code is: %v", appCode)

commentReporterManager.UpdateComment(fmt.Sprintf("Generating terraform..."))
terraformCode, err := utils.GenerateTerraformCode(appCode, generationEndpoint, webhookSecret)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: could not generate terraform code: %v", err))
log.Printf("could not generate terraform code: %v", err)
return fmt.Errorf("could not generate terraform code: %v", err)
}

commentReporterManager.UpdateComment(fmt.Sprintf(":white_check_mark: Generated terraform"))

// comment terraform code to project dir
//project.Dir
log.Printf("terraform code is %v", terraformCode)

baseTree, _, err := ghService.Client.Git.GetTree(context.Background(), repoOwner, repoName, *commitSha, false)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to get base tree: %v", err))
log.Printf("Error getting base tree: %v", err)
return fmt.Errorf("error getting base tree: %v", err)
}

// Create a new tree with the new file
treeEntries := []*github.TreeEntry{
{
Path: github.String(filepath.Join(project.Dir, fmt.Sprintf("generated_%v.tf", issueNumber))),
Mode: github.String("100644"),
Type: github.String("blob"),
Content: github.String(terraformCode),
},
}

newTree, _, err := ghService.Client.Git.CreateTree(context.Background(), repoOwner, repoName, *baseTree.SHA, treeEntries)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to create new tree: %v", err))
log.Printf("Error creating new tree: %v", err)
return fmt.Errorf("error creating new tree: %v", err)
}

// Create the commit
commitMsg := fmt.Sprintf("Add generated Terraform code for %v", projectName)
commit := &github.Commit{
Message: &commitMsg,
Tree: newTree,
Parents: []*github.Commit{{SHA: commitSha}},
}

newCommit, _, err := ghService.Client.Git.CreateCommit(context.Background(), repoOwner, repoName, commit, nil)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to commit Terraform file: %v", err))
log.Printf("Error committing Terraform file: %v", err)
return fmt.Errorf("error committing Terraform file: %v", err)
}

// Update the reference to point to the new commit
ref := &github.Reference{
Ref: github.String(fmt.Sprintf("refs/heads/%s", *branch)),
Object: &github.GitObject{
SHA: newCommit.SHA,
},
}
_, _, err = ghService.Client.Git.UpdateRef(context.Background(), repoOwner, repoName, ref, false)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to update branch reference: %v", err))
log.Printf("Error updating branch reference: %v", err)
return fmt.Errorf("error updating branch reference: %v", err)
}

commentReporterManager.UpdateComment(":white_check_mark: Successfully generated and committed Terraform code")
return nil
}
err = GenerateTerraformFromCode(payload, commentReporterManager, config, defaultBranch, ghService, repoOwner, repoName, commitSha, issueNumber, branch)
return err
}

commentIdStr := strconv.FormatInt(userCommentId, 10)
Expand Down Expand Up @@ -1120,6 +974,158 @@ func handleIssueCommentEvent(gh utils.GithubClientProvider, payload *github.Issu
return nil
}

func GenerateTerraformFromCode(payload *github.IssueCommentEvent, commentReporterManager utils.CommentReporterManager, config *dg_configuration.DiggerConfig, defaultBranch string, ghService *dg_github.GithubService, repoOwner string, repoName string, commitSha *string, issueNumber int, branch *string) error {
if strings.HasPrefix(*payload.Comment.Body, "digger generate") {
projectName := ci.ParseProjectName(*payload.Comment.Body)
if projectName == "" {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: generate requires argument -p <project_name>"))
log.Printf("missing project in command: %v", *payload.Comment.Body)
return fmt.Errorf("generate requires argument -p <project_name>")
}

project := config.GetProject(projectName)
if project == nil {
commentReporterManager.UpdateComment(fmt.Sprintf("could not find project %v in digger.yml", projectName))
log.Printf("could not find project %v in digger.yml", projectName)
return fmt.Errorf("could not find project %v in digger.yml", projectName)
}

commentReporterManager.UpdateComment(fmt.Sprintf(":white_check_mark: Successfully loaded project"))

generationEndpoint := os.Getenv("DIGGER_GENERATION_ENDPOINT")
if generationEndpoint == "" {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: server does not have generation endpoint configured, please verify"))
log.Printf("server does not have generation endpoint configured, please verify")
return fmt.Errorf("server does not have generation endpoint configured, please verify")
}
apiToken := os.Getenv("DIGGER_GENERATION_API_TOKEN")

// Get all code content from the repository at a specific commit
getCodeFromCommit := func(ghService *dg_github.GithubService, repoOwner, repoName string, commitSha *string, projectDir string) (string, error) {
const MaxPatchSize = 1024 * 1024 // 1MB limit

// Get the commit's changes compared to default branch
comparison, _, err := ghService.Client.Repositories.CompareCommits(
context.Background(),
repoOwner,
repoName,
defaultBranch,
*commitSha,
nil,
)
if err != nil {
return "", fmt.Errorf("error comparing commits: %v", err)
}

var appCode strings.Builder
for _, file := range comparison.Files {
if file.Patch == nil {
continue // Skip files without patches
}
log.Printf("Processing patch for file: %s", *file.Filename)
if *file.Additions > 0 {
lines := strings.Split(*file.Patch, "\n")
for _, line := range lines {
if strings.HasPrefix(line, "+") && !strings.HasPrefix(line, "+++") {
appCode.WriteString(strings.TrimPrefix(line, "+"))
appCode.WriteString("\n")
}
}
}
appCode.WriteString("\n")
}

if appCode.Len() == 0 {
return "", fmt.Errorf("no code changes found in commit %s. Please ensure the PR contains added or modified code", *commitSha)
}

return appCode.String(), nil
}

appCode, err := getCodeFromCommit(ghService, repoOwner, repoName, commitSha, project.Dir)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to get code content: %v", err))
log.Printf("Error getting code content: %v", err)
return fmt.Errorf("error getting code content: %v", err)
}

commentReporterManager.UpdateComment(fmt.Sprintf(":white_check_mark: Successfully loaded code from commit"))

log.Printf("the app code is: %v", appCode)

commentReporterManager.UpdateComment(fmt.Sprintf("Generating terraform..."))
terraformCode, err := utils.GenerateTerraformCode(appCode, generationEndpoint, apiToken)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: could not generate terraform code: %v", err))
log.Printf("could not generate terraform code: %v", err)
return fmt.Errorf("could not generate terraform code: %v", err)
}

commentReporterManager.UpdateComment(fmt.Sprintf(":white_check_mark: Generated terraform"))

// comment terraform code to project dir
//project.Dir
log.Printf("terraform code is %v", terraformCode)

baseTree, _, err := ghService.Client.Git.GetTree(context.Background(), repoOwner, repoName, *commitSha, false)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to get base tree: %v", err))
log.Printf("Error getting base tree: %v", err)
return fmt.Errorf("error getting base tree: %v", err)
}

// Create a new tree with the new file
treeEntries := []*github.TreeEntry{
{
Path: github.String(filepath.Join(project.Dir, fmt.Sprintf("generated_%v.tf", issueNumber))),
Mode: github.String("100644"),
Type: github.String("blob"),
Content: github.String(terraformCode),
},
}

newTree, _, err := ghService.Client.Git.CreateTree(context.Background(), repoOwner, repoName, *baseTree.SHA, treeEntries)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to create new tree: %v", err))
log.Printf("Error creating new tree: %v", err)
return fmt.Errorf("error creating new tree: %v", err)
}

// Create the commit
commitMsg := fmt.Sprintf("Add generated Terraform code for %v", projectName)
commit := &github.Commit{
Message: &commitMsg,
Tree: newTree,
Parents: []*github.Commit{{SHA: commitSha}},
}

newCommit, _, err := ghService.Client.Git.CreateCommit(context.Background(), repoOwner, repoName, commit, nil)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to commit Terraform file: %v", err))
log.Printf("Error committing Terraform file: %v", err)
return fmt.Errorf("error committing Terraform file: %v", err)
}

// Update the reference to point to the new commit
ref := &github.Reference{
Ref: github.String(fmt.Sprintf("refs/heads/%s", *branch)),
Object: &github.GitObject{
SHA: newCommit.SHA,
},
}
_, _, err = ghService.Client.Git.UpdateRef(context.Background(), repoOwner, repoName, ref, false)
if err != nil {
commentReporterManager.UpdateComment(fmt.Sprintf(":x: Failed to update branch reference: %v", err))
log.Printf("Error updating branch reference: %v", err)
return fmt.Errorf("error updating branch reference: %v", err)
}

commentReporterManager.UpdateComment(":white_check_mark: Successfully generated and committed Terraform code")
return nil
}
return nil
}

func TriggerDiggerJobs(ciBackend ci_backends.CiBackend, repoFullName string, repoOwner string, repoName string, batchId *uuid.UUID, prNumber int, prService ci.PullRequestService, gh utils.GithubClientProvider) error {
_, err := models.DB.GetDiggerBatch(batchId)
if err != nil {
Expand Down
4 changes: 2 additions & 2 deletions backend/utils/ai.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import (
"net/http"
)

func GenerateTerraformCode(appCode string, generationEndpoint string, webhookSecret string) (string, error) {
func GenerateTerraformCode(appCode string, generationEndpoint string, apiToken string) (string, error) {

payload := map[string]string{
"code": appCode,
Expand All @@ -28,7 +28,7 @@ func GenerateTerraformCode(appCode string, generationEndpoint string, webhookSec

// Set headers
req.Header.Set("Content-Type", "application/json")
req.Header.Set("X-Webhook-Secret", webhookSecret) // Replace with your webhook secret
req.Header.Set("Authorization", "Bearer "+apiToken)

// Make the request
client := &http.Client{}
Expand Down

0 comments on commit 34d1db1

Please sign in to comment.