Compare commits

..

7 Commits

531 changed files with 2230 additions and 2807 deletions

View File

@@ -59,13 +59,6 @@ You can manually trigger the bot by commenting on your Pull Request:
* `/gemini summary`: Posts a summary of the changes in the pull request.
* `/gemini help`: Overview of the available commands
## Guidelines for Pull Requests
1. Please keep your PR small for more thorough review and easier updates. In case of regression, it also allows us to roll back a single feature instead of multiple ones.
1. For non-trivial changes, consider opening an issue and discussing it with the code owners first.
1. Provide a good PR description as a record of what change is being made and why it was made. Link to a GitHub issue if it exists.
1. Make sure your code is thoroughly tested with unit tests and integration tests. Remember to clean up the test instances properly in your code to avoid memory leaks.
## Adding a New Database Source or Tool
Please create an
@@ -92,11 +85,11 @@ implementation](https://github.com/googleapis/genai-toolbox/blob/main/internal/s
`newdb.go`. Create a `Config` struct to include all the necessary parameters
for connecting to the database (e.g., host, port, username, password, database
name) and a `Source` struct to store necessary parameters for tools (e.g.,
Name, Type, connection object, additional config).
Name, Kind, connection object, additional config).
* **Implement the
[`SourceConfig`](https://github.com/googleapis/genai-toolbox/blob/fd300dc606d88bf9f7bba689e2cee4e3565537dd/internal/sources/sources.go#L57)
interface**. This interface requires two methods:
* `SourceConfigType() string`: Returns a unique string identifier for your
* `SourceConfigKind() string`: Returns a unique string identifier for your
data source (e.g., `"newdb"`).
* `Initialize(ctx context.Context, tracer trace.Tracer) (Source, error)`:
Creates a new instance of your data source and establishes a connection to
@@ -104,7 +97,7 @@ implementation](https://github.com/googleapis/genai-toolbox/blob/main/internal/s
* **Implement the
[`Source`](https://github.com/googleapis/genai-toolbox/blob/fd300dc606d88bf9f7bba689e2cee4e3565537dd/internal/sources/sources.go#L63)
interface**. This interface requires one method:
* `SourceType() string`: Returns the same string identifier as `SourceConfigType()`.
* `SourceKind() string`: Returns the same string identifier as `SourceConfigKind()`.
* **Implement `init()`** to register the new Source.
* **Implement Unit Tests** in a file named `newdb_test.go`.
@@ -117,8 +110,6 @@ implementation](https://github.com/googleapis/genai-toolbox/blob/main/internal/s
We recommend looking at an [example tool
implementation](https://github.com/googleapis/genai-toolbox/tree/main/internal/tools/postgres/postgressql).
Remember to keep your PRs small. For example, if you are contributing a new Source, only include one or two core Tools within the same PR, the rest of the Tools can come in subsequent PRs.
* **Create a new directory** under `internal/tools` for your tool type (e.g., `internal/tools/newdb/newdbtool`).
* **Define a configuration struct** for your tool in a file named `newdbtool.go`.
Create a `Config` struct and a `Tool` struct to store necessary parameters for
@@ -126,7 +117,7 @@ tools.
* **Implement the
[`ToolConfig`](https://github.com/googleapis/genai-toolbox/blob/fd300dc606d88bf9f7bba689e2cee4e3565537dd/internal/tools/tools.go#L61)
interface**. This interface requires one method:
* `ToolConfigType() string`: Returns a unique string identifier for your tool
* `ToolConfigKind() string`: Returns a unique string identifier for your tool
(e.g., `"newdb-tool"`).
* `Initialize(sources map[string]Source) (Tool, error)`: Creates a new
instance of your tool and validates that it can connect to the specified
@@ -172,8 +163,6 @@ tools.
parameters][temp-param-doc]. Only run this test if template
parameters apply to your tool.
* **Add additional tests** for the tools that are not covered by the predefined tests. Every tool must be tested!
* **Add the new database to the integration test workflow** in
[integration.cloudbuild.yaml](.ci/integration.cloudbuild.yaml).
@@ -243,7 +232,7 @@ resources.
| style | Update src code, with only formatting and whitespace updates (e.g. code formatter or linter changes). |
Pull requests should always add scope whenever possible. The scope is
formatted as `<scope-resource>/<scope-type>` (e.g., `sources/postgres`, or
formatted as `<scope-type>/<scope-kind>` (e.g., `sources/postgres`, or
`tools/mssql-sql`).
Ideally, **each PR covers only one scope**, if this is
@@ -255,4 +244,4 @@ resources.
* **PR Description:** PR description should **always** be included. It should
include a concise description of the changes, it's impact, along with a
summary of the solution. If the PR is related to a specific issue, the issue
number should be mentioned in the PR description (e.g. `Fixes #1`).
number should be mentioned in the PR description (e.g. `Fixes #1`).

View File

@@ -954,7 +954,7 @@ For more details on configuring different types of sources, see the
### Tools
The `tools` section of a `tools.yaml` define the actions an agent can take: what
type of tool it is, which source(s) it affects, what parameters it uses, etc.
kind of tool it is, which source(s) it affects, what parameters it uses, etc.
```yaml
tools:

View File

@@ -92,7 +92,6 @@ import (
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaresearchdicomstudies"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudmonitoring"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlcloneinstance"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlcreatebackup"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlcreatedatabase"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlcreateusers"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlgetinstances"

View File

@@ -535,7 +535,7 @@ func TestParseToolFile(t *testing.T) {
Sources: server.SourceConfigs{
"my-pg-instance": cloudsqlpgsrc.Config{
Name: "my-pg-instance",
Type: cloudsqlpgsrc.SourceType,
Kind: cloudsqlpgsrc.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -548,7 +548,7 @@ func TestParseToolFile(t *testing.T) {
Tools: server.ToolConfigs{
"example_tool": postgressql.Config{
Name: "example_tool",
Type: "postgres-sql",
Kind: "postgres-sql",
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
@@ -689,7 +689,7 @@ func TestParseToolFileWithAuth(t *testing.T) {
Sources: server.SourceConfigs{
"my-pg-instance": cloudsqlpgsrc.Config{
Name: "my-pg-instance",
Type: cloudsqlpgsrc.SourceType,
Kind: cloudsqlpgsrc.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -702,19 +702,19 @@ func TestParseToolFileWithAuth(t *testing.T) {
AuthServices: server.AuthServiceConfigs{
"my-google-service": google.Config{
Name: "my-google-service",
Type: google.AuthServiceType,
Kind: google.AuthServiceKind,
ClientID: "my-client-id",
},
"other-google-service": google.Config{
Name: "other-google-service",
Type: google.AuthServiceType,
Kind: google.AuthServiceKind,
ClientID: "other-client-id",
},
},
Tools: server.ToolConfigs{
"example_tool": postgressql.Config{
Name: "example_tool",
Type: "postgres-sql",
Kind: "postgres-sql",
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
@@ -789,7 +789,7 @@ func TestParseToolFileWithAuth(t *testing.T) {
Sources: server.SourceConfigs{
"my-pg-instance": cloudsqlpgsrc.Config{
Name: "my-pg-instance",
Type: cloudsqlpgsrc.SourceType,
Kind: cloudsqlpgsrc.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -802,19 +802,19 @@ func TestParseToolFileWithAuth(t *testing.T) {
AuthSources: server.AuthServiceConfigs{
"my-google-service": google.Config{
Name: "my-google-service",
Type: google.AuthServiceType,
Kind: google.AuthServiceKind,
ClientID: "my-client-id",
},
"other-google-service": google.Config{
Name: "other-google-service",
Type: google.AuthServiceType,
Kind: google.AuthServiceKind,
ClientID: "other-client-id",
},
},
Tools: server.ToolConfigs{
"example_tool": postgressql.Config{
Name: "example_tool",
Type: "postgres-sql",
Kind: "postgres-sql",
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
@@ -891,7 +891,7 @@ func TestParseToolFileWithAuth(t *testing.T) {
Sources: server.SourceConfigs{
"my-pg-instance": cloudsqlpgsrc.Config{
Name: "my-pg-instance",
Type: cloudsqlpgsrc.SourceType,
Kind: cloudsqlpgsrc.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -904,19 +904,19 @@ func TestParseToolFileWithAuth(t *testing.T) {
AuthServices: server.AuthServiceConfigs{
"my-google-service": google.Config{
Name: "my-google-service",
Type: google.AuthServiceType,
Kind: google.AuthServiceKind,
ClientID: "my-client-id",
},
"other-google-service": google.Config{
Name: "other-google-service",
Type: google.AuthServiceType,
Kind: google.AuthServiceKind,
ClientID: "other-client-id",
},
},
Tools: server.ToolConfigs{
"example_tool": postgressql.Config{
Name: "example_tool",
Type: "postgres-sql",
Kind: "postgres-sql",
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
@@ -1062,7 +1062,7 @@ func TestEnvVarReplacement(t *testing.T) {
Sources: server.SourceConfigs{
"my-http-instance": httpsrc.Config{
Name: "my-http-instance",
Type: httpsrc.SourceType,
Kind: httpsrc.SourceKind,
BaseURL: "http://test_server/",
Timeout: "10s",
DefaultHeaders: map[string]string{"Authorization": "ACTUAL_HEADER"},
@@ -1072,19 +1072,19 @@ func TestEnvVarReplacement(t *testing.T) {
AuthServices: server.AuthServiceConfigs{
"my-google-service": google.Config{
Name: "my-google-service",
Type: google.AuthServiceType,
Kind: google.AuthServiceKind,
ClientID: "ACTUAL_CLIENT_ID",
},
"other-google-service": google.Config{
Name: "other-google-service",
Type: google.AuthServiceType,
Kind: google.AuthServiceKind,
ClientID: "ACTUAL_CLIENT_ID_2",
},
},
Tools: server.ToolConfigs{
"example_tool": http.Config{
Name: "example_tool",
Type: "http",
Kind: "http",
Source: "my-instance",
Method: "GET",
Path: "search?name=alice&pet=cat",
@@ -1493,7 +1493,7 @@ func TestPrebuiltTools(t *testing.T) {
wantToolset: server.ToolsetConfigs{
"cloud_sql_postgres_admin_tools": tools.ToolsetConfig{
Name: "cloud_sql_postgres_admin_tools",
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "postgres_upgrade_precheck", "clone_instance", "create_backup"},
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "postgres_upgrade_precheck", "clone_instance"},
},
},
},
@@ -1503,7 +1503,7 @@ func TestPrebuiltTools(t *testing.T) {
wantToolset: server.ToolsetConfigs{
"cloud_sql_mysql_admin_tools": tools.ToolsetConfig{
Name: "cloud_sql_mysql_admin_tools",
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup"},
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance"},
},
},
},
@@ -1513,7 +1513,7 @@ func TestPrebuiltTools(t *testing.T) {
wantToolset: server.ToolsetConfigs{
"cloud_sql_mssql_admin_tools": tools.ToolsetConfig{
Name: "cloud_sql_mssql_admin_tools",
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup"},
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance"},
},
},
},

View File

@@ -509,7 +509,7 @@
},
"outputs": [],
"source": [
"! pip install toolbox-core --quiet\n",
"! pip install toolbox-adk --quiet\n",
"! pip install google-adk --quiet"
]
},
@@ -525,14 +525,18 @@
"from google.adk.runners import Runner\n",
"from google.adk.sessions import InMemorySessionService\n",
"from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService\n",
"from google.adk.tools.toolbox_toolset import ToolboxToolset\n",
"from google.genai import types\n",
"from toolbox_core import ToolboxSyncClient\n",
"\n",
"import os\n",
"# TODO(developer): replace this with your Google API key\n",
"os.environ['GOOGLE_API_KEY'] = \"<GOOGLE_API_KEY>\"\n",
"\n",
"toolbox_client = ToolboxSyncClient(\"http://127.0.0.1:5000\")\n",
"# Configure toolset\n",
"toolset = ToolboxToolset(\n",
" server_url=\"http://127.0.0.1:5000\",\n",
" toolset_name=\"my-toolset\"\n",
")\n",
"\n",
"prompt = \"\"\"\n",
" You're a helpful hotel assistant. You handle hotel searching, booking and\n",
@@ -549,7 +553,7 @@
" name='hotel_agent',\n",
" description='A helpful AI assistant.',\n",
" instruction=prompt,\n",
" tools=toolbox_client.load_toolset(\"my-toolset\"),\n",
" tools=[toolset],\n",
")\n",
"\n",
"session_service = InMemorySessionService()\n",

View File

@@ -52,7 +52,7 @@ runtime](https://research.google.com/colaboratory/local-runtimes.html).
{{< tabpane persist=header >}}
{{< tab header="ADK" lang="bash" >}}
pip install toolbox-core
pip install toolbox-adk
{{< /tab >}}
{{< tab header="Langchain" lang="bash" >}}

View File

@@ -1,15 +1,17 @@
from google.adk import Agent
from google.adk.apps import App
from toolbox_core import ToolboxSyncClient
from google.adk.tools.toolbox_toolset import ToolboxToolset
# TODO(developer): update the TOOLBOX_URL to your toolbox endpoint
client = ToolboxSyncClient("http://127.0.0.1:5000")
toolset = ToolboxToolset(
server_url="http://127.0.0.1:5000",
)
root_agent = Agent(
name='root_agent',
model='gemini-2.5-flash',
instruction="You are a helpful AI assistant designed to provide accurate and useful information.",
tools=client.load_toolset(),
tools=[toolset],
)
app = App(root_agent=root_agent, name="my_agent")

View File

@@ -1,3 +1,3 @@
google-adk==1.21.0
toolbox-core==0.5.4
toolbox-adk>=0.1.0
pytest==9.0.2

View File

@@ -48,7 +48,6 @@ instance, database and users:
* `roles/cloudsql.editor`: Provides permissions to manage existing resources.
* All `viewer` tools
* `create_database`
* `create_backup`
* `roles/cloudsql.admin`: Provides full control over all resources.
* All `editor` and `viewer` tools
* `create_instance`
@@ -300,7 +299,6 @@ instances and interacting with your database:
* **create_user**: Creates a new user in a Cloud SQL instance.
* **wait_for_operation**: Waits for a Cloud SQL operation to complete.
* **clone_instance**: Creates a clone of an existing Cloud SQL for SQL Server instance.
* **create_backup**: Creates a backup on a Cloud SQL instance.
{{< notice note >}}
Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs

View File

@@ -48,7 +48,6 @@ database and users:
* `roles/cloudsql.editor`: Provides permissions to manage existing resources.
* All `viewer` tools
* `create_database`
* `create_backup`
* `roles/cloudsql.admin`: Provides full control over all resources.
* All `editor` and `viewer` tools
* `create_instance`
@@ -300,7 +299,6 @@ instances and interacting with your database:
* **create_user**: Creates a new user in a Cloud SQL instance.
* **wait_for_operation**: Waits for a Cloud SQL operation to complete.
* **clone_instance**: Creates a clone of an existing Cloud SQL for MySQL instance.
* **create_backup**: Creates a backup on a Cloud SQL instance.
{{< notice note >}}
Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs

View File

@@ -48,7 +48,6 @@ instance, database and users:
* `roles/cloudsql.editor`: Provides permissions to manage existing resources.
* All `viewer` tools
* `create_database`
* `create_backup`
* `roles/cloudsql.admin`: Provides full control over all resources.
* All `editor` and `viewer` tools
* `create_instance`
@@ -300,7 +299,6 @@ instances and interacting with your database:
* **create_user**: Creates a new user in a Cloud SQL instance.
* **wait_for_operation**: Waits for a Cloud SQL operation to complete.
* **clone_instance**: Creates a clone of an existing Cloud SQL for PostgreSQL instance.
* **create_backup**: Creates a backup on a Cloud SQL instance.
{{< notice note >}}
Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs

View File

@@ -49,7 +49,7 @@ with the necessary configuration for deployment to Vertex AI Agent Engine.
4. Add `toolbox-core` as a dependency to the new project:
```bash
uv add toolbox-core
uv add toolbox-adk
```
## Step 3: Configure Google Cloud Authentication
@@ -95,22 +95,23 @@ authentication token.
```python
from google.adk import Agent
from google.adk.apps import App
from toolbox_core import ToolboxSyncClient, auth_methods
from google.adk.tools.toolbox_toolset import ToolboxToolset
from toolbox_adk import CredentialStrategy
# TODO(developer): Replace with your Toolbox Cloud Run Service URL
TOOLBOX_URL = "https://your-toolbox-service-xyz.a.run.app"
# Initialize the client with the Cloud Run URL and Auth headers
client = ToolboxSyncClient(
TOOLBOX_URL,
client_headers={"Authorization": auth_methods.get_google_id_token(TOOLBOX_URL)}
# Initialize the toolset with Workload Identity (generates ID token for the URL)
toolset = ToolboxToolset(
server_url=TOOLBOX_URL,
credentials=CredentialStrategy.workload_identity(target_audience=TOOLBOX_URL)
)
root_agent = Agent(
name='root_agent',
model='gemini-2.5-flash',
instruction="You are a helpful AI assistant designed to provide accurate and useful information.",
tools=client.load_toolset(),
tools=[toolset],
)
app = App(root_agent=root_agent, name="my_agent")

View File

@@ -187,7 +187,6 @@ See [Usage Examples](../reference/cli.md#examples).
manage existing resources.
* All `viewer` tools
* `create_database`
* `create_backup`
* **Cloud SQL Admin** (`roles/cloudsql.admin`): Provides full control over
all resources.
* All `editor` and `viewer` tools
@@ -204,7 +203,6 @@ See [Usage Examples](../reference/cli.md#examples).
* `create_user`: Creates a new user in a Cloud SQL instance.
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
* `clone_instance`: Creates a clone for an existing Cloud SQL for MySQL instance.
* `create_backup`: Creates a backup on a Cloud SQL instance.
## Cloud SQL for PostgreSQL
@@ -277,7 +275,6 @@ See [Usage Examples](../reference/cli.md#examples).
manage existing resources.
* All `viewer` tools
* `create_database`
* `create_backup`
* **Cloud SQL Admin** (`roles/cloudsql.admin`): Provides full control over
all resources.
* All `editor` and `viewer` tools
@@ -293,7 +290,6 @@ See [Usage Examples](../reference/cli.md#examples).
* `create_user`: Creates a new user in a Cloud SQL instance.
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
* `clone_instance`: Creates a clone for an existing Cloud SQL for PostgreSQL instance.
* `create_backup`: Creates a backup on a Cloud SQL instance.
## Cloud SQL for SQL Server
@@ -340,7 +336,6 @@ See [Usage Examples](../reference/cli.md#examples).
manage existing resources.
* All `viewer` tools
* `create_database`
* `create_backup`
* **Cloud SQL Admin** (`roles/cloudsql.admin`): Provides full control over
all resources.
* All `editor` and `viewer` tools
@@ -356,7 +351,6 @@ See [Usage Examples](../reference/cli.md#examples).
* `create_user`: Creates a new user in a Cloud SQL instance.
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
* `clone_instance`: Creates a clone for an existing Cloud SQL for SQL Server instance.
* `create_backup`: Creates a backup on a Cloud SQL instance.
## Dataplex

View File

@@ -1,45 +0,0 @@
---
title: cloud-sql-create-backup
type: docs
weight: 10
description: "Creates a backup on a Cloud SQL instance."
---
The `cloud-sql-create-backup` tool creates an on-demand backup on a Cloud SQL instance using the Cloud SQL Admin API.
{{< notice info dd>}}
This tool uses a `source` of kind `cloud-sql-admin`.
{{< /notice >}}
## Examples
Basic backup creation (current state)
```yaml
tools:
backup-creation-basic:
kind: cloud-sql-create-backup
source: cloud-sql-admin-source
description: "Creates a backup on the given Cloud SQL instance."
```
## Reference
### Tool Configuration
| **field** | **type** | **required** | **description** |
| -------------- | :------: | :----------: | ------------------------------------------------------------- |
| kind | string | true | Must be "cloud-sql-create-backup". |
| source | string | true | The name of the `cloud-sql-admin` source to use. |
| description | string | false | A description of the tool. |
### Tool Inputs
| **parameter** | **type** | **required** | **description** |
| -------------------------- | :------: | :----------: | ------------------------------------------------------------------------------- |
| project | string | true | The project ID. |
| instance | string | true | The name of the instance to take a backup on. Does not include the project ID. |
| location | string | false | (Optional) Location of the backup run. |
| backup_description | string | false | (Optional) The description of this backup run. |
## See Also
- [Cloud SQL Admin API documentation](https://cloud.google.com/sql/docs/mysql/admin-api)
- [Toolbox Cloud SQL tools documentation](../cloudsql)
- [Cloud SQL Backup API documentation](https://cloud.google.com/sql/docs/mysql/backup-recovery/backups)

View File

@@ -365,7 +365,7 @@ pip install llama-index-llms-google-genai
{{< /tab >}}
{{< tab header="ADK" lang="bash" >}}
pip install toolbox-core
pip install toolbox-adk
{{< /tab >}}
{{< /tabpane >}}
@@ -607,8 +607,8 @@ from google.adk.agents import Agent
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
from google.adk.tools.toolbox_toolset import ToolboxToolset
from google.genai import types # For constructing message content
from toolbox_core import ToolboxSyncClient
import os
os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'True'
@@ -623,48 +623,47 @@ os.environ['GOOGLE_CLOUD_LOCATION'] = 'us-central1'
# --- Load Tools from Toolbox ---
# TODO(developer): Ensure the Toolbox server is running at <http://127.0.0.1:5000>
# TODO(developer): Ensure the Toolbox server is running at http://127.0.0.1:5000
toolset = ToolboxToolset(server_url="http://127.0.0.1:5000")
with ToolboxSyncClient("<http://127.0.0.1:5000>") as toolbox_client:
# TODO(developer): Replace "my-toolset" with the actual ID of your toolset as configured in your MCP Toolbox server.
agent_toolset = toolbox_client.load_toolset("my-toolset")
# --- Define the Agent's Prompt ---
prompt = """
You're a helpful hotel assistant. You handle hotel searching, booking and
cancellations. When the user searches for a hotel, mention it's name, id,
location and price tier. Always mention hotel ids while performing any
searches. This is very important for any operations. For any bookings or
cancellations, please provide the appropriate confirmation. Be sure to
update checkin or checkout dates if mentioned by the user.
Don't ask for confirmations from the user.
"""
# --- Define the Agent's Prompt ---
prompt = """
You're a helpful hotel assistant. You handle hotel searching, booking and
cancellations. When the user searches for a hotel, mention it's name, id,
location and price tier. Always mention hotel ids while performing any
searches. This is very important for any operations. For any bookings or
cancellations, please provide the appropriate confirmation. Be sure to
update checkin or checkout dates if mentioned by the user.
Don't ask for confirmations from the user.
"""
# --- Configure the Agent ---
# --- Configure the Agent ---
root_agent = Agent(
model='gemini-2.0-flash-001',
name='hotel_agent',
description='A helpful AI assistant that can search and book hotels.',
instruction=prompt,
tools=[toolset], # Pass the loaded toolset
)
root_agent = Agent(
model='gemini-2.0-flash-001',
name='hotel_agent',
description='A helpful AI assistant that can search and book hotels.',
instruction=prompt,
tools=agent_toolset, # Pass the loaded toolset
)
# --- Initialize Services for Running the Agent ---
session_service = InMemorySessionService()
artifacts_service = InMemoryArtifactService()
# --- Initialize Services for Running the Agent ---
session_service = InMemorySessionService()
artifacts_service = InMemoryArtifactService()
runner = Runner(
app_name='hotel_agent',
agent=root_agent,
artifact_service=artifacts_service,
session_service=session_service,
)
async def main():
# Create a new session for the interaction.
session = session_service.create_session(
session = await session_service.create_session(
state={}, app_name='hotel_agent', user_id='123'
)
runner = Runner(
app_name='hotel_agent',
agent=root_agent,
artifact_service=artifacts_service,
session_service=session_service,
)
# --- Define Queries and Run the Agent ---
queries = [
"Find hotels in Basel with Basel in it's name.",
@@ -687,6 +686,10 @@ with ToolboxSyncClient("<http://127.0.0.1:5000>") as toolbox_client:
for text in responses:
print(text)
import asyncio
if __name__ == "__main__":
asyncio.run(main())
{{< /tab >}}
{{< /tabpane >}}

View File

@@ -21,13 +21,13 @@ import (
// AuthServiceConfig is the interface for configuring authentication services.
type AuthServiceConfig interface {
AuthServiceConfigType() string
AuthServiceConfigKind() string
Initialize() (AuthService, error)
}
// AuthService is the interface for authentication services.
type AuthService interface {
AuthServiceType() string
AuthServiceKind() string
GetName() string
GetClaimsFromHeader(context.Context, http.Header) (map[string]any, error)
ToConfig() AuthServiceConfig

View File

@@ -23,7 +23,7 @@ import (
"google.golang.org/api/idtoken"
)
const AuthServiceType string = "google"
const AuthServiceKind string = "google"
// validate interface
var _ auth.AuthServiceConfig = Config{}
@@ -31,13 +31,13 @@ var _ auth.AuthServiceConfig = Config{}
// Auth service configuration
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
ClientID string `yaml:"clientId" validate:"required"`
}
// Returns the auth service type
func (cfg Config) AuthServiceConfigType() string {
return AuthServiceType
// Returns the auth service kind
func (cfg Config) AuthServiceConfigKind() string {
return AuthServiceKind
}
// Initialize a Google auth service
@@ -55,9 +55,9 @@ type AuthService struct {
Config
}
// Returns the auth service type
func (a AuthService) AuthServiceType() string {
return AuthServiceType
// Returns the auth service kind
func (a AuthService) AuthServiceKind() string {
return AuthServiceKind
}
func (a AuthService) ToConfig() auth.AuthServiceConfig {

View File

@@ -22,12 +22,12 @@ import (
// EmbeddingModelConfig is the interface for configuring embedding models.
type EmbeddingModelConfig interface {
EmbeddingModelConfigType() string
EmbeddingModelConfigKind() string
Initialize(context.Context) (EmbeddingModel, error)
}
type EmbeddingModel interface {
EmbeddingModelType() string
EmbeddingModelKind() string
ToConfig() EmbeddingModelConfig
EmbedParameters(context.Context, []string) ([][]float32, error)
}

View File

@@ -23,22 +23,22 @@ import (
"google.golang.org/genai"
)
const EmbeddingModelType string = "gemini"
const EmbeddingModelKind string = "gemini"
// validate interface
var _ embeddingmodels.EmbeddingModelConfig = Config{}
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Model string `yaml:"model" validate:"required"`
ApiKey string `yaml:"apiKey"`
Dimension int32 `yaml:"dimension"`
}
// Returns the embedding model type
func (cfg Config) EmbeddingModelConfigType() string {
return EmbeddingModelType
// Returns the embedding model kind
func (cfg Config) EmbeddingModelConfigKind() string {
return EmbeddingModelKind
}
// Initialize a Gemini embedding model
@@ -69,9 +69,9 @@ type EmbeddingModel struct {
Config
}
// Returns the embedding model type
func (m EmbeddingModel) EmbeddingModelType() string {
return EmbeddingModelType
// Returns the embedding model kind
func (m EmbeddingModel) EmbeddingModelKind() string {
return EmbeddingModelKind
}
func (m EmbeddingModel) ToConfig() embeddingmodels.EmbeddingModelConfig {

View File

@@ -42,7 +42,7 @@ func TestParseFromYamlGemini(t *testing.T) {
want: map[string]embeddingmodels.EmbeddingModelConfig{
"my-gemini-model": gemini.Config{
Name: "my-gemini-model",
Type: gemini.EmbeddingModelType,
Kind: gemini.EmbeddingModelKind,
Model: "text-embedding-004",
},
},
@@ -60,7 +60,7 @@ func TestParseFromYamlGemini(t *testing.T) {
want: map[string]embeddingmodels.EmbeddingModelConfig{
"complex-gemini": gemini.Config{
Name: "complex-gemini",
Type: gemini.EmbeddingModelType,
Kind: gemini.EmbeddingModelKind,
Model: "text-embedding-004",
ApiKey: "test-api-key",
Dimension: 768,

View File

@@ -43,9 +43,6 @@ tools:
clone_instance:
kind: cloud-sql-clone-instance
source: cloud-sql-admin-source
create_backup:
kind: cloud-sql-create-backup
source: cloud-sql-admin-source
toolsets:
cloud_sql_mssql_admin_tools:
@@ -57,4 +54,3 @@ toolsets:
- create_user
- wait_for_operation
- clone_instance
- create_backup

View File

@@ -43,9 +43,6 @@ tools:
clone_instance:
kind: cloud-sql-clone-instance
source: cloud-sql-admin-source
create_backup:
kind: cloud-sql-create-backup
source: cloud-sql-admin-source
toolsets:
cloud_sql_mysql_admin_tools:
@@ -57,4 +54,3 @@ toolsets:
- create_user
- wait_for_operation
- clone_instance
- create_backup

View File

@@ -46,9 +46,6 @@ tools:
postgres_upgrade_precheck:
kind: postgres-upgrade-precheck
source: cloud-sql-admin-source
create_backup:
kind: cloud-sql-create-backup
source: cloud-sql-admin-source
toolsets:
cloud_sql_postgres_admin_tools:
@@ -61,4 +58,3 @@ toolsets:
- wait_for_operation
- postgres_upgrade_precheck
- clone_instance
- create_backup

View File

@@ -27,10 +27,10 @@ type Message = prompts.Message
const kind = "custom"
// init registers this prompt type with the prompt framework.
// init registers this prompt kind with the prompt framework.
func init() {
if !prompts.Register(kind, newConfig) {
panic(fmt.Sprintf("prompt type %q already registered", kind))
panic(fmt.Sprintf("prompt kind %q already registered", kind))
}
}
@@ -56,7 +56,7 @@ type Config struct {
var _ prompts.PromptConfig = Config{}
var _ prompts.Prompt = Prompt{}
func (c Config) PromptConfigType() string {
func (c Config) PromptConfigKind() string {
return kind
}

View File

@@ -50,8 +50,8 @@ func TestConfig(t *testing.T) {
if p == nil {
t.Fatal("Initialize() returned a nil prompt")
}
if cfg.PromptConfigType() != "custom" {
t.Errorf("PromptConfigType() = %q, want %q", cfg.PromptConfigType(), "custom")
if cfg.PromptConfigKind() != "custom" {
t.Errorf("PromptConfigKind() = %q, want %q", cfg.PromptConfigKind(), "custom")
}
t.Run("Manifest", func(t *testing.T) {

View File

@@ -52,7 +52,7 @@ func DecodeConfig(ctx context.Context, kind, name string, decoder *yaml.Decoder)
}
if !found {
return nil, fmt.Errorf("unknown prompt type: %q", kind)
return nil, fmt.Errorf("unknown prompt kind: %q", kind)
}
promptConfig, err := factory(ctx, name, decoder)
@@ -63,7 +63,7 @@ func DecodeConfig(ctx context.Context, kind, name string, decoder *yaml.Decoder)
}
type PromptConfig interface {
PromptConfigType() string
PromptConfigKind() string
Initialize() (Prompt, error)
}

View File

@@ -29,16 +29,16 @@ import (
type mockPromptConfig struct {
name string
Type string
kind string
}
func (m *mockPromptConfig) PromptConfigType() string { return m.Type }
func (m *mockPromptConfig) PromptConfigKind() string { return m.kind }
func (m *mockPromptConfig) Initialize() (prompts.Prompt, error) { return nil, nil }
var errMockFactory = errors.New("mock factory error")
func mockFactory(ctx context.Context, name string, decoder *yaml.Decoder) (prompts.PromptConfig, error) {
return &mockPromptConfig{name: name, Type: "mockType"}, nil
return &mockPromptConfig{name: name, kind: "mockKind"}, nil
}
func mockErrorFactory(ctx context.Context, name string, decoder *yaml.Decoder) (prompts.PromptConfig, error) {
@@ -50,7 +50,7 @@ func TestRegistry(t *testing.T) {
ctx := context.Background()
t.Run("RegisterAndDecodeSuccess", func(t *testing.T) {
kind := "testTypeSuccess"
kind := "testKindSuccess"
if !prompts.Register(kind, mockFactory) {
t.Fatal("expected registration to succeed")
}
@@ -69,19 +69,19 @@ func TestRegistry(t *testing.T) {
}
})
t.Run("DecodeUnknownType", func(t *testing.T) {
t.Run("DecodeUnknownKind", func(t *testing.T) {
decoder := yaml.NewDecoder(strings.NewReader(""))
_, err := prompts.DecodeConfig(ctx, "unregisteredType", "testPrompt", decoder)
_, err := prompts.DecodeConfig(ctx, "unregisteredKind", "testPrompt", decoder)
if err == nil {
t.Fatal("expected an error for unknown kind, but got nil")
}
if !strings.Contains(err.Error(), "unknown prompt type") {
t.Errorf("expected error to contain 'unknown prompt type', but got: %v", err)
if !strings.Contains(err.Error(), "unknown prompt kind") {
t.Errorf("expected error to contain 'unknown prompt kind', but got: %v", err)
}
})
t.Run("FactoryReturnsError", func(t *testing.T) {
kind := "testTypeError"
kind := "testKindError"
if !prompts.Register(kind, mockErrorFactory) {
t.Fatal("expected registration to succeed")
}
@@ -105,8 +105,8 @@ func TestRegistry(t *testing.T) {
if config == nil {
t.Fatal("expected a non-nil config for default kind")
}
if config.PromptConfigType() != "custom" {
t.Errorf("expected default kind to be 'custom', but got %q", config.PromptConfigType())
if config.PromptConfigKind() != "custom" {
t.Errorf("expected default kind to be 'custom', but got %q", config.PromptConfigKind())
}
})
}

View File

@@ -198,7 +198,7 @@ func (c *AuthServiceConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(i
return fmt.Errorf("error creating decoder: %w", err)
}
switch kind {
case google.AuthServiceType:
case google.AuthServiceKind:
actual := google.Config{Name: name}
if err := dec.DecodeContext(ctx, &actual); err != nil {
return fmt.Errorf("unable to parse as %q: %w", kind, err)
@@ -242,7 +242,7 @@ func (c *EmbeddingModelConfigs) UnmarshalYAML(ctx context.Context, unmarshal fun
return fmt.Errorf("error creating decoder: %w", err)
}
switch kind {
case gemini.EmbeddingModelType:
case gemini.EmbeddingModelKind:
actual := gemini.Config{Name: name}
if err := dec.DecodeContext(ctx, &actual); err != nil {
return fmt.Errorf("unable to parse as %q: %w", kind, err)

View File

@@ -32,7 +32,7 @@ func TestUpdateServer(t *testing.T) {
"example-source": &alloydbpg.Source{
Config: alloydbpg.Config{
Name: "example-alloydb-source",
Type: "alloydb-postgres",
Kind: "alloydb-postgres",
},
},
}
@@ -92,7 +92,7 @@ func TestUpdateServer(t *testing.T) {
"example-source2": &alloydbpg.Source{
Config: alloydbpg.Config{
Name: "example-alloydb-source2",
Type: "alloydb-postgres",
Kind: "alloydb-postgres",
},
},
}

View File

@@ -82,7 +82,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
childCtx, span := instrumentation.Tracer.Start(
ctx,
"toolbox/server/source/init",
trace.WithAttributes(attribute.String("source_type", sc.SourceConfigType())),
trace.WithAttributes(attribute.String("source_kind", sc.SourceConfigKind())),
trace.WithAttributes(attribute.String("source_name", name)),
)
defer span.End()
@@ -110,7 +110,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
_, span := instrumentation.Tracer.Start(
ctx,
"toolbox/server/auth/init",
trace.WithAttributes(attribute.String("auth_type", sc.AuthServiceConfigType())),
trace.WithAttributes(attribute.String("auth_kind", sc.AuthServiceConfigKind())),
trace.WithAttributes(attribute.String("auth_name", name)),
)
defer span.End()
@@ -138,7 +138,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
_, span := instrumentation.Tracer.Start(
ctx,
"toolbox/server/embeddingmodel/init",
trace.WithAttributes(attribute.String("model_type", ec.EmbeddingModelConfigType())),
trace.WithAttributes(attribute.String("model_kind", ec.EmbeddingModelConfigKind())),
trace.WithAttributes(attribute.String("model_name", name)),
)
defer span.End()
@@ -166,7 +166,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
_, span := instrumentation.Tracer.Start(
ctx,
"toolbox/server/tool/init",
trace.WithAttributes(attribute.String("tool_type", tc.ToolConfigType())),
trace.WithAttributes(attribute.String("tool_kind", tc.ToolConfigKind())),
trace.WithAttributes(attribute.String("tool_name", name)),
)
defer span.End()
@@ -235,7 +235,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
_, span := instrumentation.Tracer.Start(
ctx,
"toolbox/server/prompt/init",
trace.WithAttributes(attribute.String("prompt_type", pc.PromptConfigType())),
trace.WithAttributes(attribute.String("prompt_kind", pc.PromptConfigKind())),
trace.WithAttributes(attribute.String("prompt_name", name)),
)
defer span.End()

View File

@@ -141,7 +141,7 @@ func TestUpdateServer(t *testing.T) {
"example-source": &alloydbpg.Source{
Config: alloydbpg.Config{
Name: "example-alloydb-source",
Type: "alloydb-postgres",
Kind: "alloydb-postgres",
},
},
}

View File

@@ -32,14 +32,14 @@ import (
"google.golang.org/api/option"
)
const SourceType string = "alloydb-admin"
const SourceKind string = "alloydb-admin"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -53,13 +53,13 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
DefaultProject string `yaml:"defaultProject"`
UseClientOAuth bool `yaml:"useClientOAuth"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -106,8 +106,8 @@ type Source struct {
Service *alloydbrestapi.Service
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {

View File

@@ -41,7 +41,7 @@ func TestParseFromYamlAlloyDBAdmin(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-alloydb-admin-instance": alloydbadmin.Config{
Name: "my-alloydb-admin-instance",
Type: alloydbadmin.SourceType,
Kind: alloydbadmin.SourceKind,
UseClientOAuth: false,
},
},
@@ -57,7 +57,7 @@ func TestParseFromYamlAlloyDBAdmin(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-alloydb-admin-instance": alloydbadmin.Config{
Name: "my-alloydb-admin-instance",
Type: alloydbadmin.SourceType,
Kind: alloydbadmin.SourceKind,
UseClientOAuth: true,
},
},

View File

@@ -29,14 +29,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "alloydb-postgres"
const SourceKind string = "alloydb-postgres"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -50,7 +50,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Region string `yaml:"region" validate:"required"`
Cluster string `yaml:"cluster" validate:"required"`
@@ -61,8 +61,8 @@ type Config struct {
Database string `yaml:"database" validate:"required"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -90,8 +90,8 @@ type Source struct {
Pool *pgxpool.Pool
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -183,7 +183,7 @@ func getConnectionConfig(ctx context.Context, user, pass, dbname string) (string
func initAlloyDBPgConnectionPool(ctx context.Context, tracer trace.Tracer, name, project, region, cluster, instance, ipType, user, pass, dbname string) (*pgxpool.Pool, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
dsn, useIAM, err := getConnectionConfig(ctx, user, pass, dbname)

View File

@@ -48,7 +48,7 @@ func TestParseFromYamlAlloyDBPg(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-pg-instance": alloydbpg.Config{
Name: "my-pg-instance",
Type: alloydbpg.SourceType,
Kind: alloydbpg.SourceKind,
Project: "my-project",
Region: "my-region",
Cluster: "my-cluster",
@@ -78,7 +78,7 @@ func TestParseFromYamlAlloyDBPg(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-pg-instance": alloydbpg.Config{
Name: "my-pg-instance",
Type: alloydbpg.SourceType,
Kind: alloydbpg.SourceKind,
Project: "my-project",
Region: "my-region",
Cluster: "my-cluster",
@@ -108,7 +108,7 @@ func TestParseFromYamlAlloyDBPg(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-pg-instance": alloydbpg.Config{
Name: "my-pg-instance",
Type: alloydbpg.SourceType,
Kind: alloydbpg.SourceKind,
Project: "my-project",
Region: "my-region",
Cluster: "my-cluster",

View File

@@ -41,7 +41,7 @@ import (
"google.golang.org/api/option"
)
const SourceType string = "bigquery"
const SourceKind string = "bigquery"
// CloudPlatformScope is a broad scope for Google Cloud Platform services.
const CloudPlatformScope = "https://www.googleapis.com/auth/cloud-platform"
@@ -65,8 +65,8 @@ type BigQuerySessionProvider func(ctx context.Context) (*Session, error)
type DataplexClientCreator func(tokenString string) (*dataplexapi.CatalogClient, error)
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -81,7 +81,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
// BigQuery configs
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Location string `yaml:"location"`
WriteMode string `yaml:"writeMode"`
@@ -119,9 +119,9 @@ func (s *StringOrStringSlice) UnmarshalYAML(unmarshal func(any) error) error {
return fmt.Errorf("cannot unmarshal %T into StringOrStringSlice", v)
}
func (r Config) SourceConfigType() string {
func (r Config) SourceConfigKind() string {
// Returns BigQuery source kind
return SourceType
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
if r.WriteMode == "" {
@@ -302,9 +302,9 @@ type Session struct {
LastUsed time.Time
}
func (s *Source) SourceType() string {
func (s *Source) SourceKind() string {
// Returns BigQuery Google SQL source kind
return SourceType
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -665,7 +665,7 @@ func initBigQueryConnection(
impersonateServiceAccount string,
scopes []string,
) (*bigqueryapi.Client, *bigqueryrestapi.Service, oauth2.TokenSource, error) {
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)
@@ -741,7 +741,7 @@ func initBigQueryConnectionWithOAuthToken(
tokenString string,
wantRestService bool,
) (*bigqueryapi.Client, *bigqueryrestapi.Service, error) {
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Construct token source
token := &oauth2.Token{
@@ -801,7 +801,7 @@ func initDataplexConnection(
var clientCreator DataplexClientCreator
var err error
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)

View File

@@ -46,7 +46,7 @@ func TestParseFromYamlBigQuery(t *testing.T) {
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Type: bigquery.SourceType,
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "",
WriteMode: "",
@@ -66,7 +66,7 @@ func TestParseFromYamlBigQuery(t *testing.T) {
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Type: bigquery.SourceType,
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "asia",
WriteMode: "blocked",
@@ -87,7 +87,7 @@ func TestParseFromYamlBigQuery(t *testing.T) {
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Type: bigquery.SourceType,
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "us",
UseClientOAuth: true,
@@ -108,7 +108,7 @@ func TestParseFromYamlBigQuery(t *testing.T) {
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Type: bigquery.SourceType,
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "us",
AllowedDatasets: []string{"my_dataset"},
@@ -128,7 +128,7 @@ func TestParseFromYamlBigQuery(t *testing.T) {
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Type: bigquery.SourceType,
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "us",
ImpersonateServiceAccount: "service-account@my-project.iam.gserviceaccount.com",
@@ -150,7 +150,7 @@ func TestParseFromYamlBigQuery(t *testing.T) {
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Type: bigquery.SourceType,
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "us",
Scopes: []string{"https://www.googleapis.com/auth/bigquery", "https://www.googleapis.com/auth/cloud-platform"},
@@ -170,7 +170,7 @@ func TestParseFromYamlBigQuery(t *testing.T) {
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Type: bigquery.SourceType,
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "us",
MaxQueryResultRows: 10,
@@ -260,7 +260,7 @@ func TestInitialize_MaxQueryResultRows(t *testing.T) {
desc: "default value",
cfg: bigquery.Config{
Name: "test-default",
Type: bigquery.SourceType,
Kind: bigquery.SourceKind,
Project: "test-project",
UseClientOAuth: true,
},
@@ -270,7 +270,7 @@ func TestInitialize_MaxQueryResultRows(t *testing.T) {
desc: "configured value",
cfg: bigquery.Config{
Name: "test-configured",
Type: bigquery.SourceType,
Kind: bigquery.SourceKind,
Project: "test-project",
UseClientOAuth: true,
MaxQueryResultRows: 100,

View File

@@ -27,14 +27,14 @@ import (
"google.golang.org/api/option"
)
const SourceType string = "bigtable"
const SourceKind string = "bigtable"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -48,13 +48,13 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Instance string `yaml:"instance" validate:"required"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -77,8 +77,8 @@ type Source struct {
Client *bigtable.Client
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -179,7 +179,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, configParam param
func initBigtableClient(ctx context.Context, tracer trace.Tracer, name, project, instance string) (*bigtable.Client, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Set up Bigtable data operations client.

View File

@@ -43,7 +43,7 @@ func TestParseFromYamlBigtableDb(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-bigtable-instance": bigtable.Config{
Name: "my-bigtable-instance",
Type: bigtable.SourceType,
Kind: bigtable.SourceKind,
Project: "my-project",
Instance: "my-instance",
},

View File

@@ -25,11 +25,11 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "cassandra"
const SourceKind string = "cassandra"
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -43,7 +43,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Hosts []string `yaml:"hosts" validate:"required"`
Keyspace string `yaml:"keyspace"`
ProtoVersion int `yaml:"protoVersion"`
@@ -68,9 +68,9 @@ func (c Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
return s, nil
}
// SourceConfigType implements sources.SourceConfig.
func (c Config) SourceConfigType() string {
return SourceType
// SourceConfigKind implements sources.SourceConfig.
func (c Config) SourceConfigKind() string {
return SourceKind
}
var _ sources.SourceConfig = Config{}
@@ -89,9 +89,9 @@ func (s *Source) ToConfig() sources.SourceConfig {
return s.Config
}
// SourceType implements sources.Source.
func (s *Source) SourceType() string {
return SourceType
// SourceKind implements sources.Source.
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) RunSQL(ctx context.Context, statement string, params parameters.ParamValues) (any, error) {
@@ -120,7 +120,7 @@ var _ sources.Source = &Source{}
func initCassandraSession(ctx context.Context, tracer trace.Tracer, c Config) (*gocql.Session, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, c.Name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, c.Name)
defer span.End()
// Validate authentication configuration

View File

@@ -43,7 +43,7 @@ func TestParseFromYamlCassandra(t *testing.T) {
want: server.SourceConfigs{
"my-cassandra-instance": cassandra.Config{
Name: "my-cassandra-instance",
Type: cassandra.SourceType,
Kind: cassandra.SourceKind,
Hosts: []string{"my-host1", "my-host2"},
Username: "",
Password: "",
@@ -77,7 +77,7 @@ func TestParseFromYamlCassandra(t *testing.T) {
want: server.SourceConfigs{
"my-cassandra-instance": cassandra.Config{
Name: "my-cassandra-instance",
Type: cassandra.SourceType,
Kind: cassandra.SourceKind,
Hosts: []string{"my-host1", "my-host2"},
Username: "user",
Password: "pass",

View File

@@ -28,14 +28,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "clickhouse"
const SourceKind string = "clickhouse"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -49,7 +49,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
Database string `yaml:"database" validate:"required"`
@@ -59,8 +59,8 @@ type Config struct {
Secure bool `yaml:"secure"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -88,8 +88,8 @@ type Source struct {
Pool *sql.DB
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -174,7 +174,7 @@ func validateConfig(protocol string) error {
func initClickHouseConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, protocol string, secure bool) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
if protocol == "" {

View File

@@ -25,10 +25,10 @@ import (
"go.opentelemetry.io/otel"
)
func TestConfigSourceConfigType(t *testing.T) {
func TestConfigSourceConfigKind(t *testing.T) {
config := Config{}
if config.SourceConfigType() != SourceType {
t.Errorf("Expected %s, got %s", SourceType, config.SourceConfigType())
if config.SourceConfigKind() != SourceKind {
t.Errorf("Expected %s, got %s", SourceKind, config.SourceConfigKind())
}
}
@@ -53,7 +53,7 @@ func TestNewConfig(t *testing.T) {
`,
expected: Config{
Name: "test-clickhouse",
Type: "clickhouse",
Kind: "clickhouse",
Host: "localhost",
Port: "8443",
User: "default",
@@ -75,7 +75,7 @@ func TestNewConfig(t *testing.T) {
`,
expected: Config{
Name: "minimal-clickhouse",
Type: "clickhouse",
Kind: "clickhouse",
Host: "127.0.0.1",
Port: "8123",
User: "testuser",
@@ -100,7 +100,7 @@ func TestNewConfig(t *testing.T) {
`,
expected: Config{
Name: "http-clickhouse",
Type: "clickhouse",
Kind: "clickhouse",
Host: "clickhouse.example.com",
Port: "8123",
User: "analytics",
@@ -125,7 +125,7 @@ func TestNewConfig(t *testing.T) {
`,
expected: Config{
Name: "secure-clickhouse",
Type: "clickhouse",
Kind: "clickhouse",
Host: "secure.clickhouse.io",
Port: "8443",
User: "secureuser",
@@ -196,10 +196,10 @@ func TestNewConfigInvalidYAML(t *testing.T) {
}
}
func TestSource_SourceType(t *testing.T) {
func TestSource_SourceKind(t *testing.T) {
source := &Source{}
if source.SourceType() != SourceType {
t.Errorf("Expected %s, got %s", SourceType, source.SourceType())
if source.SourceKind() != SourceKind {
t.Errorf("Expected %s, got %s", SourceKind, source.SourceKind())
}
}

View File

@@ -29,15 +29,15 @@ import (
"golang.org/x/oauth2/google"
)
const SourceType string = "cloud-gemini-data-analytics"
const SourceKind string = "cloud-gemini-data-analytics"
const Endpoint string = "https://geminidataanalytics.googleapis.com"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -51,13 +51,13 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
ProjectID string `yaml:"projectId" validate:"required"`
UseClientOAuth bool `yaml:"useClientOAuth"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
// Initialize initializes a Gemini Data Analytics Source instance.
@@ -102,8 +102,8 @@ type Source struct {
userAgent string
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {

View File

@@ -47,7 +47,7 @@ func TestParseFromYamlCloudGDA(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-gda-instance": cloudgda.Config{
Name: "my-gda-instance",
Type: cloudgda.SourceType,
Kind: cloudgda.SourceKind,
ProjectID: "test-project-id",
UseClientOAuth: false,
},
@@ -65,7 +65,7 @@ func TestParseFromYamlCloudGDA(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-gda-instance": cloudgda.Config{
Name: "my-gda-instance",
Type: cloudgda.SourceType,
Kind: cloudgda.SourceKind,
ProjectID: "another-project",
UseClientOAuth: true,
},
@@ -153,12 +153,12 @@ func TestInitialize(t *testing.T) {
}{
{
desc: "initialize with ADC",
cfg: cloudgda.Config{Name: "test-gda", Type: cloudgda.SourceType, ProjectID: "test-proj"},
cfg: cloudgda.Config{Name: "test-gda", Kind: cloudgda.SourceKind, ProjectID: "test-proj"},
wantClientOAuth: false,
},
{
desc: "initialize with client OAuth",
cfg: cloudgda.Config{Name: "test-gda-oauth", Type: cloudgda.SourceType, ProjectID: "test-proj", UseClientOAuth: true},
cfg: cloudgda.Config{Name: "test-gda-oauth", Kind: cloudgda.SourceKind, ProjectID: "test-proj", UseClientOAuth: true},
wantClientOAuth: true,
},
}

View File

@@ -34,7 +34,7 @@ import (
"google.golang.org/api/option"
)
const SourceType string = "cloud-healthcare"
const SourceKind string = "cloud-healthcare"
// validate interface
var _ sources.SourceConfig = Config{}
@@ -42,8 +42,8 @@ var _ sources.SourceConfig = Config{}
type HealthcareServiceCreator func(tokenString string) (*healthcare.Service, error)
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -58,7 +58,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
// Healthcare configs
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Region string `yaml:"region" validate:"required"`
Dataset string `yaml:"dataset" validate:"required"`
@@ -67,8 +67,8 @@ type Config struct {
UseClientOAuth bool `yaml:"useClientOAuth"`
}
func (c Config) SourceConfigType() string {
return SourceType
func (c Config) SourceConfigKind() string {
return SourceKind
}
func (c Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -144,7 +144,7 @@ func newHealthcareServiceCreator(ctx context.Context, tracer trace.Tracer, name
}
func initHealthcareConnectionWithOAuthToken(ctx context.Context, tracer trace.Tracer, name string, userAgent string, tokenString string) (*healthcare.Service, error) {
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Construct token source
token := &oauth2.Token{
@@ -162,7 +162,7 @@ func initHealthcareConnectionWithOAuthToken(ctx context.Context, tracer trace.Tr
}
func initHealthcareConnection(ctx context.Context, tracer trace.Tracer, name string) (*healthcare.Service, oauth2.TokenSource, error) {
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
cred, err := google.FindDefaultCredentials(ctx, healthcare.CloudHealthcareScope)
@@ -194,8 +194,8 @@ type Source struct {
allowedDICOMStores map[string]struct{}
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -517,14 +517,14 @@ func (s *Source) RetrieveRenderedDICOMInstance(storeID, study, series, sop strin
return base64String, nil
}
func (s *Source) SearchDICOM(toolType, storeID, dicomWebPath, tokenStr string, opts []googleapi.CallOption) (any, error) {
func (s *Source) SearchDICOM(toolKind, storeID, dicomWebPath, tokenStr string, opts []googleapi.CallOption) (any, error) {
svc, err := s.getService(tokenStr)
if err != nil {
return nil, err
}
name := fmt.Sprintf("projects/%s/locations/%s/datasets/%s/dicomStores/%s", s.Project(), s.Region(), s.DatasetID(), storeID)
var resp *http.Response
switch toolType {
switch toolKind {
case "cloud-healthcare-search-dicom-instances":
resp, err = svc.Projects.Locations.Datasets.DicomStores.SearchForInstances(name, dicomWebPath).Do(opts...)
case "cloud-healthcare-search-dicom-series":
@@ -532,7 +532,7 @@ func (s *Source) SearchDICOM(toolType, storeID, dicomWebPath, tokenStr string, o
case "cloud-healthcare-search-dicom-studies":
resp, err = svc.Projects.Locations.Datasets.DicomStores.SearchForStudies(name, dicomWebPath).Do(opts...)
default:
return nil, fmt.Errorf("incompatible tool type: %s", toolType)
return nil, fmt.Errorf("incompatible tool kind: %s", toolKind)
}
if err != nil {
return nil, fmt.Errorf("failed to search dicom series: %w", err)

View File

@@ -43,7 +43,7 @@ func TestParseFromYamlCloudHealthcare(t *testing.T) {
want: server.SourceConfigs{
"my-instance": cloudhealthcare.Config{
Name: "my-instance",
Type: cloudhealthcare.SourceType,
Kind: cloudhealthcare.SourceKind,
Project: "my-project",
Region: "us-central1",
Dataset: "my-dataset",
@@ -65,7 +65,7 @@ func TestParseFromYamlCloudHealthcare(t *testing.T) {
want: server.SourceConfigs{
"my-instance": cloudhealthcare.Config{
Name: "my-instance",
Type: cloudhealthcare.SourceType,
Kind: cloudhealthcare.SourceKind,
Project: "my-project",
Region: "us",
Dataset: "my-dataset",
@@ -91,7 +91,7 @@ func TestParseFromYamlCloudHealthcare(t *testing.T) {
want: server.SourceConfigs{
"my-instance": cloudhealthcare.Config{
Name: "my-instance",
Type: cloudhealthcare.SourceType,
Kind: cloudhealthcare.SourceKind,
Project: "my-project",
Region: "us",
Dataset: "my-dataset",

View File

@@ -29,14 +29,14 @@ import (
monitoring "google.golang.org/api/monitoring/v3"
)
const SourceType string = "cloud-monitoring"
const SourceKind string = "cloud-monitoring"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -50,12 +50,12 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
UseClientOAuth bool `yaml:"useClientOAuth"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
// Initialize initializes a Cloud Monitoring Source instance.
@@ -99,8 +99,8 @@ type Source struct {
userAgent string
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {

View File

@@ -42,7 +42,7 @@ func TestParseFromYamlCloudMonitoring(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-cloud-monitoring-instance": cloudmonitoring.Config{
Name: "my-cloud-monitoring-instance",
Type: cloudmonitoring.SourceType,
Kind: cloudmonitoring.SourceKind,
UseClientOAuth: false,
},
},
@@ -58,7 +58,7 @@ func TestParseFromYamlCloudMonitoring(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-cloud-monitoring-instance": cloudmonitoring.Config{
Name: "my-cloud-monitoring-instance",
Type: cloudmonitoring.SourceType,
Kind: cloudmonitoring.SourceKind,
UseClientOAuth: true,
},
},

View File

@@ -34,7 +34,7 @@ import (
sqladmin "google.golang.org/api/sqladmin/v1"
)
const SourceType string = "cloud-sql-admin"
const SourceKind string = "cloud-sql-admin"
var targetLinkRegex = regexp.MustCompile(`/projects/([^/]+)/instances/([^/]+)/databases/([^/]+)`)
@@ -42,8 +42,8 @@ var targetLinkRegex = regexp.MustCompile(`/projects/([^/]+)/instances/([^/]+)/da
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -57,13 +57,13 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
DefaultProject string `yaml:"defaultProject"`
UseClientOAuth bool `yaml:"useClientOAuth"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
// Initialize initializes a CloudSQL Admin Source instance.
@@ -110,8 +110,8 @@ type Source struct {
Service *sqladmin.Service
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -352,28 +352,6 @@ func (s *Source) GetWaitForOperations(ctx context.Context, service *sqladmin.Ser
return nil, nil
}
func (s *Source) InsertBackupRun(ctx context.Context, project, instance, location, backupDescription, accessToken string) (any, error) {
backupRun := &sqladmin.BackupRun{}
if location != "" {
backupRun.Location = location
}
if backupDescription != "" {
backupRun.Description = backupDescription
}
service, err := s.GetService(ctx, string(accessToken))
if err != nil {
return nil, err
}
resp, err := service.BackupRuns.Insert(project, instance, backupRun).Do()
if err != nil {
return nil, fmt.Errorf("error creating backup: %w", err)
}
return resp, nil
}
func generateCloudSQLConnectionMessage(ctx context.Context, source *Source, logger log.Logger, opResponse map[string]any, connectionMessageTemplate string) (string, bool) {
operationType, ok := opResponse["operationType"].(string)
if !ok || operationType != "CREATE_DATABASE" {

View File

@@ -42,7 +42,7 @@ func TestParseFromYamlCloudSQLAdmin(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-cloud-sql-admin-instance": cloudsqladmin.Config{
Name: "my-cloud-sql-admin-instance",
Type: cloudsqladmin.SourceType,
Kind: cloudsqladmin.SourceKind,
UseClientOAuth: false,
},
},
@@ -58,7 +58,7 @@ func TestParseFromYamlCloudSQLAdmin(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-cloud-sql-admin-instance": cloudsqladmin.Config{
Name: "my-cloud-sql-admin-instance",
Type: cloudsqladmin.SourceType,
Kind: cloudsqladmin.SourceKind,
UseClientOAuth: true,
},
},

View File

@@ -29,14 +29,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "cloud-sql-mssql"
const SourceKind string = "cloud-sql-mssql"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -51,7 +51,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
// Cloud SQL MSSQL configs
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Region string `yaml:"region" validate:"required"`
Instance string `yaml:"instance" validate:"required"`
@@ -62,9 +62,9 @@ type Config struct {
Database string `yaml:"database" validate:"required"`
}
func (r Config) SourceConfigType() string {
func (r Config) SourceConfigKind() string {
// Returns Cloud SQL MSSQL source kind
return SourceType
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -94,9 +94,9 @@ type Source struct {
Db *sql.DB
}
func (s *Source) SourceType() string {
func (s *Source) SourceKind() string {
// Returns Cloud SQL MSSQL source kind
return SourceType
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -152,7 +152,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
func initCloudSQLMssqlConnection(ctx context.Context, tracer trace.Tracer, name, project, region, instance, ipType, user, pass, dbname string) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)

View File

@@ -46,7 +46,7 @@ func TestParseFromYamlCloudSQLMssql(t *testing.T) {
want: server.SourceConfigs{
"my-instance": cloudsqlmssql.Config{
Name: "my-instance",
Type: cloudsqlmssql.SourceType,
Kind: cloudsqlmssql.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -74,7 +74,7 @@ func TestParseFromYamlCloudSQLMssql(t *testing.T) {
want: server.SourceConfigs{
"my-instance": cloudsqlmssql.Config{
Name: "my-instance",
Type: cloudsqlmssql.SourceType,
Kind: cloudsqlmssql.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -102,7 +102,7 @@ func TestParseFromYamlCloudSQLMssql(t *testing.T) {
want: server.SourceConfigs{
"my-instance": cloudsqlmssql.Config{
Name: "my-instance",
Type: cloudsqlmssql.SourceType,
Kind: cloudsqlmssql.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",

View File

@@ -30,14 +30,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "cloud-sql-mysql"
const SourceKind string = "cloud-sql-mysql"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -51,7 +51,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Region string `yaml:"region" validate:"required"`
Instance string `yaml:"instance" validate:"required"`
@@ -61,8 +61,8 @@ type Config struct {
Database string `yaml:"database" validate:"required"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -90,8 +90,8 @@ type Source struct {
Pool *sql.DB
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -184,7 +184,7 @@ func getConnectionConfig(ctx context.Context, user, pass string) (string, string
func initCloudSQLMySQLConnectionPool(ctx context.Context, tracer trace.Tracer, name, project, region, instance, ipType, user, pass, dbname string) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Configure the driver to connect to the database

View File

@@ -46,7 +46,7 @@ func TestParseFromYamlCloudSQLMySQL(t *testing.T) {
want: server.SourceConfigs{
"my-mysql-instance": cloudsqlmysql.Config{
Name: "my-mysql-instance",
Type: cloudsqlmysql.SourceType,
Kind: cloudsqlmysql.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -74,7 +74,7 @@ func TestParseFromYamlCloudSQLMySQL(t *testing.T) {
want: server.SourceConfigs{
"my-mysql-instance": cloudsqlmysql.Config{
Name: "my-mysql-instance",
Type: cloudsqlmysql.SourceType,
Kind: cloudsqlmysql.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -102,7 +102,7 @@ func TestParseFromYamlCloudSQLMySQL(t *testing.T) {
want: server.SourceConfigs{
"my-mysql-instance": cloudsqlmysql.Config{
Name: "my-mysql-instance",
Type: cloudsqlmysql.SourceType,
Kind: cloudsqlmysql.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -130,7 +130,7 @@ func TestParseFromYamlCloudSQLMySQL(t *testing.T) {
want: server.SourceConfigs{
"my-mysql-instance": cloudsqlmysql.Config{
Name: "my-mysql-instance",
Type: cloudsqlmysql.SourceType,
Kind: cloudsqlmysql.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",

View File

@@ -28,14 +28,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "cloud-sql-postgres"
const SourceKind string = "cloud-sql-postgres"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -49,7 +49,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Region string `yaml:"region" validate:"required"`
Instance string `yaml:"instance" validate:"required"`
@@ -59,8 +59,8 @@ type Config struct {
Password string `yaml:"password"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -88,8 +88,8 @@ type Source struct {
Pool *pgxpool.Pool
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -162,7 +162,7 @@ func getConnectionConfig(ctx context.Context, user, pass, dbname string) (string
func initCloudSQLPgConnectionPool(ctx context.Context, tracer trace.Tracer, name, project, region, instance, ipType, user, pass, dbname string) (*pgxpool.Pool, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Configure the driver to connect to the database

View File

@@ -46,7 +46,7 @@ func TestParseFromYamlCloudSQLPg(t *testing.T) {
want: server.SourceConfigs{
"my-pg-instance": cloudsqlpg.Config{
Name: "my-pg-instance",
Type: cloudsqlpg.SourceType,
Kind: cloudsqlpg.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -74,7 +74,7 @@ func TestParseFromYamlCloudSQLPg(t *testing.T) {
want: server.SourceConfigs{
"my-pg-instance": cloudsqlpg.Config{
Name: "my-pg-instance",
Type: cloudsqlpg.SourceType,
Kind: cloudsqlpg.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -102,7 +102,7 @@ func TestParseFromYamlCloudSQLPg(t *testing.T) {
want: server.SourceConfigs{
"my-pg-instance": cloudsqlpg.Config{
Name: "my-pg-instance",
Type: cloudsqlpg.SourceType,
Kind: cloudsqlpg.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",
@@ -130,7 +130,7 @@ func TestParseFromYamlCloudSQLPg(t *testing.T) {
want: server.SourceConfigs{
"my-pg-instance": cloudsqlpg.Config{
Name: "my-pg-instance",
Type: cloudsqlpg.SourceType,
Kind: cloudsqlpg.SourceKind,
Project: "my-project",
Region: "my-region",
Instance: "my-instance",

View File

@@ -29,14 +29,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "couchbase"
const SourceKind string = "couchbase"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -50,7 +50,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
ConnectionString string `yaml:"connectionString" validate:"required"`
Bucket string `yaml:"bucket" validate:"required"`
Scope string `yaml:"scope" validate:"required"`
@@ -66,8 +66,8 @@ type Config struct {
QueryScanConsistency uint `yaml:"queryScanConsistency"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -96,8 +96,8 @@ type Source struct {
Scope *gocb.Scope
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {

View File

@@ -45,7 +45,7 @@ func TestParseFromYamlCouchbase(t *testing.T) {
want: server.SourceConfigs{
"my-couchbase-instance": couchbase.Config{
Name: "my-couchbase-instance",
Type: couchbase.SourceType,
Kind: couchbase.SourceKind,
ConnectionString: "localhost",
Username: "Administrator",
Password: "password",
@@ -74,7 +74,7 @@ func TestParseFromYamlCouchbase(t *testing.T) {
want: server.SourceConfigs{
"my-couchbase-instance": couchbase.Config{
Name: "my-couchbase-instance",
Type: couchbase.SourceType,
Kind: couchbase.SourceKind,
ConnectionString: "couchbases://localhost",
Bucket: "travel-sample",
Scope: "inventory",

View File

@@ -29,14 +29,14 @@ import (
"google.golang.org/api/option"
)
const SourceType string = "dataplex"
const SourceKind string = "dataplex"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -51,13 +51,13 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
// Dataplex configs
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
}
func (r Config) SourceConfigType() string {
func (r Config) SourceConfigKind() string {
// Returns Dataplex source kind
return SourceType
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -81,9 +81,9 @@ type Source struct {
Client *dataplexapi.CatalogClient
}
func (s *Source) SourceType() string {
func (s *Source) SourceKind() string {
// Returns Dataplex source kind
return SourceType
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -104,7 +104,7 @@ func initDataplexConnection(
name string,
project string,
) (*dataplexapi.CatalogClient, error) {
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
cred, err := google.FindDefaultCredentials(ctx)

View File

@@ -41,7 +41,7 @@ func TestParseFromYamlDataplex(t *testing.T) {
want: server.SourceConfigs{
"my-instance": dataplex.Config{
Name: "my-instance",
Type: dataplex.SourceType,
Kind: dataplex.SourceKind,
Project: "my-project",
},
},

View File

@@ -30,14 +30,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "dgraph"
const SourceKind string = "dgraph"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -67,7 +67,7 @@ type DgraphClient struct {
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
DgraphUrl string `yaml:"dgraphUrl" validate:"required"`
User string `yaml:"user"`
Password string `yaml:"password"`
@@ -75,8 +75,8 @@ type Config struct {
ApiKey string `yaml:"apiKey"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -103,8 +103,8 @@ type Source struct {
Client *DgraphClient `yaml:"client"`
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -139,7 +139,7 @@ func (s *Source) RunSQL(statement string, params parameters.ParamValues, isQuery
func initDgraphHttpClient(ctx context.Context, tracer trace.Tracer, r Config) (*DgraphClient, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, r.Name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, r.Name)
defer span.End()
if r.DgraphUrl == "" {

View File

@@ -45,7 +45,7 @@ func TestParseFromYamlDgraph(t *testing.T) {
want: server.SourceConfigs{
"my-dgraph-instance": dgraph.Config{
Name: "my-dgraph-instance",
Type: dgraph.SourceType,
Kind: dgraph.SourceKind,
DgraphUrl: "https://localhost:8080",
ApiKey: "abc123",
Password: "pass@123",
@@ -65,7 +65,7 @@ func TestParseFromYamlDgraph(t *testing.T) {
want: server.SourceConfigs{
"my-dgraph-instance": dgraph.Config{
Name: "my-dgraph-instance",
Type: dgraph.SourceType,
Kind: dgraph.SourceKind,
DgraphUrl: "https://localhost:8080",
},
},

View File

@@ -30,14 +30,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "elasticsearch"
const SourceKind string = "elasticsearch"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -51,15 +51,15 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Addresses []string `yaml:"addresses" validate:"required"`
Username string `yaml:"username"`
Password string `yaml:"password"`
APIKey string `yaml:"apikey"`
}
func (c Config) SourceConfigType() string {
return SourceType
func (c Config) SourceConfigKind() string {
return SourceKind
}
type EsClient interface {
@@ -139,9 +139,9 @@ func (c Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
return s, nil
}
// SourceType returns the kind string for this source.
func (s *Source) SourceType() string {
return SourceType
// SourceKind returns the kind string for this source.
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {

View File

@@ -43,7 +43,7 @@ func TestParseFromYamlElasticsearch(t *testing.T) {
want: server.SourceConfigs{
"my-es-instance": elasticsearch.Config{
Name: "my-es-instance",
Type: elasticsearch.SourceType,
Kind: elasticsearch.SourceKind,
Addresses: []string{"http://localhost:9200"},
APIKey: "somekey",
},

View File

@@ -27,13 +27,13 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
)
const SourceType string = "firebird"
const SourceKind string = "firebird"
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -47,7 +47,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user" validate:"required"`
@@ -55,8 +55,8 @@ type Config struct {
Database string `yaml:"database" validate:"required"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -84,8 +84,8 @@ type Source struct {
Db *sql.DB
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -144,7 +144,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
}
func initFirebirdConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string) (*sql.DB, error) {
_, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
_, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// urlExample := "user:password@host:port/path/to/database.fdb"

View File

@@ -45,7 +45,7 @@ func TestParseFromYamlFirebird(t *testing.T) {
want: server.SourceConfigs{
"my-fdb-instance": firebird.Config{
Name: "my-fdb-instance",
Type: firebird.SourceType,
Kind: firebird.SourceKind,
Host: "my-host",
Port: "my-port",
Database: "my_db",

View File

@@ -31,14 +31,14 @@ import (
"google.golang.org/genproto/googleapis/type/latlng"
)
const SourceType string = "firestore"
const SourceKind string = "firestore"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -53,14 +53,14 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
// Firestore configs
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Database string `yaml:"database"` // Optional, defaults to "(default)"
}
func (r Config) SourceConfigType() string {
func (r Config) SourceConfigKind() string {
// Returns Firestore source kind
return SourceType
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -92,9 +92,9 @@ type Source struct {
RulesClient *firebaserules.Service
}
func (s *Source) SourceType() string {
func (s *Source) SourceKind() string {
// Returns Firestore source kind
return SourceType
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -594,7 +594,7 @@ func initFirestoreConnection(
project string,
database string,
) (*firestore.Client, error) {
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)

View File

@@ -42,7 +42,7 @@ func TestParseFromYamlFirestore(t *testing.T) {
want: server.SourceConfigs{
"my-firestore": firestore.Config{
Name: "my-firestore",
Type: firestore.SourceType,
Kind: firestore.SourceKind,
Project: "my-project",
Database: "",
},
@@ -60,7 +60,7 @@ func TestParseFromYamlFirestore(t *testing.T) {
want: server.SourceConfigs{
"my-firestore": firestore.Config{
Name: "my-firestore",
Type: firestore.SourceType,
Kind: firestore.SourceKind,
Project: "my-project",
Database: "my-database",
},

View File

@@ -29,14 +29,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "http"
const SourceKind string = "http"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -50,7 +50,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
BaseURL string `yaml:"baseUrl"`
Timeout string `yaml:"timeout"`
DefaultHeaders map[string]string `yaml:"headers"`
@@ -58,8 +58,8 @@ type Config struct {
DisableSslVerification bool `yaml:"disableSslVerification"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
// Initialize initializes an HTTP Source instance.
@@ -122,8 +122,8 @@ type Source struct {
client *http.Client
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {

View File

@@ -42,7 +42,7 @@ func TestParseFromYamlHttp(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-http-instance": http.Config{
Name: "my-http-instance",
Type: http.SourceType,
Kind: http.SourceKind,
BaseURL: "http://test_server/",
Timeout: "30s",
DisableSslVerification: false,
@@ -68,7 +68,7 @@ func TestParseFromYamlHttp(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-http-instance": http.Config{
Name: "my-http-instance",
Type: http.SourceType,
Kind: http.SourceKind,
BaseURL: "http://test_server/",
Timeout: "10s",
DefaultHeaders: map[string]string{"Authorization": "test_header", "Custom-Header": "custom"},

View File

@@ -33,14 +33,14 @@ import (
v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4"
)
const SourceType string = "looker"
const SourceKind string = "looker"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -64,7 +64,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
BaseURL string `yaml:"base_url" validate:"required"`
ClientId string `yaml:"client_id"`
ClientSecret string `yaml:"client_secret"`
@@ -79,8 +79,8 @@ type Config struct {
SessionLength int64 `yaml:"sessionLength"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
// Initialize initializes a Looker Source instance.
@@ -154,8 +154,8 @@ type Source struct {
AuthTokenHeaderName string
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {

View File

@@ -44,7 +44,7 @@ func TestParseFromYamlLooker(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-looker-instance": looker.Config{
Name: "my-looker-instance",
Type: looker.SourceType,
Kind: looker.SourceKind,
BaseURL: "http://example.looker.com/",
ClientId: "jasdl;k;tjl",
ClientSecret: "sdakl;jgflkasdfkfg",

View File

@@ -27,14 +27,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "mindsdb"
const SourceKind string = "mindsdb"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -48,7 +48,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user" validate:"required"`
@@ -57,8 +57,8 @@ type Config struct {
QueryTimeout string `yaml:"queryTimeout"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -86,8 +86,8 @@ type Source struct {
Pool *sql.DB
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -159,7 +159,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
func initMindsDBConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Configure the driver to connect to the database

View File

@@ -45,7 +45,7 @@ func TestParseFromYamlMindsDB(t *testing.T) {
want: server.SourceConfigs{
"my-mindsdb-instance": mindsdb.Config{
Name: "my-mindsdb-instance",
Type: mindsdb.SourceType,
Kind: mindsdb.SourceKind,
Host: "0.0.0.0",
Port: "my-port",
Database: "my_db",
@@ -70,7 +70,7 @@ func TestParseFromYamlMindsDB(t *testing.T) {
want: server.SourceConfigs{
"my-mindsdb-instance": mindsdb.Config{
Name: "my-mindsdb-instance",
Type: mindsdb.SourceType,
Kind: mindsdb.SourceKind,
Host: "0.0.0.0",
Port: "my-port",
Database: "my_db",

View File

@@ -29,14 +29,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "mongodb"
const SourceKind string = "mongodb"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -50,12 +50,12 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Uri string `yaml:"uri" validate:"required"` // MongoDB Atlas connection URI
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -84,8 +84,8 @@ type Source struct {
Client *mongo.Client
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -293,7 +293,7 @@ func (s *Source) DeleteOne(ctx context.Context, filterString, database, collecti
func initMongoDBClient(ctx context.Context, tracer trace.Tracer, name, uri string) (*mongo.Client, error) {
// Start a tracing span
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)

View File

@@ -41,7 +41,7 @@ func TestParseFromYamlMongoDB(t *testing.T) {
want: server.SourceConfigs{
"mongo-db": mongodb.Config{
Name: "mongo-db",
Type: mongodb.SourceType,
Kind: mongodb.SourceKind,
Uri: "mongodb+srv://username:password@host/dbname",
},
},

View File

@@ -28,14 +28,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "mssql"
const SourceKind string = "mssql"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -50,7 +50,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
// Cloud SQL MSSQL configs
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user" validate:"required"`
@@ -59,9 +59,9 @@ type Config struct {
Encrypt string `yaml:"encrypt"`
}
func (r Config) SourceConfigType() string {
func (r Config) SourceConfigKind() string {
// Returns Cloud SQL MSSQL source kind
return SourceType
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -91,9 +91,9 @@ type Source struct {
Db *sql.DB
}
func (s *Source) SourceType() string {
func (s *Source) SourceKind() string {
// Returns Cloud SQL MSSQL source kind
return SourceType
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -156,7 +156,7 @@ func initMssqlConnection(
error,
) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)

View File

@@ -45,7 +45,7 @@ func TestParseFromYamlMssql(t *testing.T) {
want: server.SourceConfigs{
"my-mssql-instance": mssql.Config{
Name: "my-mssql-instance",
Type: mssql.SourceType,
Kind: mssql.SourceKind,
Host: "0.0.0.0",
Port: "my-port",
Database: "my_db",
@@ -70,7 +70,7 @@ func TestParseFromYamlMssql(t *testing.T) {
want: server.SourceConfigs{
"my-mssql-instance": mssql.Config{
Name: "my-mssql-instance",
Type: mssql.SourceType,
Kind: mssql.SourceKind,
Host: "0.0.0.0",
Port: "my-port",
Database: "my_db",

View File

@@ -30,14 +30,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "mysql"
const SourceKind string = "mysql"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -51,7 +51,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user" validate:"required"`
@@ -61,8 +61,8 @@ type Config struct {
QueryParams map[string]string `yaml:"queryParams"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -90,8 +90,8 @@ type Source struct {
Pool *sql.DB
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -158,7 +158,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
func initMySQLConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string, queryParams map[string]string) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Build query parameters via url.Values for deterministic order and proper escaping.

View File

@@ -50,7 +50,7 @@ func TestParseFromYamlCloudSQLMySQL(t *testing.T) {
want: server.SourceConfigs{
"my-mysql-instance": mysql.Config{
Name: "my-mysql-instance",
Type: mysql.SourceType,
Kind: mysql.SourceKind,
Host: "0.0.0.0",
Port: "my-port",
Database: "my_db",
@@ -75,7 +75,7 @@ func TestParseFromYamlCloudSQLMySQL(t *testing.T) {
want: server.SourceConfigs{
"my-mysql-instance": mysql.Config{
Name: "my-mysql-instance",
Type: mysql.SourceType,
Kind: mysql.SourceKind,
Host: "0.0.0.0",
Port: "my-port",
Database: "my_db",
@@ -103,7 +103,7 @@ func TestParseFromYamlCloudSQLMySQL(t *testing.T) {
want: server.SourceConfigs{
"my-mysql-instance": mysql.Config{
Name: "my-mysql-instance",
Type: mysql.SourceType,
Kind: mysql.SourceKind,
Host: "0.0.0.0",
Port: "my-port",
Database: "my_db",
@@ -211,7 +211,7 @@ func TestFailInitialization(t *testing.T) {
cfg := mysql.Config{
Name: "instance",
Type: "mysql",
Kind: "mysql",
Host: "localhost",
Port: "3306",
Database: "db",

View File

@@ -29,7 +29,7 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "neo4j"
const SourceKind string = "neo4j"
var sourceClassifier *classifier.QueryClassifier = classifier.NewQueryClassifier()
@@ -37,8 +37,8 @@ var sourceClassifier *classifier.QueryClassifier = classifier.NewQueryClassifier
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -52,15 +52,15 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Uri string `yaml:"uri" validate:"required"`
User string `yaml:"user" validate:"required"`
Password string `yaml:"password" validate:"required"`
Database string `yaml:"database" validate:"required"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -91,8 +91,8 @@ type Source struct {
Driver neo4j.DriverWithContext
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -182,7 +182,7 @@ func addPlanChildren(p neo4j.Plan) []map[string]any {
func initNeo4jDriver(ctx context.Context, tracer trace.Tracer, uri, user, password, name string) (neo4j.DriverWithContext, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
auth := neo4j.BasicAuth(user, password, "")

View File

@@ -44,7 +44,7 @@ func TestParseFromYamlNeo4j(t *testing.T) {
want: server.SourceConfigs{
"my-neo4j-instance": neo4j.Config{
Name: "my-neo4j-instance",
Type: neo4j.SourceType,
Kind: neo4j.SourceKind,
Uri: "neo4j+s://my-host:7687",
Database: "my_db",
User: "my_user",

View File

@@ -27,14 +27,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "oceanbase"
const SourceKind string = "oceanbase"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -48,7 +48,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user" validate:"required"`
@@ -57,8 +57,8 @@ type Config struct {
QueryTimeout string `yaml:"queryTimeout"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -86,8 +86,8 @@ type Source struct {
Pool *sql.DB
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -153,7 +153,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
}
func initOceanBaseConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string) (*sql.DB, error) {
_, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
_, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?parseTime=true", user, pass, host, port, dbname)

View File

@@ -46,7 +46,7 @@ func TestParseFromYamlOceanBase(t *testing.T) {
want: server.SourceConfigs{
"my-oceanbase-instance": oceanbase.Config{
Name: "my-oceanbase-instance",
Type: oceanbase.SourceType,
Kind: oceanbase.SourceKind,
Host: "0.0.0.0",
Port: "2881",
Database: "ob_db",
@@ -71,7 +71,7 @@ func TestParseFromYamlOceanBase(t *testing.T) {
want: server.SourceConfigs{
"my-oceanbase-instance": oceanbase.Config{
Name: "my-oceanbase-instance",
Type: oceanbase.SourceType,
Kind: oceanbase.SourceKind,
Host: "0.0.0.0",
Port: "2881",
Database: "ob_db",

View File

@@ -18,14 +18,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "oracle"
const SourceKind string = "oracle"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -45,7 +45,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
ConnectionString string `yaml:"connectionString,omitempty"`
TnsAlias string `yaml:"tnsAlias,omitempty"`
TnsAdmin string `yaml:"tnsAdmin,omitempty"`
@@ -95,8 +95,8 @@ func (c Config) validate() error {
return nil
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -124,8 +124,8 @@ type Source struct {
DB *sql.DB
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -239,7 +239,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
func initOracleConnection(ctx context.Context, tracer trace.Tracer, config Config) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, config.Name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, config.Name)
defer span.End()
logger, err := util.LoggerFromContext(ctx)

View File

@@ -33,7 +33,7 @@ func TestParseFromYamlOracle(t *testing.T) {
want: server.SourceConfigs{
"my-oracle-cs": oracle.Config{
Name: "my-oracle-cs",
Type: oracle.SourceType,
Kind: oracle.SourceKind,
ConnectionString: "my-host:1521/XEPDB1",
User: "my_user",
Password: "my_pass",
@@ -56,7 +56,7 @@ func TestParseFromYamlOracle(t *testing.T) {
want: server.SourceConfigs{
"my-oracle-host": oracle.Config{
Name: "my-oracle-host",
Type: oracle.SourceType,
Kind: oracle.SourceKind,
Host: "my-host",
Port: 1521,
ServiceName: "ORCLPDB",
@@ -81,7 +81,7 @@ func TestParseFromYamlOracle(t *testing.T) {
want: server.SourceConfigs{
"my-oracle-tns-oci": oracle.Config{
Name: "my-oracle-tns-oci",
Type: oracle.SourceType,
Kind: oracle.SourceKind,
TnsAlias: "FINANCE_DB",
TnsAdmin: "/opt/oracle/network/admin",
User: "my_user",

View File

@@ -28,14 +28,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "postgres"
const SourceKind string = "postgres"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -49,7 +49,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user" validate:"required"`
@@ -58,8 +58,8 @@ type Config struct {
QueryParams map[string]string `yaml:"queryParams"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -87,8 +87,8 @@ type Source struct {
Pool *pgxpool.Pool
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -128,7 +128,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
func initPostgresConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, queryParams map[string]string) (*pgxpool.Pool, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)
if err != nil {

View File

@@ -47,7 +47,7 @@ func TestParseFromYamlPostgres(t *testing.T) {
want: server.SourceConfigs{
"my-pg-instance": postgres.Config{
Name: "my-pg-instance",
Type: postgres.SourceType,
Kind: postgres.SourceKind,
Host: "my-host",
Port: "my-port",
Database: "my_db",
@@ -74,7 +74,7 @@ func TestParseFromYamlPostgres(t *testing.T) {
want: server.SourceConfigs{
"my-pg-instance": postgres.Config{
Name: "my-pg-instance",
Type: postgres.SourceType,
Kind: postgres.SourceKind,
Host: "my-host",
Port: "my-port",
Database: "my_db",

View File

@@ -24,14 +24,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "redis"
const SourceKind string = "redis"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -45,7 +45,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Address []string `yaml:"address" validate:"required"`
Username string `yaml:"username"`
Password string `yaml:"password"`
@@ -54,8 +54,8 @@ type Config struct {
ClusterEnabled bool `yaml:"clusterEnabled"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
// RedisClient is an interface for `redis.Client` and `redis.ClusterClient
@@ -141,8 +141,8 @@ type Source struct {
Client RedisClient
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {

View File

@@ -43,7 +43,7 @@ func TestParseFromYamlRedis(t *testing.T) {
want: server.SourceConfigs{
"my-redis-instance": redis.Config{
Name: "my-redis-instance",
Type: redis.SourceType,
Kind: redis.SourceKind,
Address: []string{"127.0.0.1"},
ClusterEnabled: false,
UseGCPIAM: false,
@@ -66,7 +66,7 @@ func TestParseFromYamlRedis(t *testing.T) {
want: server.SourceConfigs{
"my-redis-instance": redis.Config{
Name: "my-redis-instance",
Type: redis.SourceType,
Kind: redis.SourceKind,
Address: []string{"127.0.0.1"},
Password: "my-pass",
Database: 1,

View File

@@ -33,14 +33,14 @@ import (
"google.golang.org/protobuf/encoding/protojson"
)
const SourceType string = "serverless-spark"
const SourceKind string = "serverless-spark"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -54,13 +54,13 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Location string `yaml:"location" validate:"required"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -94,8 +94,8 @@ type Source struct {
OpsClient *longrunning.OperationsClient
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {

View File

@@ -42,7 +42,7 @@ func TestParseFromYamlServerlessSpark(t *testing.T) {
want: server.SourceConfigs{
"my-instance": serverlessspark.Config{
Name: "my-instance",
Type: serverlessspark.SourceType,
Kind: serverlessspark.SourceKind,
Project: "my-project",
Location: "my-location",
},

View File

@@ -29,15 +29,15 @@ import (
"go.opentelemetry.io/otel/trace"
)
// SourceType for SingleStore source
const SourceType string = "singlestore"
// SourceKind for SingleStore source
const SourceKind string = "singlestore"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -52,7 +52,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
// Config holds the configuration parameters for connecting to a SingleStore database.
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user" validate:"required"`
@@ -61,9 +61,9 @@ type Config struct {
QueryTimeout string `yaml:"queryTimeout"`
}
// SourceConfigType returns the kind of the source configuration.
func (r Config) SourceConfigType() string {
return SourceType
// SourceConfigKind returns the kind of the source configuration.
func (r Config) SourceConfigKind() string {
return SourceKind
}
// Initialize sets up the SingleStore connection pool and returns a Source.
@@ -93,9 +93,9 @@ type Source struct {
Pool *sql.DB
}
// SourceType returns the kind of the source configuration.
func (s *Source) SourceType() string {
return SourceType
// SourceKind returns the kind of the source configuration.
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -162,7 +162,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
func initSingleStoreConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Configure the driver to connect to the database

View File

@@ -45,7 +45,7 @@ func TestParseFromYaml(t *testing.T) {
want: server.SourceConfigs{
"my-s2-instance": singlestore.Config{
Name: "my-s2-instance",
Type: singlestore.SourceType,
Kind: singlestore.SourceKind,
Host: "0.0.0.0",
Port: "my-port",
Database: "my_db",
@@ -70,7 +70,7 @@ func TestParseFromYaml(t *testing.T) {
want: server.SourceConfigs{
"my-s2-instance": singlestore.Config{
Name: "my-s2-instance",
Type: singlestore.SourceType,
Kind: singlestore.SourceKind,
Host: "0.0.0.0",
Port: "my-port",
Database: "my_db",

View File

@@ -25,14 +25,14 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceType string = "snowflake"
const SourceKind string = "snowflake"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -46,7 +46,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Account string `yaml:"account" validate:"required"`
User string `yaml:"user" validate:"required"`
Password string `yaml:"password" validate:"required"`
@@ -56,8 +56,8 @@ type Config struct {
Role string `yaml:"role"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -85,8 +85,8 @@ type Source struct {
DB *sqlx.DB
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -137,7 +137,7 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
func initSnowflakeConnection(ctx context.Context, tracer trace.Tracer, name, account, user, password, database, schema, warehouse, role string) (*sqlx.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Set defaults for optional parameters

View File

@@ -45,7 +45,7 @@ func TestParseFromYamlSnowflake(t *testing.T) {
want: server.SourceConfigs{
"my-snowflake-instance": snowflake.Config{
Name: "my-snowflake-instance",
Type: snowflake.SourceType,
Kind: snowflake.SourceKind,
Account: "my-account",
User: "my_user",
Password: "my_pass",

View File

@@ -31,46 +31,46 @@ var sourceRegistry = make(map[string]SourceConfigFactory)
// Register registers a new source kind with its factory.
// It returns false if the kind is already registered.
func Register(sourceType string, factory SourceConfigFactory) bool {
if _, exists := sourceRegistry[sourceType]; exists {
func Register(kind string, factory SourceConfigFactory) bool {
if _, exists := sourceRegistry[kind]; exists {
// Source with this kind already exists, do not overwrite.
return false
}
sourceRegistry[sourceType] = factory
sourceRegistry[kind] = factory
return true
}
// DecodeConfig decodes a source configuration using the registered factory for the given kind.
func DecodeConfig(ctx context.Context, sourceType string, name string, decoder *yaml.Decoder) (SourceConfig, error) {
factory, found := sourceRegistry[sourceType]
func DecodeConfig(ctx context.Context, kind string, name string, decoder *yaml.Decoder) (SourceConfig, error) {
factory, found := sourceRegistry[kind]
if !found {
return nil, fmt.Errorf("unknown source kind: %q", sourceType)
return nil, fmt.Errorf("unknown source kind: %q", kind)
}
sourceConfig, err := factory(ctx, name, decoder)
if err != nil {
return nil, fmt.Errorf("unable to parse source %q as %q: %w", name, sourceType, err)
return nil, fmt.Errorf("unable to parse source %q as %q: %w", name, kind, err)
}
return sourceConfig, err
}
// SourceConfig is the interface for configuring a source.
type SourceConfig interface {
SourceConfigType() string
SourceConfigKind() string
Initialize(ctx context.Context, tracer trace.Tracer) (Source, error)
}
// Source is the interface for the source itself.
type Source interface {
SourceType() string
SourceKind() string
ToConfig() SourceConfig
}
// InitConnectionSpan adds a span for database pool connection initialization
func InitConnectionSpan(ctx context.Context, tracer trace.Tracer, sourceType, sourceName string) (context.Context, trace.Span) {
func InitConnectionSpan(ctx context.Context, tracer trace.Tracer, sourceKind, sourceName string) (context.Context, trace.Span) {
ctx, span := tracer.Start(
ctx,
"toolbox/server/source/connect",
trace.WithAttributes(attribute.String("source_type", sourceType)),
trace.WithAttributes(attribute.String("source_kind", sourceKind)),
trace.WithAttributes(attribute.String("source_name", sourceName)),
)
return ctx, span

View File

@@ -28,14 +28,14 @@ import (
"google.golang.org/api/iterator"
)
const SourceType string = "spanner"
const SourceKind string = "spanner"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
@@ -49,15 +49,15 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"kind" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Instance string `yaml:"instance" validate:"required"`
Dialect sources.Dialect `yaml:"dialect" validate:"required"`
Database string `yaml:"database" validate:"required"`
}
func (r Config) SourceConfigType() string {
return SourceType
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
@@ -80,8 +80,8 @@ type Source struct {
Client *spanner.Client
}
func (s *Source) SourceType() string {
return SourceType
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) ToConfig() sources.SourceConfig {
@@ -171,7 +171,7 @@ func (s *Source) RunSQL(ctx context.Context, readOnly bool, statement string, pa
func initSpannerClient(ctx context.Context, tracer trace.Tracer, name, project, instance, dbname string) (*spanner.Client, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Configure the connection to the database

View File

@@ -44,7 +44,7 @@ func TestParseFromYamlSpannerDb(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-spanner-instance": spanner.Config{
Name: "my-spanner-instance",
Type: spanner.SourceType,
Kind: spanner.SourceKind,
Project: "my-project",
Instance: "my-instance",
Dialect: "googlesql",
@@ -66,7 +66,7 @@ func TestParseFromYamlSpannerDb(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-spanner-instance": spanner.Config{
Name: "my-spanner-instance",
Type: spanner.SourceType,
Kind: spanner.SourceKind,
Project: "my-project",
Instance: "my-instance",
Dialect: "googlesql",
@@ -88,7 +88,7 @@ func TestParseFromYamlSpannerDb(t *testing.T) {
want: map[string]sources.SourceConfig{
"my-spanner-instance": spanner.Config{
Name: "my-spanner-instance",
Type: spanner.SourceType,
Kind: spanner.SourceKind,
Project: "my-project",
Instance: "my-instance",
Dialect: "postgresql",

Some files were not shown because too many files have changed in this diff Show More