Added separate executor by kosstbarz · Pull Request #252 · JetBrains/databao-agent

kosstbarz · 2026-03-10T15:28:54Z

No description provided.

Copilot

Pull request overview

Adds a new “separate” execution mode that runs SQL directly against each datasource’s native SQLAlchemy engine (instead of routing through DuckDB), including Snowflake engine creation and schema introspection to populate the system prompt.

Changes:

Introduce SeparateExecutor + SeparateGraph with run_sql_query/submit_result tools and a dedicated system prompt template.
Add SQLAlchemy-based Snowflake schema inspection (information_schema.tables/columns) to generate TableInfo/ColumnInfo for prompt schema.
Extend database adapter plumbing with try_create_sqlalchemy_engine() and a Snowflake create_sqlalchemy_engine() implementation.

Reviewed changes

Copilot reviewed 7 out of 8 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
databao/agent/sqlalchemy/schema_inspection.py	New SQLAlchemy schema inspector (currently Snowflake-only) producing `TableInfo`/`ColumnInfo`.
databao/agent/sqlalchemy/init.py	Initializes `databao.agent.sqlalchemy` package.
databao/agent/executors/separate/system_prompt.jinja	New system prompt template for the separate executor (includes Snowflake quoting guidance).
databao/agent/executors/separate/separate_executor.py	New executor that builds per-datasource SQLAlchemy engines and injects schema into the prompt.
databao/agent/executors/separate/graph.py	New LangGraph tool loop to run SQL via SQLAlchemy engines and submit results.
databao/agent/databases/snowflake_adapter.py	Adds Snowflake SQLAlchemy engine creation from connection config (password/keypair/SSO).
databao/agent/databases/databases.py	Adds `try_create_sqlalchemy_engine()` helper to delegate engine creation to adapters.
databao/agent/databases/database_adapter.py	Adds optional `create_sqlalchemy_engine()` hook on adapters (default `None`).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

databao/agent/databases/snowflake_adapter.py

Copilot · 2026-03-11T15:35:06Z

databao/agent/databases/snowflake_adapter.py

+    @classmethod
+    def create_sqlalchemy_engine(cls, config: DBConnectionConfig) -> Engine | None:
+        if not isinstance(config, SnowflakeConnectionProperties):
+            return None
+
+        from snowflake.sqlalchemy import URL  # type: ignore[import-untyped]
+
+        url_kwargs: dict[str, str] = {"account": config.account}
+        if config.user:
+            url_kwargs["user"] = config.user
+        if config.database:
+            url_kwargs["database"] = config.database
+        if config.warehouse:
+            url_kwargs["warehouse"] = config.warehouse
+        if config.role:
+            url_kwargs["role"] = config.role
+
+        connect_args: dict[str, Any] = {}
+        auth = config.auth
+        if isinstance(auth, SnowflakePasswordAuth):
+            url_kwargs["password"] = auth.password
+        elif isinstance(auth, SnowflakeKeyPairAuth):
+            connect_args["private_key"] = cls._load_private_key_bytes(auth)
+        elif isinstance(auth, SnowflakeSSOAuth):
+            url_kwargs["authenticator"] = auth.authenticator
+        else:
+            return None
+
+        if connect_args:
+            return create_engine(URL(**url_kwargs), connect_args=connect_args)
+        return create_engine(URL(**url_kwargs))


New create_sqlalchemy_engine() behavior is not covered by tests. There is already comprehensive coverage for SnowflakeAdapter helpers in tests/test_snowflake_adapter.py; please add unit tests that validate URL/connect_args generation for password, key-pair, and SSO auth (and that additional_properties are preserved).

databao/agent/executors/separate/separate_executor.py

databao/agent/executors/separate/graph.py

databao/agent/sqlalchemy/schema_inspection.py

databao/agent/executors/separate/system_prompt.jinja

SokolovYaroslav · 2026-03-16T13:17:18Z

@kosstbarz please address the copilot's comments somehow :)

I personally don't look into the code until AI is happy

… additional properties

Copilot

Pull request overview

Copilot reviewed 10 out of 13 changed files in this pull request and generated 4 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-02T00:30:41Z

databao/agent/executors/separate/separate_executor.py

+                db_tables = inspect_sqlalchemy_schema(engine)
+                # Use the registered name as table_catalog so the LLM can derive
+                # the datasource argument directly from the schema prefix.
+                tables.extend(replace(t, table_catalog=name) for t in db_tables)


_inspect_database_schema() rewrites TableInfo.table_catalog to the datasource name, but leaves columns_catalog unchanged. Downstream helpers (e.g. summarize_duckdb_schema_overview) assume a 1:1 mapping between table_catalog and columns_catalog, so this can lead to duplicated/incorrect grouping. Consider rewriting both fields (or constructing a new TableInfo) so columns_catalog stays consistent with the synthetic datasource catalog name.

Suggested change

tables.extend(replace(t, table_catalog=name) for t in db_tables)

tables.extend(

replace(t, table_catalog=name, columns_catalog=name) for t in db_tables

)

Copilot · 2026-04-02T00:30:42Z

databao/agent/executors/separate/graph.py

+
+
+def exception_to_string(e: Exception | str) -> str:
+    if isinstance(e, str):
+        return e
+    return f"Exception Name: {type(e).__name__}. Exception Desc: {e}"
+
+


This file reimplements exception_to_string() even though an identical helper exists in databao.agent.executors.utils. Reusing the shared helper avoids divergence and keeps error formatting consistent across executors.

Suggested change

def exception_to_string(e: Exception | str) -> str:

