Python client usage

from astrapy.data_types import DataAPIVector

# Initialize a vector
vector = DataAPIVector([0.1, -0.2, 0.3])

# Access a component by index
print(vector[1])

# Loop over the vector components
for x in vector:
    print(x)

# Compute the vector Euclidean norm
print(sum(x*x for x in vector)**0.5)

# When inserting to a collection, binary encoding is always used
collection.insert_one({"$vector": DataAPIVector([1, 2, 3])})
collection.insert_one({"$vector": [4, 5, 6]})

# When reading from a collection, DataAPIVector is always returned
# The following outputs (reformatted for clarity):
#   [
#       {'$vector': DataAPIVector([4.0, 5.0, 6.0])},
#       {'$vector': DataAPIVector([1.0, 2.0, 3.0])}
#   ]
collection.find({}, projection={"_id": False, "$vector": True}).to_list()

# When inserting to a table, binary encoding is only used with DataAPIVector
my_table.insert_one({'primary_column': 'A', 'vector_column': DataAPIVector([9, 8, 7])})
my_table.insert_one({'primary_column': 'B', 'vector_column': [6, 5, 4]})

# When reading from a table, DataAPIVector is always returned
# The following outputs (reformatted for clarity):
#   [
#       {'primary_column': 'B', 'vector_column': DataAPIVector([6.0, 5.0, 4.0])},
#       {'primary_column': 'A', 'vector_column': DataAPIVector([9.0, 8.0, 7.0])}
#   ]
my_table.find({}).to_list()

cursor = my_table.find(
    {"match_id": "challenge6"},
    projection={"winner": True},
    limit=3,
)
for row in cursor:
    print(row)
# Output:
#   {'winner': 'Donna'}
#   {'winner': 'Erick'}
#   {'winner': 'Fiona'}

rows = my_table.find(
    {"match_id": "challenge6"},
    projection={"winner": True},
    limit=3,
).to_list()
print(rows)
# [{'winner': 'Donna'}, {'winner': 'Erick'}, {'winner': 'Fiona'}]

# A simple cursor from 'find':
cursor = my_table.find(
    {"match_id": "challenge6"},
    projection={"winner": True},
    limit=3,
)
for row in cursor:
    print(row)
# Output:
#   {'winner': 'Donna'}
#   {'winner': 'Erick'}
#   {'winner': 'Fiona'}

# Applying 'map' to a cursor:
cursor_mapped = my_table.find(
    {"match_id": "challenge6"},
    projection={"winner": True},
    limit=3,
).map(lambda row: row["winner"])
for value in cursor_mapped:
    print(value)
# Output:
#   Donna
#   Erick
#   Fiona

# Applying 'map' twice:
cursor_mapped_twice = my_table.find(
    {"match_id": "challenge6"},
    projection={"winner": True},
    limit=3,
).map(lambda row: row["winner"]).map(lambda w: w.upper())
for value in cursor_mapped_twice:
    print(value)
# Output:
#   DONNA
#   ERICK
#   FIONA

# Calling 'to_list' on an IDLE cursor:
cursor_tl = my_table.find(
    {"match_id": "challenge6"},
    projection={"winner": True},
    limit=2,
)
print(cursor_tl.to_list())
# Output:
#   [{'winner': 'Donna'}, {'winner': 'Erick'}]
print(cursor_tl.state)
# Output:
#   FindCursorState.CLOSED

# Calling 'to_list' on a partially-consumed (ACTIVE) cursor:
cursor_pc = my_table.find(
    {"match_id": "challenge6"},
    projection={"winner": True},
    limit=2,
)
cursor_pc.next()
# Output:
#   {'winner': 'Donna'}
print(cursor_pc.consumed)
# Output:
#   1
print(cursor_pc.state)
# Output:
#   FindCursorState.STARTED
print(cursor_pc.to_list())
# Output:
#   [{'winner': 'Erick'}]

# Calling 'for_each' across the whole of a cursor:
cursor_fe = my_table.find(
    {"match_id": "challenge6"},
    projection={"winner": True},
    limit=3,
)

def printer(row):
    print(f"-> {row['winner']}")

cursor_fe.for_each(printer)
# Output:
#   -> Donna
#   -> Erick
#   -> Fiona
print(cursor_fe.state)
# Output:
#   FindCursorState.CLOSED

# Calling 'for_each' with an early-stop callback:
cursor_es = my_table.find(
    {"match_id": "challenge6"},
    projection={"winner": True},
    limit=3,
)

def printer_es(row):
    go_on = row["winner"] != "Erick"
    print(f"-> {row['winner']} (go_on={go_on})")
    return go_on

cursor_es.for_each(printer_es)
# Output:
#   -> Donna (go_on=True)
#   -> Erick (go_on=False)
print(cursor_es.consumed)
# Output:
#   2
print(cursor_es.to_list())
# Output:
#   [{'winner': 'Fiona'}]

import asyncio
import os

from astrapy import DataAPIClient


