Hyper-Converged Database (HCD) quickstart for collections (Go)

package shared

import (
	"fmt"
	"log"
	"os"

	"github.com/datastax/astra-db-go/v2/astra"
	"github.com/datastax/astra-db-go/v2/astra/options"
)

// ConnectToDatabase connects to a database.
// This function retrieves the database endpoint, username, and password
// from the environment variables `API_ENDPOINT`, `USERNAME`, and
// `PASSWORD`.
//
// Returns an instance of the connected database.
// Exits with an error if the environment variables
// `API_ENDPOINT`, `USERNAME`, or `PASSWORD` are not defined.
func ConnectToDatabase() *astra.Db {
	endpoint := os.Getenv("API_ENDPOINT") (1)
	username := os.Getenv("USERNAME")
	password := os.Getenv("PASSWORD")

	if endpoint == "" || username == "" || password == "" {
		log.Fatal(
			"Environment variables API_ENDPOINT, USERNAME, and PASSWORD must be defined.",
		)
	}

	// Create an instance of DataAPIClient
	client := astra.NewClient(
		options.API().SetEnvironment(options.EnvironmentHCD),
	)

	// Get the database specified by your endpoint and provide the token
	database := client.Database(
		endpoint,
		options.API().
			SetUsernamePasswordTokenProvider(
				"**USERNAME**",
				"**PASSWORD**",
			),
	)

	fmt.Printf("Connected to database %s\n", database.Endpoint())

	return database
}

package main

import (
	"context"
	"fmt"
	"log"

	"quickstart/shared"
)

func main() {
	ctx := context.Background()

	database := shared.ConnectToDatabase() (1)

	// Get an admin object
	admin, err := database.DatabaseAdmin()

	if err != nil {
		log.Fatal(err)
	}

	// Create a keyspace
	err = admin.CreateKeyspace(ctx, "quickstart_keyspace") (2)

	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("Created keyspace")
}

package main

import (
	"context"
	"fmt"
	"log"

	"quickstart/shared"

	"github.com/datastax/astra-db-go/v2/astra/options"
)

func main() {
	ctx := context.Background()
	database := shared.ConnectToDatabase() (1)

	collection, err := database.CreateCollection(
		ctx,
		"quickstart_collection", (2)
		options.CreateCollection().
			SetKeyspace("quickstart_keyspace"). (3)
			UpdateVector(
				options.Vector(). (4)
							SetDimension(5).
							SetMetric(options.MetricCosine),
			),
	)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("Created collection %s\n", collection.Name())
}

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"time"

	"quickstart/shared"

	"github.com/datastax/astra-db-go/v2/astra"
	"github.com/datastax/astra-db-go/v2/astra/options"
)

func main() {
	ctx := context.Background()
	database := shared.ConnectToDatabase() (1)

	collection := database.Collection(
		"quickstart_collection",
		options.GetCollection().SetKeyspace("quickstart_keyspace"),
	) (2)

	dataFilePath := "PATH_TO_DATA_FILE" (3)

	// Read the JSON file and parse it into a JSON array
	rawData, err := os.ReadFile(dataFilePath)
	if err != nil {
		log.Fatal(err)
	}

	var jsonData []map[string]any
	if err := json.Unmarshal(rawData, &jsonData); err != nil {
		log.Fatal(err)
	}

	// Assemble the documents to insert:
	// - Convert the date string into a time.Time
	// - Add a $vector field
	documents := make([]astra.Document, len(jsonData))
	for i, data := range jsonData {
		// Convert due_date string to time.Time if present
		if dueDateStr, ok := data["due_date"].(string); ok &&
			dueDateStr != "" {
			dueDate, err := time.Parse("2006-01-02", dueDateStr)
			if err != nil {
				log.Fatalf("invalid due_date %q: %v", dueDateStr, err)
			}
			data["due_date"] = dueDate
		}

		// Populate the reserved $vector field
		if vector, ok := data["summary_genres_vector"]; ok {
			data["$vector"] = vector
		} else {
			log.Fatalf("record %d missing summary_genres_vector", i)
		}

		documents[i] = astra.NewDocument(data)
	}

	// Insert the data
	insertedResult, err := collection.InsertMany(
		ctx,
		documents,
	)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf(
		"Inserted %d documents.\n",
		insertedResult.InsertedCount(),
	)
}

package main

import (
	"context"
	"fmt"
	"log"

	"quickstart/shared"

	"github.com/datastax/astra-db-go/v2/astra"
	"github.com/datastax/astra-db-go/v2/astra/filter"
	"github.com/datastax/astra-db-go/v2/astra/options"
	"github.com/datastax/astra-db-go/v2/astra/sort"
)

func main() {
	ctx := context.Background()

	database := shared.ConnectToDatabase() (1)

	collection := database.Collection(
		"quickstart_collection",
		options.GetCollection().SetKeyspace("quickstart_keyspace"),
	) (2)

	// Find documents that match a filter
	fmt.Println("\nFinding books with rating greater than 4.7...")

	ratingCursor := collection.Find(
		filter.Gt("rating", 4.7),
		options.CollectionFind().SetLimit(10),
	)
	defer ratingCursor.Close()

	for ratingCursor.Next(ctx) {
		var document astra.Document
		if err := ratingCursor.Decode(&document); err != nil {
			log.Fatal(err)
		}
		fmt.Printf(
			"%v is rated %v\n",
			document.MustGet("title"),
			document.MustGet("rating"),
		)
	}

	if err := ratingCursor.Err(); err != nil {
		log.Fatal(err)
	}

	// Perform a vector search to find the closest match to a given vector
	fmt.Println("\nUsing vector search to find a book...")

	var singleVectorMatch astra.Document
	err := collection.FindOne(
		ctx,
		nil,
		options.CollectionFindOne().
			SetSort(sort.Vector([]float32{0.016326904, -0.031677246, 0.04815674, 0.0033435822, 0.01876831})),
	).Decode(&singleVectorMatch)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf(
		"%v is the best match\n",
		singleVectorMatch.MustGet("title"),
	)

	// Combine a filter, vector search, and projection to find the 3 books
	// with
	// more than 400 pages that are the closest matches to a search vector,
	// and just return the title and author
	fmt.Println(
		"\nUsing filters and vector search to find 3 books with more than 400 pages, returning just the title and author...",
	)

	vectorCursor := collection.Find(
		filter.Gt("number_of_pages", 400),
		options.CollectionFind().
			SetSort(sort.Vector([]float32{0.016326904, -0.031677246, 0.04815674, 0.0033435822, 0.01876831})).
			SetLimit(3).
			SetProjection(map[string]any{"title": true, "author": true}),
	)
	defer vectorCursor.Close()

	for vectorCursor.Next(ctx) {
		var document astra.Document
		if err := vectorCursor.Decode(&document); err != nil {
			log.Fatal(err)
		}
		fmt.Println(document.ToMap())
	}

	if err := vectorCursor.Err(); err != nil {
		log.Fatal(err)
	}
}