Skip to content

Commit

Permalink
add dgraph tool and source (#1)
Browse files Browse the repository at this point in the history
* add dgraph tool and source

Co-authored-by: Ryan Fox-Tyler <60440289+ryanfoxtyler@users.noreply.github.com>
  • Loading branch information
shivaji-kharse and ryanfoxtyler authored Jan 23, 2025
1 parent 1de3853 commit 54844f4
Show file tree
Hide file tree
Showing 13 changed files with 1,064 additions and 3 deletions.
13 changes: 12 additions & 1 deletion docs/en/resources/sources/_index.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,4 +26,15 @@ sources:
In implementation, each source is a different connection pool or client that used
to connect to the database and execute the tool.
## Available Sources
## Kinds of Sources
We currently support the following types of kinds of sources:
* [alloydb-postgres](./alloydb-pg.md) - Connect to an AlloyDB for PostgreSQL
cluster.
* [cloud-sql-postgres](./cloud-sql-pg.md) - Connect to a Cloud SQL for
PostgreSQL instance.
* [postgres](./postgres.md) - Connect to any PostgreSQL compatible database.
* [spanner](./spanner.md) - Connect to a Spanner database.
* [neo4j](./neo4j.md) - Connect to a Neo4j instance.
* [dgraph](./dgraph.md) - Connect to a Dgraph instance.
37 changes: 37 additions & 0 deletions docs/en/resources/sources/dgraph.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
# Dgraph Source

[Dgraph](https://dgraph.io/docs) is a horizontally scalable and distributed graph database.
It provides ACID transactions, consistent replication, and linearizable reads.

This source can connect to either a self-managed Dgraph cluster or one hosted on Dgraph Cloud.
If you're new to Dgraph, the fastest way to get started is to [sign up for Dgraph Cloud](https://cloud.dgraph.io/login).

## Requirements

### Database User

When **connecting to a hosted Dgraph database**, this source uses the API key for access. If you are using a dedicated environment, you will additionally need the namespace and user credentials for that namespace.

For **connecting to a local or self-hosted Dgraph database**, use the namespace and user credentials for that namespace.

```yaml
sources:
my-dgraph-source:
kind: "dgraph"
dgraphUrl: "https://xxxx.cloud.dgraph.io"
user: "groot"
password: "password"
apiKey: abc123
namepace : 0
```
## Reference
| **Field** | **Type** | **Required** | **Description** |
|-------------|:--------:|:------------:|--------------------------------------------------------------------------------------------------|
| kind | string | true | Must be "dgraph". |
| dgraphurl | string | true | Connection URI (e.g. "https://xxx.cloud.dgraph.io", "https://localhost:8080"). |
| user | string | false | Name of the Dgraph user to connect as (e.g., "groot"). |
| password | string | false | Password of the Dgraph user (e.g., "password"). |
| apiKey | string | false | API key to connect to a Dgraph Cloud instance. |
| namespace | uint64 | false | Dgraph namespace (not required for Dgraph Cloud Shared Clusters). |
2 changes: 2 additions & 0 deletions docs/tools/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,8 @@ We currently support the following types of kinds of tools:
statement againts Spanner database.
* [neo4j-cypher](./neo4j-cypher.md) - Run a Cypher statement against a
Neo4j database.
* [dgraph-dql](./dgraph-dql.md) - Run a DQL statement against a
Dgraph database.
## Specifying Parameters
Expand Down
103 changes: 103 additions & 0 deletions docs/tools/dgraph-dql.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
# dgraph DQL Tool


A "dgraph-dql" tool executes a pre-defined DQL statement against a Dgraph database. It's compatible with any of the following
sources:
- [dgraph](../sources/dgraph.md)

To run a statement as a query, you need to set the config isQuery=true. For upserts or mutations, set isQuery=false.
You can also configure timeout for a query.
## Example
Query:

```yaml
tools:
search_user:
kind: dgraph-dql
source: dgraph-user-instance
statement: |
query all($role: string){
users(func: has(name)) @filter(eq(role, $role) AND ge(age, 30) AND le(age, 50)) {
uid
name
email
role
age
}
}
isQuery: true
timeout: 20s
description: |
Use this tool to retrieve the details of users who are admins and are between 30 and 50 years old.
The query returns the user's name, email, role, and age.
This can be helpful when you want to fetch admin users within a specific age range.
Example: Fetch admins aged between 30 and 50:
[
{
"name": "Alice",
"role": "admin",
"age": 35
},
{
"name": "Bob",
"role": "admin",
"age": 45
}
]
parameters:
- name: $role
type: string
description: admin
```
Mutation:
```yaml
tools:
dgraph-manage-user-instance:
kind: dgraph-dql
source: dgraph-manage-user-instance
isQuery: false
statement: |
{
set {
_:user1 <name> $user1 .
_:user1 <email> $email1 .
_:user1 <role> "admin" .
_:user1 <age> "35" .
_:user2 <name> $user2 .
_:user2 <email> $email2 .
_:user2 <role> "admin" .
_:user2 <age> "45" .
}
}
description: |
Use this tool to insert or update user data into the Dgraph database.
The mutation adds or updates user details like name, email, role, and age.
Example: Add users Alice and Bob as admins with specific ages.
parameters:
- name: $user1
type: string
description: Alice
- name: $email1
type: string
description: alice@email.com
- name: $user2
type: string
description: Bob
- name: $email2
type: string
description: bob@email.com
```
## Reference
| **field** | **type** | **required** | **description** |
|-------------|----------:|:------------:|----------------------------------------------------------------------------------------------------|
| kind | string | true | Must be "dgraph-dql". |
| source | string | true | Name of the source the dql query should execute on. |
| description | string | true | Description of the tool |
| statement | string | true | dql statement to execute |
| isQuery | boolean | false | To run statment as query set true otherwise false |
| timeout | string | false | To set timout for query |
| parameters | parameter | true | List of [parameters](README.md#specifying-parameters) that will be used with the dql statement. |
2 changes: 2 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ require (
github.com/census-instrumentation/opencensus-proto v0.4.1 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/cncf/xds/go v0.0.0-20240905190251-b4127c9b8d78 // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
github.com/envoyproxy/go-control-plane v0.13.1 // indirect
github.com/envoyproxy/protoc-gen-validate v1.1.0 // indirect
github.com/felixge/httpsnoop v1.0.4 // indirect
Expand All @@ -66,6 +67,7 @@ require (
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
github.com/jackc/puddle/v2 v2.2.2 // indirect
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
github.com/spf13/pflag v1.0.5 // indirect
go.opencensus.io v0.24.0 // indirect
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
Expand Down
6 changes: 4 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -687,8 +687,9 @@ github.com/cncf/xds/go v0.0.0-20240905190251-b4127c9b8d78/go.mod h1:W+zGtBO5Y1Ig
github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE=
github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
Expand Down Expand Up @@ -933,8 +934,9 @@ github.com/pkg/sftp v1.10.1/go.mod h1:lYOWFsE0bwd1+KfKJaKeuokY15vzFx25BLbzYYoAxZ
github.com/pkg/sftp v1.13.1/go.mod h1:3HaPG6Dq1ILlpPZRO0HVMrsydcdLt6HRDccSgb87qRg=
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo=
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
github.com/prometheus/client_model v0.2.0/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
github.com/prometheus/client_model v0.3.0/go.mod h1:LDGWKZIo7rky3hgvBe+caln+Dr3dPggB5dvjtD7w9+w=
Expand Down
14 changes: 14 additions & 0 deletions internal/server/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,12 @@ import (
alloydbpgsrc "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
cloudsqlmysqlsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql"
cloudsqlpgsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
dgraphsrc "github.com/googleapis/genai-toolbox/internal/sources/dgraph"
neo4jrc "github.com/googleapis/genai-toolbox/internal/sources/neo4j"
postgressrc "github.com/googleapis/genai-toolbox/internal/sources/postgres"
spannersrc "github.com/googleapis/genai-toolbox/internal/sources/spanner"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/dgraph"
"github.com/googleapis/genai-toolbox/internal/tools/mysql"
neo4jtool "github.com/googleapis/genai-toolbox/internal/tools/neo4j"
"github.com/googleapis/genai-toolbox/internal/tools/postgressql"
Expand Down Expand Up @@ -173,6 +175,12 @@ func (c *SourceConfigs) UnmarshalYAML(unmarshal func(interface{}) error) error {
return fmt.Errorf("unable to parse as %q: %w", k.Kind, err)
}
(*c)[name] = actual
case dgraphsrc.SourceKind:
actual := dgraphsrc.Config{Name: name}
if err := u.Unmarshal(&actual); err != nil {
return fmt.Errorf("unable to parse as %q: %w", k.Kind, err)
}
(*c)[name] = actual
default:
return fmt.Errorf("%q is not a valid kind of data source", k.Kind)
}
Expand Down Expand Up @@ -264,6 +272,12 @@ func (c *ToolConfigs) UnmarshalYAML(unmarshal func(interface{}) error) error {
return fmt.Errorf("unable to parse as %q: %w", k.Kind, err)
}
(*c)[name] = actual
case dgraph.ToolKind:
actual := dgraph.Config{Name: name}
if err := u.Unmarshal(&actual); err != nil {
return fmt.Errorf("unable to parse as %q: %w", k.Kind, err)
}
(*c)[name] = actual
default:
return fmt.Errorf("%q is not a valid kind of tool", k.Kind)
}
Expand Down
Loading

0 comments on commit 54844f4

Please sign in to comment.