async def amain():
    # get an ordinary DataAPIClient as the starting point
    my_client = DataAPIClient(token=os.environ["ASTRA_DB_APPLICATION_TOKEN"])

    # get an ordinary AstraDBAdmin and call an async method on it
    admin = my_client.get_admin()
    db_list = await admin.async_list_databases()
    print(f"Databases: {', '.join(db.name for db in db_list)}")

    # Get an AsyncDatabase from the client
    adatabase = my_client.get_async_database(
        os.environ["ASTRA_DB_API_ENDPOINT"],
    )

    # Create a collection on DB and get an AsyncCollection object:
    acollection = await adatabase.create_collection("my_collection")

    # Use the AsyncCollection object:
    insertion_result = await acollection.insert_many(
        [{"_id": i} for i in range(5)]
    )
    print(f"Inserted ids: {insertion_result.inserted_ids}")
    acursor = acollection.find({})
    async for doc in acursor:
        print(f"  * Document: {doc}")

    acursor2 = acollection.find({}, limit=2)
    the_docs = await acursor2.to_list()
    print(f"Some docs: {the_docs}")

    await acollection.delete_many({})

if __name__ == "__main__":
    asyncio.run(amain())

# ***************#
# No type hints: #
# ***************#

# Create and get a Table:
my_table = database.create_table(
    "friends_t",
    definition=table_definition,
)
the_same_table = database.get_table("friends_t")

# Create and get a Collection:
my_collection = database.create_collection("friends_c")
the_same_collection = database.get_collection("friends_c")

# Run a find and get a cursor:
projected_cursor_c = my_collection.find(
    {"city_id": 6},
    projection={"name": True, "nickname": True, "_id": False},
)


# *****************#
# With type hints: #
# *****************#

# Create and get a Table:
my_table: Table[MyFullTableDict] = database.create_table(
    "friends_t",
    definition=table_definition,
    row_type=MyFullTableDict,
)
the_same_table: Table[MyFullTableDict] = database.get_table(
    "friends_t",
    row_type=MyFullTableDict,
)

# Create and get a Collection:
my_collection: Collection[MyFullCollectionDict] = database.create_collection(
    "friends_c",
    document_type=MyFullCollectionDict,
)
the_same_collection: Collection[MyFullCollectionDict] = database.get_collection(
    "friends_c",
    document_type=MyFullCollectionDict,
)

# Run a find and get a cursor:
projected_cursor_c = my_collection.find(
    {"city_id": 6},
    projection={"name": True, "nickname": True, "_id": False},
    document_type=MyNameNickDict,
)

# *****************#
# With type hints: #
# *****************#

from typing import TypedDict

# Preliminaries
from astrapy import DataAPIClient
client = DataAPIClient("TOKEN")
database = client.get_database("API_ENDPOINT")

# TypedDict objects for the involved rows/document
class NickDict(TypedDict):
    nickname: str

class MyNameNickDict(TypedDict):
    name: str
    nickname: str

class MyFullTableDict(TypedDict):
    city_id: int
    name: str
    age: int

class MyFullEnrichedTableDict(TypedDict):
    city_id: int
    name: str
    age: int
    nickname: str

class MyFullCollectionDict(TypedDict, total=False):
    _id: str | None
    city_id: int
    name: str
    age: int
    nickname: str

# Additional import of Table/Collection needed for type hints
from astrapy import Collection, Table

# Functions to format and transform (for demonstration purposes)
def format_name(fname: str, fnick: str) -> str:
    return f"{fname} '{fnick}'"

def extract_nick(row: NickDict) -> str:
    return row["nickname"]

def format_nick(nick: str) -> str:
    return f"<{nick.upper()} !!!>"


# Interaction with TABLES
from astrapy.constants import SortMode
from astrapy.info import (
    CreateTableDefinition,
    ColumnType,
    TableScalarColumnTypeDescriptor,
    AlterTableAddColumns,
)

table_definition = (
    CreateTableDefinition.builder()
    .add_column("city_id", ColumnType.TINYINT)
    .add_column("name", ColumnType.TEXT)
    .add_column("age", ColumnType.INT)
    .add_partition_by(["city_id"])
    .add_partition_sort({"name": SortMode.ASCENDING})
    .build()
)

# Create a table
my_table: Table[MyFullTableDict] = database.create_table(
    "friends_t",
    definition=table_definition,
    row_type=MyFullTableDict,
)

# Use get_table to obtain an equivalent Table object
the_same_table: Table[MyFullTableDict] = database.get_table(
    "friends_t",
    row_type=MyFullTableDict,
)

# Insert rows
my_table.insert_many([
    {"city_id": 6, "name": "Paula", "age": 39},
    {"city_id": 6, "name": "Liam", "age": 25},
    {"city_id": 6, "name": "Dana", "age": 31},
])

# Get a row
paula_row = my_table.find_one({"city_id": 6, "name": "Paula"})
if paula_row is None:
   raise ValueError("Paula not found. Hmm.")

