From 8f933ce76a0fffe05a1ab6d73678bd4436e6b433 Mon Sep 17 00:00:00 2001
From: Neil Macneale IV <neil@dolthub.com>
Date: Tue, 11 Feb 2025 17:35:35 -0800
Subject: [PATCH] Clone tests passing again

---
 go/libraries/doltcore/remotestorage/chunk_store.go | 12 +++++++++++-
 go/store/datas/pull/clone.go                       |  7 +------
 go/store/nbs/store.go                              | 10 +++++++---
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/go/libraries/doltcore/remotestorage/chunk_store.go b/go/libraries/doltcore/remotestorage/chunk_store.go
index 02b06060b5..75fa614c2c 100644
--- a/go/libraries/doltcore/remotestorage/chunk_store.go
+++ b/go/libraries/doltcore/remotestorage/chunk_store.go
@@ -1051,6 +1051,11 @@ func (dcs *DoltChunkStore) SupportedOperations() chunks.TableFileStoreOps {
 
 // WriteTableFile reads a table file from the provided reader and writes it to the chunk store.
 func (dcs *DoltChunkStore) WriteTableFile(ctx context.Context, fileId string, numChunks int, contentHash []byte, getRd func() (io.ReadCloser, uint64, error)) error {
+	// Err if the suffix is an archive file
+	if strings.HasSuffix(fileId, nbs.ArchiveFileSuffix) {
+		return errors.New("cannot write archive file ids currently.")
+	}
+
 	fileIdBytes := hash.Parse(fileId)
 	err := dcs.uploadTableFileWithRetries(ctx, fileIdBytes, uint64(numChunks), contentHash, getRd)
 	if err != nil {
@@ -1153,7 +1158,12 @@ func (drtf DoltRemoteTableFile) LocationPrefix() string {
 }
 
 func (drtf DoltRemoteTableFile) LocationSuffix() string {
-	panic("implement me LocationSuffix")
+	u, _ := url.Parse(drtf.info.Url)
+	if strings.HasSuffix(u.Path, nbs.ArchiveFileSuffix) {
+		return nbs.ArchiveFileSuffix
+	}
+
+	return ""
 }
 
 // FileID gets the id of the file
diff --git a/go/store/datas/pull/clone.go b/go/store/datas/pull/clone.go
index 8d9f9f638f..27a3ca9405 100644
--- a/go/store/datas/pull/clone.go
+++ b/go/store/datas/pull/clone.go
@@ -19,10 +19,8 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"strings"
 
 	"github.com/cenkalti/backoff/v4"
-	"github.com/dolthub/dolt/go/store/nbs"
 	"golang.org/x/sync/errgroup"
 	"golang.org/x/sync/semaphore"
 
@@ -84,9 +82,6 @@ func mapTableFiles(tblFiles []chunks.TableFile) ([]string, map[string]chunks.Tab
 
 	for i, tblFile := range tblFiles {
 		fileId := tblFile.FileID()
-		if strings.HasSuffix(fileId, nbs.ArchiveFileSuffix) {
-			fileId = fileId[:len(fileId)-len(nbs.ArchiveFileSuffix)]
-		}
 
 		fileIDtoTblFile[fileId] = tblFile
 		fileIds[i] = fileId
@@ -143,7 +138,7 @@ func clone(ctx context.Context, srcTS, sinkTS chunks.TableFileStore, sinkCS chun
 				}
 
 				report(TableFileEvent{EventType: DownloadStart, TableFiles: []chunks.TableFile{tblFile}})
-				err = sinkTS.WriteTableFile(ctx, tblFile.FileID(), tblFile.NumChunks(), nil, func() (io.ReadCloser, uint64, error) {
+				err = sinkTS.WriteTableFile(ctx, tblFile.FileID()+tblFile.LocationSuffix(), tblFile.NumChunks(), nil, func() (io.ReadCloser, uint64, error) {
 					rd, contentLength, err := tblFile.Open(ctx)
 					if err != nil {
 						return nil, 0, err
diff --git a/go/store/nbs/store.go b/go/store/nbs/store.go
index 0e4e3a3d97..a0386bd4b6 100644
--- a/go/store/nbs/store.go
+++ b/go/store/nbs/store.go
@@ -234,8 +234,12 @@ func (nbs *NomsBlockStore) GetChunkLocations(ctx context.Context, hashes hash.Ha
 	}
 	res := make(map[string]map[hash.Hash]Range, len(hashes))
 	for cs, ranges := range sourcesToRanges {
+		suffix := ""
+		if _, ok := cs.(archiveChunkSource); ok {
+			suffix = ArchiveFileSuffix
+		}
 		h := cs.hash()
-		res[h.String()] = ranges // NOT SURE about this. NM4.
+		res[h.String()+suffix] = ranges
 	}
 	return res, nil
 }
@@ -1662,7 +1666,7 @@ func (nbs *NomsBlockStore) Path() (string, bool) {
 }
 
 // WriteTableFile will read a table file from the provided reader and write it to the TableFileStore
-func (nbs *NomsBlockStore) WriteTableFile(ctx context.Context, fileId string, numChunks int, contentHash []byte, getRd func() (io.ReadCloser, uint64, error)) error {
+func (nbs *NomsBlockStore) WriteTableFile(ctx context.Context, fileName string, numChunks int, contentHash []byte, getRd func() (io.ReadCloser, uint64, error)) error {
 	tfp, ok := nbs.p.(tableFilePersister)
 	if !ok {
 		return errors.New("Not implemented")
@@ -1673,7 +1677,7 @@ func (nbs *NomsBlockStore) WriteTableFile(ctx context.Context, fileId string, nu
 		return err
 	}
 	defer r.Close()
-	return tfp.CopyTableFile(ctx, r, fileId, sz, uint32(numChunks))
+	return tfp.CopyTableFile(ctx, r, fileName, sz, uint32(numChunks))
 }
 
 // AddTableFilesToManifest adds table files to the manifest