Hyper-Converged Database (HCD) quickstart for tables (Python)

import os

from astrapy import DataAPIClient, Database
from astrapy.authentication import UsernamePasswordTokenProvider
from astrapy.constants import Environment


def connect_to_database() -> Database:
    """
    Connects to your database.
    This function retrieves the database endpoint, username, and password from the
    environment variables `API_ENDPOINT`, `USERNAME`, and `PASSWORD`.

    Returns:
        Database: An instance of the connected database.

    Raises:
        RuntimeError: If the environment variables `API_ENDPOINT`,
        `USERNAME`, or `PASSWORD` are not defined.
    """
    endpoint = os.environ.get("API_ENDPOINT")  (1)
    username = os.environ.get("USERNAME")
    password = os.environ.get("PASSWORD")

    if not endpoint or not username or not password:
        raise RuntimeError(
            "Environment variables API_ENDPOINT, USERNAME, and PASSWORD must be defined"
        )

    # Create an instance of the `DataAPIClient` class
    client = DataAPIClient(environment=Environment.HCD)

    # Get the database specified by your endpoint and provide the token
    database = client.get_database(
        endpoint,
        token=UsernamePasswordTokenProvider(username, password),
    )

    print(f"Connected to database {database.info().name}")

    return database

from quickstart_connect import connect_to_database  (1)


def main() -> None:
    database = connect_to_database()

    # Get an admin object
    admin = database.get_database_admin()

    # Create a keyspace
    admin.create_keyspace("quickstart_keyspace")  (2)

    print(f"Created keyspace")


if __name__ == "__main__":
    main()

from astrapy.constants import VectorMetric
from astrapy.info import (
    ColumnType,
    CreateTableDefinition,
    TableVectorIndexOptions,
)
from quickstart_connect import connect_to_database  (1)


def main() -> None:
    database = connect_to_database()

    table_definition = (
        CreateTableDefinition.builder()
        # Define all of the columns in the table
        .add_column("title", ColumnType.TEXT)
        .add_column("author", ColumnType.TEXT)
        .add_column("number_of_pages", ColumnType.INT)
        .add_column("rating", ColumnType.FLOAT)
        .add_column("publication_year", ColumnType.INT)
        .add_column("summary", ColumnType.TEXT)
        .add_set_column(
            "genres",
            ColumnType.TEXT,
        )
        .add_map_column(
            "metadata",
            # This is the key type for the map column
            ColumnType.TEXT,
            # This is the value type for the map column
            ColumnType.TEXT,
        )
        .add_column("is_checked_out", ColumnType.BOOLEAN)
        .add_column("borrower", ColumnType.TEXT)
        .add_column("due_date", ColumnType.DATE)
        # This column will store vector embeddings. (2)
        .add_vector_column("summary_genres_vector", dimension=5)
        # Define the primary key for the table.
        # In this case, the table uses a composite primary key.
        .add_partition_by(["title", "author"])
        # Finally, build the table definition.
        .build()
    )

    table = database.create_table(
        "quickstart_table",  (3)
        keyspace="quickstart_keyspace",  (4)
        definition=table_definition,
    )

    print("Created table")

    # Index any columns that you want to sort and filter on.
    table.create_index(
        "rating_index",
        column="rating",
    )

    table.create_index(
        "number_of_pages_index",
        column="number_of_pages",
    )

    table.create_vector_index(
        "summary_genres_vector_index",
        column="summary_genres_vector",
        options=TableVectorIndexOptions(
            metric=VectorMetric.COSINE,  (5)
        ),
    )

    print("Indexed columns")


if __name__ == "__main__":
    main()

import json

from astrapy.data_types import DataAPIDate, DataAPIVector
from quickstart_connect import connect_to_database  (1)


def main() -> None:
    database = connect_to_database()

    table = database.get_table(
        "quickstart_table", keyspace="quickstart_keyspace"
    )  (2)

    data_file_path = "PATH_TO_DATA_FILE"  (3)

    with open(data_file_path, "r", encoding="utf8") as file:
        json_data = json.load(file)

    rows = [
        {
            **data,
            "due_date": (
                DataAPIDate.from_string(data["due_date"])
                if data.get("due_date")
                else None
            ),
            "summary_genres_vector": DataAPIVector(
                data["summary_genres_vector"]
            ),
        }
        for data in json_data
    ]

    insert_result = table.insert_many(rows)

    print(f"Inserted {len(insert_result.inserted_ids)} rows")


if __name__ == "__main__":
    main()

from astrapy.data_types import DataAPIVector
from quickstart_connect import connect_to_database  (1)


def main() -> None:
    database = connect_to_database()

    table = database.get_table(
        "quickstart_table", keyspace="quickstart_keyspace"
    )  (2)

    # Find rows that match a filter
    print("\nFinding books with rating greater than 4.7...")

    rating_cursor = table.find(
        {"rating": {"$gt": 4.7}},
        projection={"title": True, "rating": True},
    )

    for row in rating_cursor:
        print(f"{row['title']} is rated {row['rating']}")

    # Perform a vector search to find the closest match to a search vector
    print("\nUsing vector search to find a book...")

    single_vector_match = table.find_one(
        {},
        sort={
            "summary_genres_vector": DataAPIVector(
                [
                    0.016326904,
                    -0.031677246,
                    0.04815674,
                    0.0033435822,
                    0.01876831,
                ]
            )
        },
        projection={"title": True},
    )

    print(f"{single_vector_match['title']} is the best match")

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

    vector_cursor = table.find(
        {"number_of_pages": {"$gt": 400}},
        sort={
            "summary_genres_vector": DataAPIVector(
                [
                    0.016326904,
                    -0.031677246,
                    0.04815674,
                    0.0033435822,
                    0.01876831,
                ]
            )
        },
        limit=3,
        projection={"title": True, "author": True},
    )

    for row in vector_cursor:
        print(row)


if __name__ == "__main__":
    main()