# Read a row column
# Variable paula_row is now typed as MyFullTableDict -> paula_age as `int`:
paula_age = paula_row["age"]

# Alter the table, getting a new Table object as result:
# the reason this method does return at all is to adopt the return value
# entirely, which can encode the change in the row type:
enriched_table: Table[MyFullEnrichedTableDict] = my_table.alter(
    AlterTableAddColumns(
        columns={
            "nickname": TableScalarColumnTypeDescriptor(
                column_type=ColumnType.TEXT,
            ),
        }
    ),
    row_type=MyFullEnrichedTableDict,
)

# Update rows with the newly-added column
enriched_table.update_one(
    {"city_id": 6, "name": "Paula"}, {"$set": {"nickname": "The Wise"}}
)
enriched_table.update_one(
    {"city_id": 6, "name": "Liam"}, {"$set": {"nickname": "Franz"}}
)
enriched_table.update_one(
    {"city_id": 6, "name": "Dana"}, {"$set": {"nickname": "Dee"}}
)

# Scroll through the rows from a find with a certain projection, and use them
# This `find` returns a `TableFindCursor[MyFullEnrichedTableDict, MyNameNickDict]`:
projected_cursor_t = enriched_table.find(
    {"city_id": 6},
    projection={"name": True, "nickname": True},
    row_type=MyNameNickDict,
)

print("\nWith a cursor on a table:")
for nn_row in projected_cursor_t:
    # Variable nn_row is typed as a `MyNameNickDict`
    formatted_name = format_name(nn_row["name"], nn_row["nickname"])
    print(f"*  {formatted_name}")

# Scroll through rows from a find with a mapping attached
# (1) `find` returns a `TableFindCursor[MyFullEnrichedTableDict, NickDict]`
# (2) The map function is a `Callable[[NickDict], str]`
# => The mapped cursor gets type `TableFindCursor[MyFullEnrichedTableDict, str]`
mapped_cursor_t = enriched_table.find(
    {"city_id": 6},
    projection={"nickname": True},
    row_type=NickDict,
).map(extract_nick)

print("\nWith a mapped cursor on a table:")
for nn_nickname in mapped_cursor_t:
    # ... and nn_nickname is typed as a str:
    print(f"=> {format_nick(nn_nickname)}")

# Interaction with COLLECTIONS

# Create a collection
my_collection: Collection[MyFullCollectionDict] = database.create_collection(
    "friends_c",
    document_type=MyFullCollectionDict,
)

# Use get_collection to get an equivalent Collection object
the_same_collection: Collection[MyFullCollectionDict] = database.get_collection(
    "friends_c",
    document_type=MyFullCollectionDict,
)

# Insert documents (with temporary nickname for now)
my_collection.insert_many([
    {"city_id": 6, "name": "Paula", "age": 39, "nickname": "n/a"},
    {"city_id": 6, "name": "Liam", "age": 25, "nickname": "n/a"},
    {"city_id": 6, "name": "Dana", "age": 31, "nickname": "n/a"},
])

# Get a document
paula_document = my_collection.find_one({"city_id": 6, "name": "Paula"})
if paula_document is None:
   raise ValueError("Paula not found. Hmm.")

# Read a field from the document
# Var. paula_document inferred a MyFullCollectionDict -> paula_age as `int`:
paula_age = paula_document["age"]

# Update documents with a new field
my_collection.update_one(
    {"city_id": 6, "name": "Paula"}, {"$set": {"nickname": "The Wise"}}
)
my_collection.update_one(
    {"city_id": 6, "name": "Liam"}, {"$set": {"nickname": "Franz"}}
)
my_collection.update_one(
    {"city_id": 6, "name": "Dana"}, {"$set": {"nickname": "Dee"}}
)

# Iterate over documents from a find with a certain projection and use them
# `find` returns a `CollectionFindCursor[MyFullCollectionDict, MyNameNickDict]`:
projected_cursor_c = my_collection.find(
    {"city_id": 6},
    projection={"name": True, "nickname": True, "_id": False},
    document_type=MyNameNickDict,
)

print("\nWith a cursor on a collection:")
for nn_doc in projected_cursor_c:
    # Variable nn_doc is typed as a `MyNameNickDict`
    formatted_name = format_name(nn_doc["name"], nn_doc["nickname"])
    print(f"*  {formatted_name}")

# Iterate over documents from a find with a mapping attached
# (1) `find` return value is a `CollectionFindCursor[MyFullCollectionDict, NickDict]`
# (2) The map function is a `Callable[[NickDict], str]`
# => The mapped cursor gets the type `CollectionFindCursor[MyFullCollectionDict, str]`
mapped_cursor_c = my_collection.find(
    {"city_id": 6},
    projection={"nickname": True},
    document_type=NickDict,
).map(extract_nick)

print("\nWith a mapped cursor on a collection:")
for nn_nickname in mapped_cursor_c:
    # ... and nn_nickname is typed as a str:
    print(f"=> {format_nick(nn_nickname)}")