nat.plugins.vanna.db_utils#

Attributes#

`logger`
`RequiredSecretStr`

Classes#

`SupportedDatabase`	Supported database types for Vanna text-to-SQL.
`QueryResult`	Result from executing a database query.

Functions#

`_serialize_secret`(→ str)	Serialize SecretStr to plain string for required secret fields.
`extract_sql_from_message`(→ str)	Extract clean SQL query from various input formats.
`connect_to_databricks`(→ Any)	Connect to Databricks SQL Warehouse.
`connect_to_database`(→ Any)	Connect to a database based on type.
`execute_query`(→ QueryResult)	Execute a query and return results.
`async_execute_query`(→ QueryResult)	Execute query asynchronously and return QueryResult.
`setup_vanna_db_connection`(→ None)	Set up database connection for Vanna instance.

Module Contents#

logger#

_serialize_secret(v: pydantic.SecretStr) → str#: Serialize SecretStr to plain string for required secret fields.

RequiredSecretStr#

class SupportedDatabase#

Bases: str, enum.Enum

Supported database types for Vanna text-to-SQL.

Initialize self. See help(type(self)) for accurate signature.

DATABRICKS = 'databricks'#

class QueryResult(/, **data: Any)#

Bases: pydantic.BaseModel

Result from executing a database query.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

results: list[tuple[Any, Ellipsis]] = None#

column_names: list[str] = None#

to_dataframe() → Any#: Convert query results to a pandas DataFrame.

to_records() → list[dict[str, Any]]#: Convert query results to a list of dictionaries.

property row_count: int#

Get the number of rows in the result set.

Returns:: Number of rows

extract_sql_from_message(sql_query: str | Any) → str#

Extract clean SQL query from various input formats.

Handles: 1. Direct SQL strings (passes through) 2. BaseModel objects with ‘sql’ field (Text2SQLOutput) 3. Dictionaries with ‘sql’ key 4. Tool message format with content attribute 5. String representations of tool messages

Args:: sql_query: SQL query in various formats
Returns:: Clean SQL query string

connect_to_databricks(connection_url: str) → Any#

Connect to Databricks SQL Warehouse.

Args:: connection_url: Database connection string
Returns:: Databricks connection object

connect_to_database( database_type: str | SupportedDatabase, connection_url: str, \*\*kwargs, ) → Any#

Connect to a database based on type.

Currently only Databricks is supported.

Args:: database_type: Type of database (currently only ‘databricks’ is supported) connection_url: Database connection string kwargs: Additional database-specific parameters
Returns:: Database connection object
Raises:: ValueError: If database_type is not supported

execute_query(connection: Any, query: str) → QueryResult#

Execute a query and return results.

Args:: connection: Database connection object query: SQL query to execute
Returns:: QueryResult object containing results and column names

async async_execute_query(connection: Any, query: str) → QueryResult#

Execute query asynchronously and return QueryResult.

Args:: connection: Database connection object query: SQL query to execute
Returns:: QueryResult object containing results and column names

setup_vanna_db_connection( vn: Any, database_type: str | SupportedDatabase, connection_url: str, \*\*kwargs, ) → None#

Set up database connection for Vanna instance.

Currently only Databricks is supported.

The database Engine is stored in the Vanna instance (vn.db_engine) and will persist for the lifetime of the Vanna singleton. The Engine will be disposed when the Vanna singleton is reset.

Args:: vn: Vanna instance database_type: Type of database (currently only ‘databricks’ is supported) connection_url: Database connection string kwargs: Additional connection parameters
Raises:: ValueError: If database_type is not supported