fix(pipeline): beam output path can contain a filesystem scheme (#3532)

kythe/go/serving/pipeline/beamio/leveldb.go

@@ -27,6 +27,7 @@ import (
 	"path/filepath"
 	"reflect"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/apache/beam/sdks/go/pkg/beam"
@@ -119,7 +120,7 @@ func (w *writeManifest) ProcessElement(ctx context.Context, _ beam.T, e func(*ta
 	defer func(start time.Time) { log.Printf("Manifest written in %s", time.Since(start)) }(time.Now())
 
 	// Write the CURRENT manifest to the 0'th LevelDB file.
-	f, err := openWrite(ctx, filepath.Join(w.Path, manifestName))
+	f, err := openWrite(ctx, schemePreservingPathJoin(w.Path, manifestName))
 	if err != nil {
 		return 0, err
 	}
@@ -173,7 +174,7 @@ func (w *writeManifest) ProcessElement(ctx context.Context, _ beam.T, e func(*ta
 	}
 
 	// Write the CURRENT pointer to the freshly written manifest file.
-	currentFile, err := openWrite(ctx, filepath.Join(w.Path, "CURRENT"))
+	currentFile, err := openWrite(ctx, schemePreservingPathJoin(w.Path, "CURRENT"))
 	if err != nil {
 		return 0, err
 	} else if _, err := io.WriteString(currentFile, manifestName+"\n"); err != nil {
@@ -226,6 +227,18 @@ var (
 	conflictingLevelDBValuesCounter = beam.NewCounter("kythe.beamio.leveldb", "conflicting-values")
 )
 
+const schemaSeparator = "://"
+
+// schemePreservingPathJoin is like filepath.Join, but doesn't collapse
+// the double-slash in the schema prefix, if any.
+func schemePreservingPathJoin(p, f string) string {
+	parts := strings.SplitN(p, schemaSeparator, 2)
+	if len(parts) == 2 {
+		return parts[0] + schemaSeparator + filepath.Join(parts[1], f)
+	}
+	return filepath.Join(p, f)
+}
+
 // ProcessElement writes a set of KeyValues to the an SSTable per shard.  Shards
 // should be small enough to fit into memory so that they can be sorted.
 // TODO(BEAM-4405): use SortValues extension to remove in-memory requirement
@@ -276,7 +289,7 @@ func (w *writeTable) ProcessElement(ctx context.Context, shard int, kvIter func(
 	md.Last = els[len(els)-1].Key
 
 	// Write each sorted key-value to an SSTable.
-	f, err := openWrite(ctx, filepath.Join(w.Path, fmt.Sprintf("%06d.ldb", md.Shard)))
+	f, err := openWrite(ctx, schemePreservingPathJoin(w.Path, fmt.Sprintf("%06d.ldb", md.Shard)))
 	if err != nil {
 		return err
 	}

kythe/go/serving/pipeline/beamio/leveldb_test.go

@@ -106,3 +106,23 @@ func TestLevelDBSink(t *testing.T) {
 }
 
 func extendToKey(v beam.T) (beam.T, beam.T) { return v, v }
+
+func TestSchemaPreservingPathJoin(t *testing.T) {
+	exp := "gs://foo/bar/baz"
+	res := schemePreservingPathJoin("gs://foo", "bar/baz")
+	if exp != res {
+		t.Fatalf("Expected [%s], got [%s]", exp, res)
+	}
+
+	exp = "foo/bar/baz"
+	res = schemePreservingPathJoin("foo//bar", "baz")
+	if exp != res {
+		t.Fatalf("Expected [%s], got [%s]", exp, res)
+	}
+
+	exp = "/foo/bar"
+	res = schemePreservingPathJoin("/foo/", "bar")
+	if exp != res {
+		t.Fatalf("Expected [%s], got [%s]", exp, res)
+	}
+}