From 7857dfe6341f9c143aed06cedf482d4d30512d3a Mon Sep 17 00:00:00 2001
From: Roy Reznik <roy@wiz.io>
Date: Tue, 28 Mar 2023 23:14:30 +0300
Subject: [PATCH] Cleanup only non-gqlgen packages when reloading all packages

---
 internal/code/packages.go | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/internal/code/packages.go b/internal/code/packages.go
index 3505920cbc..0b25750a52 100644
--- a/internal/code/packages.go
+++ b/internal/code/packages.go
@@ -7,10 +7,16 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"runtime/debug"
+	"strings"
+	"sync"
 
 	"golang.org/x/tools/go/packages"
 )
 
+var once = sync.Once{}
+var modInfo *debug.BuildInfo
+
 var mode = packages.NeedName |
 	packages.NeedFiles |
 	packages.NeedImports |
@@ -31,10 +37,38 @@ type Packages struct {
 	numNameCalls int // stupid test steam. ignore.
 }
 
+func (p *Packages) CleanupUserPackages() {
+	once.Do(func() {
+		var ok bool
+		modInfo, ok = debug.ReadBuildInfo()
+		if !ok {
+			modInfo = nil
+		}
+	})
+
+	// Don't cleanup github.com/99designs/gqlgen prefixed packages, they haven't changed and do not need to be reloaded
+	if modInfo != nil {
+		var toRemove []string
+		for k := range p.packages {
+			if !strings.HasPrefix(k, modInfo.Main.Path) {
+				toRemove = append(toRemove, k)
+			}
+		}
+
+		for _, k := range toRemove {
+			delete(p.packages, k)
+		}
+	} else {
+		p.packages = nil // Cleanup all packages if we don't know for some reason which ones to keep
+	}
+}
+
 // ReloadAll will call LoadAll after clearing the package cache, so we can reload
 // packages in the case that the packages have changed
 func (p *Packages) ReloadAll(importPaths ...string) []*packages.Package {
-	p.packages = nil
+	if p.packages != nil {
+		p.CleanupUserPackages()
+	}
 	return p.LoadAll(importPaths...)
 }