if isinstance(e, str):

return e

return f"Exception Name: {type(e).__name__}. Exception Desc: {e}"

from databao.agent.executors.utils import exception_to_string

Copilot · 2026-04-02T00:30:42Z

databao/agent/executors/separate/graph.py

+def trim_string_middle(
+    content: str, max_length: int | None, sep: str = "[...trimmed...]", front_percentage: float = 0.7
+) -> str:
+    if max_length is None or len(content) <= max_length:
+        return content
+    take_front = max(0, math.ceil(max_length * front_percentage) - len(sep) // 2)
+    take_end = max(0, max_length - take_front - len(sep))
+    return content[:take_front] + sep + content[len(content) - take_end :]
+
+
+def trim_dataframe_values(df: pd.DataFrame, max_cell_chars: int | None) -> pd.DataFrame:
+    df_sanitized = df.copy()
+    if max_cell_chars is None:
+        return df_sanitized
+
+    def trim_cell(val: Any) -> str:
+        return trim_string_middle(str(val), max_cell_chars)
+
+    for col, dtype in zip(df_sanitized.columns, df_sanitized.dtypes, strict=True):
+        if not pd.api.types.is_object_dtype(dtype) and not pd.api.types.is_string_dtype(dtype):
+            continue
+        df_sanitized[col] = df_sanitized[col].apply(trim_cell)
+    return df_sanitized


trim_string_middle() / trim_dataframe_values() duplicate logic already present in databao.agent.executors.utils. Consider importing and reusing the shared implementations to reduce maintenance overhead and keep truncation behavior consistent across executors.

Copilot · 2026-04-02T00:30:42Z

databao/agent/databases/snowflake_adapter.py

+        if auth.private_key:
+            pem_data = auth.private_key.encode()
+        elif auth.private_key_file:
+            pem_data = Path(auth.private_key_file).read_bytes()


_load_private_key_bytes() reads private_key_file directly and will propagate raw OSError/FileNotFoundError exceptions. For consistency with _create_secret_params() (which wraps file read errors in a ValueError with a clear message), consider catching OSError here and raising a ValueError that points to the invalid private_key_file path.

Suggested change

pem_data = Path(auth.private_key_file).read_bytes()

try:

pem_data = Path(auth.private_key_file).read_bytes()

except OSError as exc:

raise ValueError(f"Failed to read private key file at '{auth.private_key_file}'.") from exc

…comma in Snowflake OAuth example.

kosstbarz marked this pull request as ready for review March 11, 2026 15:13

kosstbarz requested a review from a team March 11, 2026 15:13

SokolovYaroslav requested a review from Copilot March 11, 2026 15:26

Copilot started reviewing on behalf of SokolovYaroslav March 11, 2026 15:26 View session

Copilot AI reviewed Mar 11, 2026

View reviewed changes

SokolovYaroslav marked this pull request as draft March 23, 2026 11:27

kosstbarz and others added 10 commits April 2, 2026 01:37

Added separate executor

08d10e8

Added sqlalchemy schema inspection

fdf4766

Refactor import: move make_search_context_tool to langchain_tools.

8de05a7

Add OAuth authentication method for Snowflake

cb36fa6

Use databao-context-engine 0.7.1.dev2 instead

43e4eca

Add Snowflake OAuth integration example

ac30379

Handle exclusion of SQLAlchemy-internal keys for Snowflake connection

2e68aeb

Ensure autocommit and port are excluded from Snowflake connection…

a94f027

… additional properties

Add support for OAuth authentication in Snowflake adapter

b1fcbcf

Address Copilot comments

863a5c9

catstrike force-pushed the kosst/separate-executor branch from ca1d38d to 863a5c9 Compare April 2, 2026 00:13

Add __init__.py for databao/agent/executors/separate

0d8d0a5

catstrike had a problem deploying to pypi April 2, 2026 00:19 — with GitHub Actions Failure

catstrike had a problem deploying to pypi April 2, 2026 00:25 — with GitHub Actions Failure

catstrike requested a review from Copilot April 2, 2026 00:26

Copilot started reviewing on behalf of catstrike April 2, 2026 00:27 View session

Copilot AI reviewed Apr 2, 2026

View reviewed changes

catstrike added 2 commits April 2, 2026 02:38

Address Copilot comments

ba1b677

Fix import order in test_snowflake_adapter.py and correct trailing …

8babacc

…comma in Snowflake OAuth example.

SokolovYaroslav had a problem deploying to pypi April 2, 2026 13:25 — with GitHub Actions Failure

SokolovYaroslav had a problem deploying to pypi April 2, 2026 13:33 — with GitHub Actions Failure

Try to fix dev build: Fix dependency version for databao-context-engine

66b342a

SokolovYaroslav temporarily deployed to pypi April 2, 2026 13:50 — with GitHub Actions Inactive

SokolovYaroslav added 2 commits April 2, 2026 15:51

Fix dependency version syntax for databao-context-engine

78cbb94

fix exzct version of dce

98bf950

SokolovYaroslav deployed to pypi April 2, 2026 14:42 — with GitHub Actions Active

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Added separate executor#252

Added separate executor#252
kosstbarz wants to merge 16 commits intomainfrom
kosst/separate-executor

kosstbarz commented Mar 10, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI Mar 11, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

SokolovYaroslav commented Mar 16, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 2, 2026

Uh oh!

Copilot AI Apr 2, 2026

Uh oh!

Copilot AI Apr 2, 2026

Uh oh!

Copilot AI Apr 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

kosstbarz commented Mar 10, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

SokolovYaroslav commented Mar 16, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants