internal/postgres: don't touch package_imports

If the read-imports experiment is on, replace all references to the
package_imports table with the imports table: don't write into or read
from package_imports.

Change-Id: I02631aba574b9938a44b06bce855997606984939
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/327749
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Julie Qiu <julie@golang.org>

3 files changed, 58 insertions, 28 deletions

diff --git a/internal/postgres/insert_module.go b/internal/postgres/insert_module.go
index 0643d2d5..70b06868 100644
--- a/internal/postgres/insert_module.go
+++ b/internal/postgres/insert_module.go
@@ -23,6 +23,7 @@ import (
 	"golang.org/x/pkgsite/internal"
 	"golang.org/x/pkgsite/internal/database"
 	"golang.org/x/pkgsite/internal/derrors"
+	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/licenses"
 	"golang.org/x/pkgsite/internal/log"
 	"golang.org/x/pkgsite/internal/stdlib"
@@ -581,24 +582,25 @@ func insertImports(ctx context.Context, tx *database.DB,
 	pathToImports map[string][]string) (err error) {
 	defer derrors.WrapStack(&err, "insertImports")
 
-	// Insert into package_imports. This table is going to go away soon,
-	// but for now it is still the table we read from.
-	var pkgImportValues []interface{}
-	for _, pkgPath := range paths {
-		imports, ok := pathToImports[pkgPath]
-		if !ok {
-			continue
+	if !experiment.IsActive(ctx, internal.ExperimentReadImports) {
+		// Insert into package_imports. This table is going to go away soon,
+		// but for now it is still the table we read from.
+		var pkgImportValues []interface{}
+		for _, pkgPath := range paths {
+			imports, ok := pathToImports[pkgPath]
+			if !ok {
+				continue
+			}
+			unitID := pathToUnitID[pkgPath]
+			for _, toPath := range imports {
+				pkgImportValues = append(pkgImportValues, unitID, toPath)
+			}
 		}
-		unitID := pathToUnitID[pkgPath]
-		for _, toPath := range imports {
-			pkgImportValues = append(pkgImportValues, unitID, toPath)
+		pkgImportCols := []string{"unit_id", "to_path"}
+		if err := tx.BulkUpsert(ctx, "package_imports", pkgImportCols, pkgImportValues, pkgImportCols); err != nil {
+			return err
 		}
 	}
-	pkgImportCols := []string{"unit_id", "to_path"}
-	if err := tx.BulkUpsert(ctx, "package_imports", pkgImportCols, pkgImportValues, pkgImportCols); err != nil {
-		return err
-	}
-
 	// Insert into imports. This will eventually replace package_imports.
 	importPathSet := map[string]bool{}
 	for _, pkgPath := range paths {
@@ -752,16 +754,31 @@ func (db *DB) ReInsertLatestVersion(ctx context.Context, modulePath string) (err
 		}
 
 		// Insert this version's imports into imports_unique.
-		if _, err := tx.Exec(ctx, `
-			INSERT INTO imports_unique (from_path, from_module_path, to_path)
-			SELECT p.path, m.module_path, i.to_path
-			FROM units u
-			INNER JOIN package_imports i ON (u.id = i.unit_id)
-			INNER JOIN paths p ON (p.id = u.path_id)
-			INNER JOIN modules m ON (m.id=u.module_id)
-			WHERE m.module_path = $1 and m.version = $2
+		if experiment.IsActive(ctx, internal.ExperimentReadImports) {
+			if _, err := tx.Exec(ctx, `
+				INSERT INTO imports_unique (from_path, from_module_path, to_path)
+				SELECT p1.path, m.module_path, p2.path
+				FROM units u
+				INNER JOIN imports i ON u.id = i.unit_id
+				INNER JOIN paths p1 ON p1.id = u.path_id
+				INNER JOIN modules m ON m.id = u.module_id
+				INNER JOIN paths p2 ON p2.id = i.to_path_id
+				WHERE m.module_path = $1 and m.version = $2
 		`, modulePath, lmv.GoodVersion); err != nil {
-			return err
+				return err
+			}
+		} else {
+			if _, err := tx.Exec(ctx, `
+				INSERT INTO imports_unique (from_path, from_module_path, to_path)
+				SELECT p.path, m.module_path, i.to_path
+				FROM units u
+				INNER JOIN package_imports i ON (u.id = i.unit_id)
+				INNER JOIN paths p ON (p.id = u.path_id)
+				INNER JOIN modules m ON (m.id=u.module_id)
+				WHERE m.module_path = $1 and m.version = $2
+		`, modulePath, lmv.GoodVersion); err != nil {
+				return err
+			}
 		}
 
 		log.Debugf(ctx, "ReInsertLatestVersion(%q): re-inserted at latest good version %s", modulePath, lmv.GoodVersion)
diff --git a/internal/postgres/insert_module_test.go b/internal/postgres/insert_module_test.go
index 5b26656c..981e7780 100644
--- a/internal/postgres/insert_module_test.go
+++ b/internal/postgres/insert_module_test.go
@@ -24,6 +24,7 @@ import (
 	"golang.org/x/pkgsite/internal"
 	"golang.org/x/pkgsite/internal/database"
 	"golang.org/x/pkgsite/internal/derrors"
+	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/licenses"
 	"golang.org/x/pkgsite/internal/source"
 	"golang.org/x/pkgsite/internal/stdlib"
@@ -659,7 +660,15 @@ func TestIsAlternativeModulePath(t *testing.T) {
 
 func TestReInsertLatestVersion(t *testing.T) {
 	t.Parallel()
-	ctx := context.Background()
+	t.Run("package_imports", func(t *testing.T) {
+		testReInsertLatestVersion(t, context.Background())
+	})
+	t.Run("imports", func(t *testing.T) {
+		testReInsertLatestVersion(t, experiment.NewContext(context.Background(), internal.ExperimentReadImports))
+	})
+}
+
+func testReInsertLatestVersion(t *testing.T, ctx context.Context) {
 	testDB, release := acquire(t)
 	defer release()
 
diff --git a/internal/postgres/unit.go b/internal/postgres/unit.go
index b2528f48..6007edf2 100644
--- a/internal/postgres/unit.go
+++ b/internal/postgres/unit.go
@@ -471,7 +471,11 @@ func (db *DB) getUnitWithAllFields(ctx context.Context, um *internal.UnitMeta, b
 		}
 	}
 	// Get README, documentation and import counts.
-	query := `
+	importTableName := "package_imports"
+	if experiment.IsActive(ctx, internal.ExperimentReadImports) {
+		importTableName = "imports"
+	}
+	query := fmt.Sprintf(`
         SELECT
 			r.file_path,
 			r.contents,
@@ -479,7 +483,7 @@ func (db *DB) getUnitWithAllFields(ctx context.Context, um *internal.UnitMeta, b
 			d.source,
 			COALESCE((
 				SELECT COUNT(unit_id)
-				FROM package_imports
+				FROM %s
 				WHERE unit_id = u.id
 				GROUP BY unit_id
 				), 0) AS num_imports,
@@ -501,7 +505,7 @@ func (db *DB) getUnitWithAllFields(ctx context.Context, um *internal.UnitMeta, b
         ) d
 		ON d.unit_id = u.id
 		WHERE u.id = $2
-	`
+	`, importTableName)
 	var (
 		r internal.Readme
 		u internal.Unit