Astra DB Serverless quickstart for tables (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 DataStax Astra database.
// This function retrieves the database endpoint and application token
// from the environment variables `API_ENDPOINT` and `APPLICATION_TOKEN`.
//
// Returns an instance of the connected database.
// Exits with an error if the environment variables
// `API_ENDPOINT` or `APPLICATION_TOKEN` are not defined.
func ConnectToDatabase() *astra.Db {
	endpoint := os.Getenv("API_ENDPOINT") (1)
	token := os.Getenv("APPLICATION_TOKEN")

	if token == "" || endpoint == "" {
		log.Fatal(
			"Environment variables API_ENDPOINT and APPLICATION_TOKEN must be defined.",
		)
	}

	// Create an instance of `DataAPIClient`
	client := astra.NewClient()

	// Get the database specified by your endpoint and provide the token
	database := client.Database(endpoint, options.API().SetToken(token))

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

	return database
}

package main

import (
	"context"
	"fmt"
	"log"

	"quickstart/shared"

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

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

	database := shared.ConnectToDatabase() (1)

	// Define all of the columns in the table
	definition := table.Definition{
		Columns: table.Columns{
			{Name: "title", Column: table.Text()},
			{Name: "author", Column: table.Text()},
			{Name: "number_of_pages", Column: table.Int()},
			{Name: "rating", Column: table.Float()},
			{Name: "publication_year", Column: table.Int()},
			{Name: "summary", Column: table.Text()},
			{Name: "genres", Column: table.Set(table.Text())},
			{Name: "metadata", Column: table.Map("text", table.Text())},
			{Name: "is_checked_out", Column: table.Boolean()},
			{Name: "borrower", Column: table.Text()},
			{Name: "due_date", Column: table.Date()},
			// This column will store vector embeddings.
			// The column will use an embedding model from NVIDIA to
			// generate the
			// vector embeddings when data is inserted to the column. (2)
			{
				Name: "summary_genres_vector",
				Column: table.VectorWithService(
					1024,
					&table.VectorService{
						Provider:  "nvidia",
						ModelName: "nvidia/nv-embedqa-e5-v5",
					},
				),
			},
		},
		// Define the primary key for the table.
		// In this case, the table uses a composite primary key.
		PrimaryKey: table.PrimaryKey{
			PartitionBy: []string{"title", "author"},
		},
	}

	table, err := database.CreateTable(
		ctx,
		"quickstart_table", (3)
		definition,
	)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("Created table")

	// Index any columns that you want to sort and filter on.
	err = table.CreateIndex(ctx, "rating_index", "rating")
	if err != nil {
		log.Fatal(err)
	}

	err = table.CreateIndex(
		ctx,
		"number_of_pages_index",
		"number_of_pages",
	)
	if err != nil {
		log.Fatal(err)
	}

	err = table.CreateVectorIndex(
		ctx,
		"summary_genres_vector_index",
		"summary_genres_vector",
		options.CreateVectorIndex().SetMetric(options.MetricCosine),
	)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("Indexed columns")
}

package main

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

	"quickstart/shared"

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

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

	database := shared.ConnectToDatabase() (1)

	table := database.Table("quickstart_table") (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)
	}

	rows := make([]astra.Row, len(jsonData))
	for i, data := range jsonData {
		// Convert due_date string to DateOnly type if present
		if dueDateStr, ok := data["due_date"].(string); ok &&
			dueDateStr != "" {
			dueDate, err := datatypes.ParseDateOnly(dueDateStr)
			if err != nil {
				log.Fatalf("invalid due_date %q: %v", dueDateStr, err)
			}
			data["due_date"] = dueDate
		}

		// Create the summary_genres_vector text to vectorize
		summary, _ := data["summary"].(string)

		genres := []string{}
		if g, ok := data["genres"].([]any); ok {
			for _, genre := range g {
				if genreStr, ok := genre.(string); ok {
					genres = append(genres, genreStr)
				}
			}
		}

		data["summary_genres_vector"] = fmt.Sprintf(
			"summary: %s | genres: %s",
			summary,
			strings.Join(genres, ", "),
		)

		rows[i] = astra.NewRow(data)
	}

	insertedResult, err := table.InsertMany(ctx, rows)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf(
		"Inserted %d rows.\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)

	table := database.Table("quickstart_table") (2)

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

	ratingCursor := table.Find(
		filter.Gt("rating", 4.7),
		options.TableFind().
			SetLimit(10).
			SetProjection(map[string]any{"title": true, "rating": true}),
	)
	defer ratingCursor.Close()

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

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

	// Perform a vector search to find the closest match to a search string
	fmt.Println("\nUsing vector search to find a single scary novel...")

	var singleVectorMatch astra.Row
	err := table.FindOne(
		ctx,
		nil,
		options.TableFindOne().
			SetSort(sort.Table.Vectorize("summary_genres_vector", "A scary novel")).
			SetProjection(map[string]any{"title": true}),
	).Decode(&singleVectorMatch)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("%s is a scary novel\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 string
	// and just return the title and author
	fmt.Println(
		"\nUsing filters and vector search to find 3 books with more than 400 pages that are set in the arctic, returning just the title and author...",
	)

	vectorCursor := table.Find(
		filter.Gt("number_of_pages", 400),
		options.TableFind().
			SetSort(sort.Table.Vectorize("summary_genres_vector", "A book set in the arctic")).
			SetLimit(3).
			SetProjection(map[string]any{"title": true, "author": true}),
	)
	defer vectorCursor.Close()

	for vectorCursor.Next(ctx) {
		var row astra.Row
		if err := vectorCursor.Decode(&row); err != nil {
			log.Fatal(err)
		}
		fmt.Printf(
			"Title: %s, Author: %s\n",
			row.MustGet("title"),
			row.MustGet("author"),
		)
	}

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