Hyper-Converged Database (HCD) quickstart for collections (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")
    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


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

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

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

    print(f"Created keyspace")


if __name__ == "__main__":
    main()

from astrapy.constants import VectorMetric
from astrapy.info import (
    CollectionDefinition,
    CollectionVectorOptions,
)
from quickstart_connect import connect_to_database


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

    collection = database.create_collection(
        "quickstart_collection",
        keyspace="quickstart_keyspace",
        definition=CollectionDefinition(
            vector=CollectionVectorOptions(
                dimension=5,
                metric=VectorMetric.COSINE,
            ),
        ),
    )

    print(f"Created collection {collection.full_name}")


if __name__ == "__main__":
    main()

import json

from astrapy.data_types import DataAPIDate
from quickstart_connect import connect_to_database


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

    collection = database.get_collection(
        "quickstart_collection", keyspace="quickstart_keyspace"
    )

    data_file_path = "PATH_TO_DATA_FILE"

    # Read the JSON file and parse it into a JSON array
    with open(data_file_path, "r", encoding="utf8") as file:
        json_data = json.load(file)

    # Assemble the documents to insert
    documents = [
        {
            **data,
            # Convert the date string into a DataAPIDate
            "due_date": (
                DataAPIDate.from_string(data["due_date"])
                if data.get("due_date")
                else None
            ),
            # Populate the reserved $vector field
            "$vector": data["summary_genres_vector"],
        }
        for data in json_data
    ]

    # Insert the data
    inserted = collection.insert_many(documents)

    print(f"Inserted {len(inserted.inserted_ids)} documents.")


if __name__ == "__main__":
    main()

from quickstart_connect import connect_to_database


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

    collection = database.get_collection(
        "quickstart_collection", keyspace="quickstart_keyspace"
    )

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

    rating_cursor = collection.find({"rating": {"$gt": 4.7}})

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

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

    single_vector_match = collection.find_one(
        sort={
            "$vector": [
                0.016326904,
                -0.031677246,
                0.04815674,
                0.0033435822,
                0.01876831,
            ]
        }
    )

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

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

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

    for document in vector_cursor:
        print(document)


if __name__ == "__main__":
    main()