feat(DATAGO-131057): Add table filtering to sam-sql-database-tool (#153)

Author: gregmeldrum

sam-sql-database-tool/README.md

@@ -64,6 +64,14 @@ tools:
       # schema_sample_size: 100
       # cache_ttl_seconds: 3600
 
+      # --- Table Filtering (glob patterns supported: *, ?, [seq]) ---
+      # include_tables:            # Only include matching tables in schema detection
+      #   - "customers*"
+      #   - "orders*"
+      # exclude_tables:            # Exclude matching tables from schema detection
+      #   - "bkp_*"
+      #   - "*_temp"
+
       # --- Connection Pool (optional tuning) ---
       # pool_size: 10
       # max_overflow: 10
@@ -105,6 +113,18 @@ tools:
 -   `max_enum_cardinality`: (Optional, default: `100`) Maximum number of distinct values to consider a column as an enum. Increase for columns like countries (190+), decrease for faster init times.
 -   `schema_sample_size`: (Optional, default: `100`) Number of rows to sample per table for schema detection. Increase for better accuracy on sparse data, decrease for faster init times.
 -   `cache_ttl_seconds`: (Optional, default: `3600`) Time-to-live for schema cache in seconds. After this duration, the schema will be re-detected on the next query. Set to `0` to disable caching.
+-   `include_tables`: (Optional) A list of glob patterns for tables to include in schema detection. If set, only tables matching at least one pattern are included. Supports wildcards: `*`, `?`, `[seq]`. Example: `["tms_trx*", "tms_alert*"]`.
+-   `exclude_tables`: (Optional) A list of glob patterns for tables to exclude from schema detection. Applied after `include_tables`. Supports the same wildcard syntax. Example: `["bkp_*", "*_temp", "*_dev"]`. Both options can be used together and matching is case-sensitive.
+
+    **Important: Table filtering is not access control.** These options only control which tables appear in the schema provided to the LLM. They do not prevent the LLM from executing queries against other tables in the database — for example, by querying database metadata or being prompted to access tables outside the filter. The underlying database connection still has full access to all tables the database user can see.
+
+    To reduce the likelihood of the LLM querying unfiltered tables, add an instruction to your agent such as:
+    ```
+    Only query tables that appear in your tool's schema description.
+    Do not query database metadata tables or any tables not listed in your schema.
+    ```
+
+    **For actual access control, configure the database user in the connection string with `SELECT` permissions restricted to only the allowed tables.** This is the only way to guarantee that the LLM cannot access tables outside the intended scope.
 
 #### Connection Pool Settings
 

sam-sql-database-tool/config.yaml

@@ -106,6 +106,15 @@ apps:
             schema_sample_size: 50
             cache_ttl_seconds: 1800 # 30 minutes
 
+            # Optional table filtering (glob patterns supported: *, ?, [seq]):
+            # include_tables:              # Only include matching tables in schema detection
+            #   - "tms_trx*"
+            #   - "tms_alert*"
+            # exclude_tables:              # Exclude matching tables from schema detection
+            #   - "bkp_*"
+            #   - "*_temp"
+            #   - "*_dev"
+
             # Optional connection pool tuning:
             # pool_size: 10           # Persistent connections kept in pool (default: 10)
             # max_overflow: 10        # Additional connections allowed during traffic bursts (default: 10)

sam-sql-database-tool/src/sam_sql_database_tool/services/database_service.py

@@ -1,6 +1,7 @@
 """Service for handling SQL database operations with parallel processing."""
 
 from contextlib import contextmanager
+from fnmatch import fnmatchcase
 from typing import List, Dict, Any, Generator, Optional, Tuple
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime, timedelta, timezone
@@ -31,6 +32,8 @@ def __init__(
         connect_args: Optional[dict] = None,
         isolation_level: Optional[str] = None,
         echo: bool = False,
+        include_tables: Optional[List[str]] = None,
+        exclude_tables: Optional[List[str]] = None,
     ):
         """Initialize the database service.
 
@@ -45,8 +48,12 @@ def __init__(
             connect_args: Extra kwargs passed directly to the DBAPI connect() call (default: {}).
             isolation_level: Transaction isolation level, e.g. 'READ_COMMITTED' (default: dialect default).
             echo: Log all SQL statements to the Python logger (default: False).
+            include_tables: List of glob patterns for tables to include in schema detection (default: None = all).
+            exclude_tables: List of glob patterns for tables to exclude from schema detection (default: None = none).
         """
         self.connection_string = connection_string
+        self.include_tables = include_tables
+        self.exclude_tables = exclude_tables
         self.pool_size = pool_size
         self.max_overflow = max_overflow
         self.pool_timeout = pool_timeout
@@ -311,6 +318,14 @@ def _get_approximate_row_count(self, table_name: str) -> Optional[int]:
             log.debug("Could not get row count for %s: %s", table_name, e)
             return None
 
+    def _filter_tables(self, tables: List[str]) -> List[str]:
+        """Filter table list by include/exclude glob patterns."""
+        if self.include_tables:
+            tables = [t for t in tables if any(fnmatchcase(t, p) for p in self.include_tables)]
+        if self.exclude_tables:
+            tables = [t for t in tables if not any(fnmatchcase(t, p) for p in self.exclude_tables)]
+        return tables
+
     def _looks_like_enum_column(self, column_name: str) -> bool:
         """Check if a column name suggests it might be an enum."""
         enum_patterns = [
@@ -430,7 +445,10 @@ def _process_table(
     def _compute_schema(self, max_enum_cardinality: int, sample_size: int) -> str:
         """Compute schema without caching logic."""
         log.info("Starting schema detection...")
-        tables = self.get_tables()
+        all_tables = self.get_tables()
+        tables = self._filter_tables(all_tables)
+        if len(tables) != len(all_tables):
+            log.info("Table filtering: %d of %d tables selected", len(tables), len(all_tables))
         log.debug("Processing %d tables in parallel with 5 workers...", len(tables))
         schema = {}
 

sam-sql-database-tool/src/sam_sql_database_tool/tools.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, List
 from pydantic import BaseModel, Field, model_validator, field_validator, SecretStr
 from google.genai import types as adk_types
 from solace_agent_mesh.agent.tools.dynamic_tool import DynamicTool
@@ -39,6 +39,20 @@ class DatabaseConfig(BaseModel):
         default=3600,
         description="Time-to-live for schema cache in seconds (default: 3600 = 1 hour).",
     )
+    include_tables: Optional[List[str]] = Field(
+        default=None,
+        description=(
+            "List of glob patterns for tables to include in schema detection. "
+            "If set, only matching tables are included. Supports wildcards: *, ?, [seq]."
+        ),
+    )
+    exclude_tables: Optional[List[str]] = Field(
+        default=None,
+        description=(
+            "List of glob patterns for tables to exclude from schema detection. "
+            "Applied after include_tables. Supports wildcards: *, ?, [seq]."
+        ),
+    )
 
     # Connection pool settings
     pool_size: int = Field(
@@ -99,8 +113,13 @@ def check_required_fields(self) -> 'DatabaseConfig':
                 raise ValueError(
                     "'schema_summary_override' is required when 'auto_detect_schema' is false"
                 )
+            if self.include_tables or self.exclude_tables:
+                log.warning(
+                    "Tool '%s': include_tables/exclude_tables have no effect when auto_detect_schema is false",
+                    self.tool_name,
+                )
         return self
-    
+
     def get(self, key: str, default: Any = None) -> Any:
         """Allows dictionary-like access to the model's attributes."""
         return getattr(self, key, default)
@@ -170,6 +189,8 @@ async def init(self, component: SamAgentComponent, tool_config: Dict):
                 connect_args=self.tool_config.connect_args,
                 isolation_level=self.tool_config.isolation_level,
                 echo=self.tool_config.echo,
+                include_tables=self.tool_config.include_tables,
+                exclude_tables=self.tool_config.exclude_tables,
             )
         except Exception as e:
             self._connection_healthy = False

sam-sql-database-tool/tests/unit/test_config.py

@@ -47,3 +47,57 @@ def test_manual_schema_with_override(self):
         )
         assert config.auto_detect_schema is False
         assert config.schema_summary_override is not None
+
+    def test_include_tables(self):
+        """Test config accepts include_tables as list of strings."""
+        config = DatabaseConfig(
+            tool_name="test",
+            connection_string="postgresql+psycopg2://user:password@localhost:5432/testdb",
+            include_tables=["tms_trx*", "tms_alert*"],
+        )
+        assert config.include_tables == ["tms_trx*", "tms_alert*"]
+
+    def test_exclude_tables(self):
+        """Test config accepts exclude_tables as list of strings."""
+        config = DatabaseConfig(
+            tool_name="test",
+            connection_string="postgresql+psycopg2://user:password@localhost:5432/testdb",
+            exclude_tables=["bkp_*", "*_temp"],
+        )
+        assert config.exclude_tables == ["bkp_*", "*_temp"]
+
+    def test_include_and_exclude_tables_together(self):
+        """Test config accepts both include_tables and exclude_tables."""
+        config = DatabaseConfig(
+            tool_name="test",
+            connection_string="postgresql+psycopg2://user:password@localhost:5432/testdb",
+            include_tables=["tms_trx*"],
+            exclude_tables=["*_temp"],
+        )
+        assert config.include_tables == ["tms_trx*"]
+        assert config.exclude_tables == ["*_temp"]
+
+    def test_table_filters_ignored_when_auto_detect_false(self):
+        """Test that a warning is logged when filters are set with auto_detect_schema=false."""
+        import logging
+        from unittest.mock import patch
+        logger = logging.getLogger("sam_sql_database_tool.tools")
+        with patch.object(logger, "warning") as mock_warning:
+            DatabaseConfig(
+                tool_name="test",
+                connection_string="postgresql+psycopg2://user:password@localhost:5432/testdb",
+                auto_detect_schema=False,
+                schema_summary_override="manual schema",
+                include_tables=["tms_trx*"],
+            )
+            mock_warning.assert_called_once()
+            assert "include_tables/exclude_tables have no effect" in mock_warning.call_args[0][0]
+
+    def test_table_filters_default_to_none(self):
+        """Test that include_tables and exclude_tables default to None."""
+        config = DatabaseConfig(
+            tool_name="test",
+            connection_string="postgresql+psycopg2://user:password@localhost:5432/testdb",
+        )
+        assert config.include_tables is None
+        assert config.exclude_tables is None

sam-sql-database-tool/tests/unit/test_database_service.py

@@ -124,3 +124,101 @@ def test_mysql_dialect_detection(self):
 
             service = DatabaseService("mysql+pymysql://test")
             assert service.engine.dialect.name == "mysql"
+
+
+class TestFilterTables:
+    """Unit tests for table filtering logic."""
+
+    @pytest.fixture
+    def service_with_filters(self):
+        """Create a DatabaseService with custom include/exclude filters."""
+        def _create(include_tables=None, exclude_tables=None):
+            with patch('sqlalchemy.create_engine') as mock_create_engine:
+                mock_engine = MagicMock()
+                mock_engine.dialect.name = "postgresql"
+                mock_create_engine.return_value = mock_engine
+                return DatabaseService(
+                    "postgresql://test",
+                    include_tables=include_tables,
+                    exclude_tables=exclude_tables,
+                )
+        return _create
+
+    def test_no_filters_returns_all(self, service_with_filters):
+        """No filters should return all tables unchanged."""
+        service = service_with_filters()
+        tables = ["users", "orders", "products"]
+        assert service._filter_tables(tables) == ["users", "orders", "products"]
+
+    def test_include_exact_match(self, service_with_filters):
+        """Include with exact names should only return matching tables."""
+        service = service_with_filters(include_tables=["users", "orders"])
+        tables = ["users", "orders", "products", "logs"]
+        assert service._filter_tables(tables) == ["users", "orders"]
+
+    def test_include_wildcard(self, service_with_filters):
+        """Include with wildcard patterns should match correctly."""
+        service = service_with_filters(include_tables=["tms_trx*", "tms_alert*"])
+        tables = ["tms_trx5min", "tms_trx60min", "tms_alert_anomaly", "tms_bank", "bkp_data"]
+        assert service._filter_tables(tables) == ["tms_trx5min", "tms_trx60min", "tms_alert_anomaly"]
+
+    def test_exclude_wildcard(self, service_with_filters):
+        """Exclude with wildcard patterns should remove matching tables."""
+        service = service_with_filters(exclude_tables=["bkp_*", "*_temp", "*_dev"])
+        tables = ["tms_trx5min", "bkp_old_data", "tms_data_temp", "tms_aggr_dev", "tms_alert"]
+        assert service._filter_tables(tables) == ["tms_trx5min", "tms_alert"]
+
+    def test_include_and_exclude_together(self, service_with_filters):
+        """Include and exclude together: include first, then exclude from result."""
+        service = service_with_filters(
+            include_tables=["tms_trx*"],
+            exclude_tables=["*_temp", "*_dev"],
+        )
+        tables = ["tms_trx5min", "tms_trx_temp", "tms_trx_dev", "tms_trx60min", "tms_bank"]
+        assert service._filter_tables(tables) == ["tms_trx5min", "tms_trx60min"]
+
+    def test_question_mark_wildcard(self, service_with_filters):
+        """The ? wildcard should match a single character."""
+        service = service_with_filters(include_tables=["table_?"])
+        tables = ["table_a", "table_b", "table_ab", "other"]
+        assert service._filter_tables(tables) == ["table_a", "table_b"]
+
+    def test_case_sensitive(self, service_with_filters):
+        """Matching should be case-sensitive."""
+        service = service_with_filters(include_tables=["Users"])
+        tables = ["Users", "users", "USERS"]
+        assert service._filter_tables(tables) == ["Users"]
+
+    def test_empty_include_list_returns_all(self, service_with_filters):
+        """An empty include list should be falsy and return all tables."""
+        service = service_with_filters(include_tables=[])
+        tables = ["users", "orders", "products"]
+        assert service._filter_tables(tables) == ["users", "orders", "products"]
+
+    def test_all_tables_filtered_out(self, service_with_filters):
+        """When all tables are excluded, result should be empty."""
+        service = service_with_filters(include_tables=["nonexistent_*"])
+        tables = ["users", "orders", "products"]
+        assert service._filter_tables(tables) == []
+
+    def test_compute_schema_uses_filter(self, service_with_filters):
+        """_compute_schema should only process tables that pass the filter."""
+        service = service_with_filters(include_tables=["users"])
+
+        mock_inspector = MagicMock()
+        mock_inspector.get_table_names.return_value = ["users", "orders", "products"]
+        mock_inspector.get_columns.return_value = [
+            {"name": "id", "type": "INTEGER", "nullable": False},
+        ]
+        mock_inspector.get_pk_constraint.return_value = {"constrained_columns": ["id"]}
+        mock_inspector.get_foreign_keys.return_value = []
+        mock_inspector.get_indexes.return_value = []
+
+        service.get_table_sample = MagicMock(return_value=[{"id": 1}])
+
+        with patch('sam_sql_database_tool.services.database_service.inspect', return_value=mock_inspector):
+            schema = service._compute_schema(max_enum_cardinality=100, sample_size=100)
+
+        assert "users" in schema
+        assert "orders" not in schema
+        assert "products" not in schema