add datasource tool (#6422)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->
Fetches tables (with schema, database and connection fields) in a
database.
Supports regex queries.

<img width="550" height="628" alt="CleanShot 2025-09-19 at 01 22 12"
src="https://github.com/user-attachments/assets/8b3a678a-2ded-48e0-a36e-3a63bd38aa9e"
/>

<img width="549" height="492" alt="CleanShot 2025-09-19 at 01 24 20"
src="https://github.com/user-attachments/assets/eeb456bd-383c-4dc5-8ae4-4f2e86fd6686"
/>

## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [x] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

Author: Light2Dark

marimo/_ai/_tools/tools/datasource.py

@@ -0,0 +1,117 @@
+# Copyright 2025 Marimo. All rights reserved.
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+from marimo import _loggers
+from marimo._ai._tools.base import ToolBase
+from marimo._ai._tools.types import SuccessResult
+from marimo._ai._tools.utils.exceptions import ToolExecutionError
+from marimo._data.models import DataTable
+from marimo._server.sessions import Session
+from marimo._types.ids import SessionId
+from marimo._utils.fuzzy_match import compile_regex, is_fuzzy_match
+
+LOGGER = _loggers.marimo_logger()
+
+
+@dataclass
+class GetDatabaseTablesArgs:
+    session_id: SessionId
+    query: Optional[str] = None
+
+
+@dataclass
+class TableDetails:
+    connection: str
+    database: str
+    schema: str
+    table: DataTable
+
+
+@dataclass
+class GetDatabaseTablesOutput(SuccessResult):
+    tables: list[TableDetails] = field(default_factory=list)
+
+
+class GetDatabaseTables(
+    ToolBase[GetDatabaseTablesArgs, GetDatabaseTablesOutput]
+):
+    """
+    Get information about tables in a database.
+
+    Args:
+        session_id: The session id.
+        query (optional): The query to match the database, schemas, and tables. Regex is supported.
+
+    If a query is provided, it will fuzzy match the query to the database, schemas, and tables available. If no query is provided, all tables are returned. Don't provide a query if you need to see the entire schema view.
+
+    The tables returned contain information about the database, schema and connection name to use in forming SQL queries.
+    """
+
+    def handle(self, args: GetDatabaseTablesArgs) -> GetDatabaseTablesOutput:
+        session_id = args.session_id
+        session = self.context.get_session(session_id)
+
+        return self._get_tables(session, args.query)
+
+    def _get_tables(
+        self, session: Session, query: Optional[str]
+    ) -> GetDatabaseTablesOutput:
+        session_view = session.session_view
+        data_connectors = session_view.data_connectors
+
+        if len(data_connectors.connections) == 0:
+            raise ToolExecutionError(
+                message="No databases found. Please create a connection first.",
+                code="NO_DATABASES_FOUND",
+                is_retryable=False,
+            )
+
+        tables: list[TableDetails] = []
+
+        # Pre-compile regex if query exists
+        compiled_pattern = None
+        is_regex = False
+        if query:
+            compiled_pattern, is_regex = compile_regex(query)
+
+        for connection in data_connectors.connections:
+            for database in connection.databases:
+                for schema in database.schemas:
+                    # If query is None, match all schemas
+                    # If matching, add all tables to the list
+                    if query is None or is_fuzzy_match(
+                        query, schema.name, compiled_pattern, is_regex
+                    ):
+                        for table in schema.tables:
+                            tables.append(
+                                TableDetails(
+                                    connection=connection.name,
+                                    database=database.name,
+                                    schema=schema.name,
+                                    table=table,
+                                )
+                            )
+                        continue
+                    for table in schema.tables:
+                        if is_fuzzy_match(
+                            query, table.name, compiled_pattern, is_regex
+                        ):
+                            tables.append(
+                                TableDetails(
+                                    connection=connection.name,
+                                    database=database.name,
+                                    schema=schema.name,
+                                    table=table,
+                                )
+                            )
+
+        return GetDatabaseTablesOutput(
+            tables=tables,
+            next_steps=[
+                'Example of an SQL query: _df = mo.sql(f"""SELECT * FROM database.schema.name LIMIT 100""")',
+            ],
+        )

marimo/_ai/_tools/tools_registry.py

@@ -6,6 +6,7 @@
     GetCellRuntimeData,
     GetLightweightCellMap,
 )
+from marimo._ai._tools.tools.datasource import GetDatabaseTables
 from marimo._ai._tools.tools.notebooks import GetActiveNotebooks
 from marimo._ai._tools.tools.tables_and_variables import GetTablesAndVariables
 
@@ -14,4 +15,5 @@
     GetCellRuntimeData,
     GetLightweightCellMap,
     GetTablesAndVariables,
+    GetDatabaseTables,
 ]

marimo/_utils/fuzzy_match.py

@@ -0,0 +1,36 @@
+# Copyright 2025 Marimo. All rights reserved.
+
+from __future__ import annotations
+
+import re
+
+
+def compile_regex(query: str) -> tuple[re.Pattern[str] | None, bool]:
+    """
+    Returns compiled regex pattern and whether the query is a valid regex.
+    """
+    try:
+        return re.compile(query, re.IGNORECASE), True
+    except re.error:
+        return None, False
+
+
+def is_fuzzy_match(
+    query: str,
+    name: str,
+    compiled_pattern: re.Pattern[str] | None,
+    is_regex: bool,
+) -> bool:
+    """
+    Fuzzy match using pre-compiled regex. If is not regex, fallback to substring match.
+
+    Args:
+        query: The query to match.
+        name: The name to match against.
+        compiled_pattern: Pre-compiled regex pattern (None if not regex).
+        is_regex: Whether the query is a valid regex.
+    """
+    if is_regex and compiled_pattern:
+        return bool(compiled_pattern.search(name))
+    else:
+        return query.lower() in name.lower()