Merge pull request #54 from datastax/astra-vector

Grab all Python CLI developments and bring to master
datastax · Oct 27, 2023 · c51d9d0 · c51d9d0
2 parents b883c5e + ececbfb
commit c51d9d0
Show file tree

Hide file tree

Showing 28 changed files with 1,906 additions and 1,686 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -0,0 +1,36 @@
+name: Run Pytest
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+    branches:
+      - master
+
+jobs:
+  test:
+    env:
+      ASTRA_DB_ID: ${{ secrets.ASTRA_DB_ID }}
+      ASTRA_DB_REGION: ${{ secrets.ASTRA_DB_REGION }}
+      ASTRA_DB_APPLICATION_TOKEN: ${{ secrets.ASTRA_DB_APPLICATION_TOKEN }}
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.11
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install pytest
+        pip install -r requirements.txt
+
+    - name: Run pytest
+      run: |
+        PYTHONPATH=. pytest -s
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -0,0 +1,7 @@
+repos:
+-   repo: https://github.com/psf/black
+    rev: main  # Replace with your desired version or use 'stable'
+    hooks:
+    - id: black
+      args: ['--quiet']
+
diff --git a/README.md b/README.md
@@ -1,144 +1,127 @@
 ## AstraPy
 
-[![Actions Status](https://github.com/datastax/astrapy/workflows/Tests/badge.svg)](https://github.com/datastax/astrapy/actions) 
+[![Actions Status](https://github.com/datastax/astrapy/workflows/Tests/badge.svg)](https://github.com/datastax/astrapy/actions)
 
 AstraPy is a Pythonic SDK for [DataStax Astra](https://astra.datastax.com) and [Stargate](https://stargate.io/)
 
 ### Resources
+
 - [DataStax Astra](https://astra.datastax.com)
 - [Stargate](https://stargate.io/)
 
-
 ### Getting Started
+
 Install AstraPy
+
 ```shell
 pip install astrapy
 ```
 
 Setup your Astra client
-```python
-from astrapy.client import create_astra_client
-
-astra_client = create_astra_client(astra_database_id=ASTRA_DB_ID,
-                                   astra_database_region=ASTRA_DB_REGION,
-                                   astra_application_token=ASTRA_DB_APPLICATION_TOKEN)
-```
 
-Take a look at the [client tests](https://github.com/datastax/astrapy/blob/master/tests/astrapy/test_client.py) and the [collection tests](https://github.com/datastax/astrapy/blob/master/tests/astrapy/test_collections.py) for specific endpoint examples.
+Create a .env file with the appropriate values, or use the 'astra' cli to do the same.
 
-#### Using the Ops Client
-You can use the Ops client to work the with Astra DevOps API. [API Reference](https://docs.datastax.com/en/astra/docs/_attachments/devopsv2.html)
-```python
-# astra_client created above
-# create a keyspace using the Ops API
-astra_client.ops.create_keyspace(database=ASTRA_DB_ID, keyspace=KEYSPACE_NAME)
+```bash
+ASTRA_DB_KEYSPACE="<keyspace>"
+ASTRA_DB_APPLICATION_TOKEN="<AstraCS:...>"
+ASTRA_DB_REGION="<region>"
+ASTRA_DB_ID=<db_id>
 ```
 
-#### Using the REST Client
-You can use the REST client to work with the Astra REST API. [API Reference](https://docs.datastax.com/en/astra/docs/_attachments/restv2.html#tag/Data)
-```python
-# astra_client created above
-# search a table
-res = astra_client.rest.search_table(keyspace=ASTRA_DB_KEYSPACE,
-                                     table=TABLE_NAME,
-                                     query={"firstname": {"$eq": "Cliff"}})
-print(res["count"]) # number of results
-print(res["data"]) # list of rows
-```
+Load the variables in and then create the client. This collections client can make non-vector and vector calls, depending on the call configuration.
 
-#### Using the Schemas Client
-You can use the Schemas client to work with the Astra Schemas API. [API Reference](https://docs.datastax.com/en/astra/docs/_attachments/restv2.html#tag/Schemas)
 ```python
-# astra_client created above
-# create a table
-astra_client.schemas.create_table(keyspace=ASTRA_DB_KEYSPACE, table_definition={
-    "name": "my_table",
-    "columnDefinitions": [
-        {
-            "name": "firstname",
-            "typeDefinition": "text"
-        },
-        {
-            "name": "lastname",
-            "typeDefinition": "text"
-        },
-        {
-            "name": "favorite_color",
-            "typeDefinition": "text",
-        }
-    ],
-    "primaryKey": {
-        "partitionKey": [
-            "firstname"
-        ],
-        "clusteringKey": [
-            "lastname"
-        ]
+import os
+import sys
+
+from astrapy.db import AstraDB, AstraDBCollection
+from astrapy.ops import AstraDBOps
+
+# First, we work with devops
+api_key = os.getenv("ASTRA_DB_APPLICATION_TOKEN")
+astra_ops = AstraDBOps(api_key)
+
+# Define a database to create
+database_definition = {
+    "name": "vector_test",
+    "tier": "serverless",
+    "cloudProvider": "GCP",
+    "keyspace": os.getenv("ASTRA_DB_KEYSPACE", "default_keyspace"),
+    "region": os.getenv("ASTRA_DB_REGION", None),
+    "capacityUnits": 1,
+    "user": "example",
+    "password": api_key,
+    "dbType": "vector",
+}
+
+# Create the database
+create_result = astra_ops.create_database(database_definition=database_definition)
+
+# Grab the new information from the database
+database_id = create_result["id"]
+database_region = astra_ops.get_database()[0]["info"]["region"]
+database_base_url = "apps.astra.datastax.com"
+
+# Build the endpoint URL:
+api_endpoint = f"https://{database_id}-{database_region}.{database_base_url}"
+
+# Initialize our vector db
+astra_db = AstraDB(api_key=api_key, api_endpoint=api_endpoint)
+
+# Possible Operations
+astra_db.create_collection(collection_name="collection_test_delete", size=5)
+astra_db.delete_collection(collection_name="collection_test_delete")
+astra_db.create_collection(collection_name="collection_test", size=5)
+
+# Collections
+astra_db_collection = AstraDBCollection(
+    collection_name="collection_test", astra_db=astra_db
+)
+# Or...
+astra_db_collection = AstraDBCollection(
+    collection_name="collection_test", api_key=api_key, api_endpoint=api_endpoint
+)
+
+astra_db_collection.insert_one(
+    {
+        "_id": "5",
+        "name": "Coded Cleats Copy",
+        "description": "ChatGPT integrated sneakers that talk to you",
+        "$vector": [0.25, 0.25, 0.25, 0.25, 0.25],
     }
-})
+)
+
+astra_db_collection.find_one({"name": "potato"})
+astra_db_collection.find_one({"name": "Coded Cleats Copy"})
 ```
 
+#### More Information
 
-#### Using the Collections Client
-You can use the Collections client to work with the Astra Document API. [API Reference](https://docs.datastax.com/en/astra/docs/_attachments/docv2.html)
-```python
-# astra_client created above
-# create multiple documents using the collections API
-my_collection = astra_client.collections.namespace(ASTRA_DB_KEYSPACE).collection(COLLECTION_NAME)
-my_collection.batch(documents=[
-    {
-        "documentId": "1",
-        "first_name": "Dang",
-        "last_name": "Son",
-    }, {
-        "documentId": "2",
-        "first_name": "Yep",
-        "last_name": "Boss",
-    }])
-```
+Check out the [notebook](https://colab.research.google.com/github/synedra/astra_vector_examples/blob/main/notebook/vector.ipynb#scrollTo=f04a1806) which has examples for finding and inserting information into the database, including vector commands.
 
-#### Using the GraphQL Client
-You can use the GraphQL client to work with the Astra GraphQL API. [API Reference](https://docs.datastax.com/en/astra/docs/using-the-astra-graphql-api.html)
-```python
-# astra_client created above
-# create multiple documents using the GraphQL API
-astra_client.gql.execute(keyspace=ASTRA_DB_KEYSPACE, query="""
-        mutation insert2Books {
-            moby: insertbook(value: {title:"Moby Dick", author:"Herman Melville"}) {
-                value {
-                    title
-                }
-            }
-            catch22: insertbook(value: {title:"Catch-22", author:"Joseph Heller"}) {
-                value {
-                    title
-                }
-            }
-        }
-    """)
-```
+Take a look at the [vector tests](https://github.com/datastax/astrapy/blob/master/tests/astrapy/test_collections.py) and the [collection tests](https://github.com/datastax/astrapy/blob/master/tests/astrapy/test_collections.py) for specific endpoint examples.
+
+#### Using the Ops Client
+
+You can use the Ops client to work with the Astra DevOps API. Check the [devops tests](https://github.com/datastax/astrapy/blob/master/tests/astrapy/test_devops.py)
+
+### For Developers
+
+#### Testing
+
+Ensure you provide all required environment variables:
 
-#### Using the HTTP Client
-You can use the HTTP client to work with any Astra/Stargate endpoint directly. [API Reference](https://docs.datastax.com/en/astra/docs/api.html)
-```python
-# astra_client created above
-# create a document on Astra using the Document API
-astra_client._rest_client.request(
-    method="PUT",
-    path=f"/api/rest/v2/namespaces/my_namespace/collections/my_collection/user_1",
-    json_data={
-        "first_name": "Cliff",
-        "last_name": "Wicklow",
-        "emails": ["[email protected]"],
-    })
+```
+export ASTRA_DB_ID="..."
+export ASTRA_DB_REGION="..."
+export ASTRA_DB_APPLICATION_TOKEN="..."
+export ASTRA_DB_KEYSPACE="..."
+export ASTRA_CLIENT_ID="..."
+export ASTRA_CLIENT_SECRET="..."
 ```
 
-#### Connecting to a local Stargate Instance
-```python
-from astrapy.client import create_astra_client
+then you can run:
 
-stargate_client = create_astra_client(base_url=http://localhost:8082,
-                                      auth_base_url=http://localhost:8081/v1/auth,
-                                      username=****,
-                                      password=****)
+```
+PYTHONPATH=. pytest
 ```
diff --git a/astrapy/__init__.py b/astrapy/__init__.py
@@ -10,4 +10,4 @@
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
-# limitations under the License.
+# limitations under the License.