airbyte.mcp.local
Local MCP operations.
1# Copyright (c) 2024 Airbyte, Inc., all rights reserved. 2"""Local MCP operations. 3 4.. include:: ../../docs/mcp-generated/local.md 5""" 6 7# No public Python API — MCP primitives are registered via decorators and 8# documented via the generated Markdown include above. Setting `__all__` to an 9# empty list tells pdoc (and other doc tools) not to surface the individual 10# tool / helper definitions as a redundant "API Documentation" list. 11__all__: list[str] = [] 12 13import sys 14import traceback 15from itertools import islice 16from pathlib import Path 17from typing import TYPE_CHECKING, Annotated, Any, Literal 18 19from fastmcp import FastMCP 20from fastmcp_extensions import mcp_tool, register_mcp_tools 21from pydantic import BaseModel, Field 22 23from airbyte import get_source 24from airbyte._util.destination_smoke_tests import ( 25 DestinationSmokeTestResult, 26 run_destination_smoke_test, 27) 28from airbyte._util.meta import is_docker_installed 29from airbyte.caches.util import get_default_cache 30from airbyte.destinations.util import get_destination 31from airbyte.mcp._arg_resolvers import resolve_connector_config, resolve_list_of_strings 32from airbyte.mcp._guards import raise_if_untrusted_execution_context 33from airbyte.registry import get_connector_metadata 34from airbyte.secrets.config import _get_secret_sources 35from airbyte.secrets.env_vars import DotenvSecretManager 36from airbyte.secrets.google_gsm import GoogleGSMSecretManager 37from airbyte.sources.base import Source 38 39 40if TYPE_CHECKING: 41 from airbyte.caches.duckdb import DuckDBCache 42 43 44_CONFIG_HELP = """ 45You can provide `config` as JSON or a Path to a YAML/JSON file. 46If a `dict` is provided, it must not contain hardcoded secrets. 47Instead, secrets should be provided using environment variables, 48and the config should reference them using the format 49`secret_reference::ENV_VAR_NAME`. 50 51You can also provide a `config_secret_name` to use a specific 52secret name for the configuration. This is useful if you want to 53validate a configuration that is stored in a secrets manager. 54 55If `config_secret_name` is provided, it should point to a string 56that contains valid JSON or YAML. 57 58If both `config` and `config_secret_name` are provided, the 59`config` will be loaded first and then the referenced secret config 60will be layered on top of the non-secret config. 61 62For declarative connectors, you can provide a `manifest_path` to 63specify a local YAML manifest file instead of using the registry 64version. This is useful for testing custom or locally-developed 65connector manifests. 66""" 67 68 69def _get_mcp_source( 70 connector_name: str, 71 override_execution_mode: Literal["auto", "docker", "python", "yaml"] = "auto", 72 *, 73 install_if_missing: bool = True, 74 manifest_path: str | Path | None, 75) -> Source: 76 """Get the MCP source for a connector. 77 78 This installs and executes a connector on the server, so it is gated by trusted 79 execution and hard-fails when trusted execution is disabled, independently of whether 80 the calling tool was hidden from the tool listing. 81 """ 82 raise_if_untrusted_execution_context("Local connector execution (`_get_mcp_source`)") 83 if manifest_path: 84 override_execution_mode = "yaml" 85 elif override_execution_mode == "auto" and is_docker_installed(): 86 override_execution_mode = "docker" 87 88 source: Source 89 if override_execution_mode == "auto": 90 # Use defaults with no overrides 91 source = get_source( 92 connector_name, 93 install_if_missing=False, 94 source_manifest=manifest_path or None, 95 ) 96 elif override_execution_mode == "python": 97 source = get_source( 98 connector_name, 99 use_python=True, 100 install_if_missing=False, 101 source_manifest=manifest_path or None, 102 ) 103 elif override_execution_mode == "docker": 104 source = get_source( 105 connector_name, 106 docker_image=True, 107 install_if_missing=False, 108 source_manifest=manifest_path or None, 109 ) 110 elif override_execution_mode == "yaml": 111 source = get_source( 112 connector_name, 113 source_manifest=manifest_path or True, 114 install_if_missing=False, 115 ) 116 else: 117 raise ValueError( 118 f"Unknown execution method: {override_execution_mode}. " 119 "Expected one of: ['auto', 'docker', 'python', 'yaml']." 120 ) 121 122 # Ensure installed: 123 if install_if_missing: 124 source.executor.ensure_installation() 125 126 return source 127 128 129@mcp_tool( 130 read_only=True, 131 idempotent=True, 132 requires_client_filesystem=True, 133 extra_help_text=_CONFIG_HELP, 134) 135def validate_connector_config( 136 connector_name: Annotated[ 137 str, 138 Field(description="The name of the connector to validate."), 139 ], 140 config: Annotated[ 141 dict | str | None, 142 Field( 143 description="The configuration for the connector as a dict object or JSON string.", 144 default=None, 145 ), 146 ], 147 config_file: Annotated[ 148 str | Path | None, 149 Field( 150 description="Path to a YAML or JSON file containing the connector configuration.", 151 default=None, 152 ), 153 ], 154 config_secret_name: Annotated[ 155 str | None, 156 Field( 157 description="The name of the secret containing the configuration.", 158 default=None, 159 ), 160 ], 161 override_execution_mode: Annotated[ 162 Literal["docker", "python", "yaml", "auto"], 163 Field( 164 description="Optionally override the execution method to use for the connector. " 165 "This parameter is ignored if manifest_path is provided (yaml mode will be used).", 166 default="auto", 167 ), 168 ], 169 manifest_path: Annotated[ 170 str | Path | None, 171 Field( 172 description="Path to a local YAML manifest file for declarative connectors.", 173 default=None, 174 ), 175 ], 176) -> tuple[bool, str]: 177 """Validate a connector configuration. 178 179 Returns a tuple of (is_valid: bool, message: str). 180 """ 181 try: 182 source: Source = _get_mcp_source( 183 connector_name, 184 override_execution_mode=override_execution_mode, 185 manifest_path=manifest_path, 186 ) 187 except Exception as ex: 188 return False, f"Failed to get connector '{connector_name}': {ex}" 189 190 try: 191 config_dict = resolve_connector_config( 192 config=config, 193 config_file=config_file, 194 config_secret_name=config_secret_name, 195 config_spec_jsonschema=source.config_spec, 196 ) 197 source.set_config(config_dict) 198 except Exception as ex: 199 return False, f"Failed to resolve configuration for {connector_name}: {ex}" 200 201 try: 202 source.check() 203 except Exception as ex: 204 return False, f"Configuration for {connector_name} is invalid: {ex}" 205 206 return True, f"Configuration for {connector_name} is valid!" 207 208 209@mcp_tool( 210 read_only=True, 211 idempotent=True, 212 requires_client_filesystem=True, 213) 214def list_connector_config_secrets( 215 connector_name: Annotated[ 216 str, 217 Field(description="The name of the connector."), 218 ], 219) -> list[str]: 220 """List all `config_secret_name` options that are known for the given connector. 221 222 This can be used to find out which already-created config secret names are available 223 for a given connector. The return value is a list of secret names, but it will not 224 return the actual secret values. 225 """ 226 raise_if_untrusted_execution_context( 227 "Listing connector config secrets (`list_connector_config_secrets`)" 228 ) 229 secrets_names: list[str] = [] 230 for secrets_mgr in _get_secret_sources(): 231 if isinstance(secrets_mgr, GoogleGSMSecretManager): 232 secrets_names.extend( 233 [ 234 secret_handle.secret_name.split("/")[-1] 235 for secret_handle in secrets_mgr.fetch_connector_secrets(connector_name) 236 ] 237 ) 238 239 return secrets_names 240 241 242@mcp_tool( 243 read_only=True, 244 idempotent=True, 245 requires_client_filesystem=True, 246 extra_help_text=_CONFIG_HELP, 247) 248def list_dotenv_secrets() -> dict[str, list[str]]: 249 """List all environment variable names declared within declared .env files. 250 251 This returns a dictionary mapping the .env file name to a list of environment 252 variable names. The values of the environment variables are not returned. 253 """ 254 raise_if_untrusted_execution_context("Listing dotenv secret names (`list_dotenv_secrets`)") 255 result: dict[str, list[str]] = {} 256 for secrets_mgr in _get_secret_sources(): 257 if isinstance(secrets_mgr, DotenvSecretManager) and secrets_mgr.dotenv_path: 258 result[str(secrets_mgr.dotenv_path.resolve())] = secrets_mgr.list_secrets_names() 259 260 return result 261 262 263@mcp_tool( 264 read_only=True, 265 idempotent=True, 266 requires_client_filesystem=True, 267 extra_help_text=_CONFIG_HELP, 268) 269def list_source_streams( 270 source_connector_name: Annotated[ 271 str, 272 Field(description="The name of the source connector."), 273 ], 274 config: Annotated[ 275 dict | str | None, 276 Field( 277 description="The configuration for the source connector as a dict or JSON string.", 278 default=None, 279 ), 280 ], 281 config_file: Annotated[ 282 str | Path | None, 283 Field( 284 description="Path to a YAML or JSON file containing the source connector config.", 285 default=None, 286 ), 287 ], 288 config_secret_name: Annotated[ 289 str | None, 290 Field( 291 description="The name of the secret containing the configuration.", 292 default=None, 293 ), 294 ], 295 override_execution_mode: Annotated[ 296 Literal["docker", "python", "yaml", "auto"], 297 Field( 298 description="Optionally override the execution method to use for the connector. " 299 "This parameter is ignored if manifest_path is provided (yaml mode will be used).", 300 default="auto", 301 ), 302 ], 303 manifest_path: Annotated[ 304 str | Path | None, 305 Field( 306 description="Path to a local YAML manifest file for declarative connectors.", 307 default=None, 308 ), 309 ], 310) -> list[str]: 311 """List all streams available in a source connector. 312 313 This operation (generally) requires a valid configuration, including any required secrets. 314 """ 315 source: Source = _get_mcp_source( 316 connector_name=source_connector_name, 317 override_execution_mode=override_execution_mode, 318 manifest_path=manifest_path, 319 ) 320 config_dict = resolve_connector_config( 321 config=config, 322 config_file=config_file, 323 config_secret_name=config_secret_name, 324 config_spec_jsonschema=source.config_spec, 325 ) 326 source.set_config(config_dict) 327 return source.get_available_streams() 328 329 330@mcp_tool( 331 read_only=True, 332 idempotent=True, 333 requires_client_filesystem=True, 334 extra_help_text=_CONFIG_HELP, 335) 336def get_source_stream_json_schema( 337 source_connector_name: Annotated[ 338 str, 339 Field(description="The name of the source connector."), 340 ], 341 stream_name: Annotated[ 342 str, 343 Field(description="The name of the stream."), 344 ], 345 config: Annotated[ 346 dict | str | None, 347 Field( 348 description="The configuration for the source connector as a dict or JSON string.", 349 default=None, 350 ), 351 ], 352 config_file: Annotated[ 353 str | Path | None, 354 Field( 355 description="Path to a YAML or JSON file containing the source connector config.", 356 default=None, 357 ), 358 ], 359 config_secret_name: Annotated[ 360 str | None, 361 Field( 362 description="The name of the secret containing the configuration.", 363 default=None, 364 ), 365 ], 366 override_execution_mode: Annotated[ 367 Literal["docker", "python", "yaml", "auto"], 368 Field( 369 description="Optionally override the execution method to use for the connector. " 370 "This parameter is ignored if manifest_path is provided (yaml mode will be used).", 371 default="auto", 372 ), 373 ], 374 manifest_path: Annotated[ 375 str | Path | None, 376 Field( 377 description="Path to a local YAML manifest file for declarative connectors.", 378 default=None, 379 ), 380 ], 381) -> dict[str, Any]: 382 """List all properties for a specific stream in a source connector.""" 383 source: Source = _get_mcp_source( 384 connector_name=source_connector_name, 385 override_execution_mode=override_execution_mode, 386 manifest_path=manifest_path, 387 ) 388 config_dict = resolve_connector_config( 389 config=config, 390 config_file=config_file, 391 config_secret_name=config_secret_name, 392 config_spec_jsonschema=source.config_spec, 393 ) 394 source.set_config(config_dict) 395 return source.get_stream_json_schema(stream_name=stream_name) 396 397 398@mcp_tool( 399 read_only=True, 400 requires_client_filesystem=True, 401 extra_help_text=_CONFIG_HELP, 402) 403def read_source_stream_records( 404 source_connector_name: Annotated[ 405 str, 406 Field(description="The name of the source connector."), 407 ], 408 config: Annotated[ 409 dict | str | None, 410 Field( 411 description="The configuration for the source connector as a dict or JSON string.", 412 default=None, 413 ), 414 ], 415 config_file: Annotated[ 416 str | Path | None, 417 Field( 418 description="Path to a YAML or JSON file containing the source connector config.", 419 default=None, 420 ), 421 ], 422 config_secret_name: Annotated[ 423 str | None, 424 Field( 425 description="The name of the secret containing the configuration.", 426 default=None, 427 ), 428 ], 429 *, 430 stream_name: Annotated[ 431 str, 432 Field(description="The name of the stream to read records from."), 433 ], 434 max_records: Annotated[ 435 int, 436 Field( 437 description="The maximum number of records to read.", 438 default=1000, 439 ), 440 ], 441 override_execution_mode: Annotated[ 442 Literal["docker", "python", "yaml", "auto"], 443 Field( 444 description="Optionally override the execution method to use for the connector. " 445 "This parameter is ignored if manifest_path is provided (yaml mode will be used).", 446 default="auto", 447 ), 448 ], 449 manifest_path: Annotated[ 450 str | Path | None, 451 Field( 452 description="Path to a local YAML manifest file for declarative connectors.", 453 default=None, 454 ), 455 ], 456) -> list[dict[str, Any]] | str: 457 """Get records from a source connector.""" 458 try: 459 source: Source = _get_mcp_source( 460 connector_name=source_connector_name, 461 override_execution_mode=override_execution_mode, 462 manifest_path=manifest_path, 463 ) 464 config_dict = resolve_connector_config( 465 config=config, 466 config_file=config_file, 467 config_secret_name=config_secret_name, 468 config_spec_jsonschema=source.config_spec, 469 ) 470 source.set_config(config_dict) 471 # First we get a generator for the records in the specified stream. 472 record_generator = source.get_records(stream_name) 473 # Next we load a limited number of records from the generator into our list. 474 records: list[dict[str, Any]] = list(islice(record_generator, max_records)) 475 476 print(f"Retrieved {len(records)} records from stream '{stream_name}'", sys.stderr) 477 478 except Exception as ex: 479 tb_str = traceback.format_exc() 480 # If any error occurs, we print the error message to stderr and return an empty list. 481 return ( 482 f"Error reading records from source '{source_connector_name}': {ex!r}, {ex!s}\n{tb_str}" 483 ) 484 485 else: 486 return records 487 488 489@mcp_tool( 490 read_only=True, 491 requires_client_filesystem=True, 492 extra_help_text=_CONFIG_HELP, 493) 494def get_stream_previews( 495 source_name: Annotated[ 496 str, 497 Field(description="The name of the source connector."), 498 ], 499 config: Annotated[ 500 dict | str | None, 501 Field( 502 description="The configuration for the source connector as a dict or JSON string.", 503 default=None, 504 ), 505 ], 506 config_file: Annotated[ 507 str | Path | None, 508 Field( 509 description="Path to a YAML or JSON file containing the source connector config.", 510 default=None, 511 ), 512 ], 513 config_secret_name: Annotated[ 514 str | None, 515 Field( 516 description="The name of the secret containing the configuration.", 517 default=None, 518 ), 519 ], 520 streams: Annotated[ 521 list[str] | str | None, 522 Field( 523 description=( 524 "The streams to get previews for. " 525 "Use '*' for all streams, or None for selected streams." 526 ), 527 default=None, 528 ), 529 ], 530 limit: Annotated[ 531 int, 532 Field( 533 description="The maximum number of sample records to return per stream.", 534 default=10, 535 ), 536 ], 537 override_execution_mode: Annotated[ 538 Literal["docker", "python", "yaml", "auto"], 539 Field( 540 description="Optionally override the execution method to use for the connector. " 541 "This parameter is ignored if manifest_path is provided (yaml mode will be used).", 542 default="auto", 543 ), 544 ], 545 manifest_path: Annotated[ 546 str | Path | None, 547 Field( 548 description="Path to a local YAML manifest file for declarative connectors.", 549 default=None, 550 ), 551 ], 552) -> dict[str, list[dict[str, Any]] | str]: 553 """Get sample records (previews) from streams in a source connector. 554 555 This operation requires a valid configuration, including any required secrets. 556 Returns a dictionary mapping stream names to lists of sample records, or an error 557 message string if an error occurred for that stream. 558 """ 559 source: Source = _get_mcp_source( 560 connector_name=source_name, 561 override_execution_mode=override_execution_mode, 562 manifest_path=manifest_path, 563 ) 564 565 config_dict = resolve_connector_config( 566 config=config, 567 config_file=config_file, 568 config_secret_name=config_secret_name, 569 config_spec_jsonschema=source.config_spec, 570 ) 571 source.set_config(config_dict) 572 573 streams_param: list[str] | Literal["*"] | None = resolve_list_of_strings( 574 streams 575 ) # pyrefly: ignore[no-matching-overload] 576 if streams_param and len(streams_param) == 1 and streams_param[0] == "*": 577 streams_param = "*" 578 579 try: 580 samples_result = source.get_samples( 581 streams=streams_param, 582 limit=limit, 583 on_error="ignore", 584 ) 585 except Exception as ex: 586 tb_str = traceback.format_exc() 587 return { 588 "ERROR": f"Error getting stream previews from source '{source_name}': " 589 f"{ex!r}, {ex!s}\n{tb_str}" 590 } 591 592 result: dict[str, list[dict[str, Any]] | str] = {} 593 for stream_name, dataset in samples_result.items(): 594 if dataset is None: 595 result[stream_name] = f"Could not retrieve stream samples for stream '{stream_name}'" 596 else: 597 result[stream_name] = list(dataset) 598 599 return result 600 601 602@mcp_tool( 603 destructive=False, 604 requires_client_filesystem=True, 605 extra_help_text=_CONFIG_HELP, 606) 607def sync_source_to_cache( 608 source_connector_name: Annotated[ 609 str, 610 Field(description="The name of the source connector."), 611 ], 612 config: Annotated[ 613 dict | str | None, 614 Field( 615 description="The configuration for the source connector as a dict or JSON string.", 616 default=None, 617 ), 618 ], 619 config_file: Annotated[ 620 str | Path | None, 621 Field( 622 description="Path to a YAML or JSON file containing the source connector config.", 623 default=None, 624 ), 625 ], 626 config_secret_name: Annotated[ 627 str | None, 628 Field( 629 description="The name of the secret containing the configuration.", 630 default=None, 631 ), 632 ], 633 streams: Annotated[ 634 list[str] | str, 635 Field( 636 description="The streams to sync.", 637 default="suggested", 638 ), 639 ], 640 override_execution_mode: Annotated[ 641 Literal["docker", "python", "yaml", "auto"], 642 Field( 643 description="Optionally override the execution method to use for the connector. " 644 "This parameter is ignored if manifest_path is provided (yaml mode will be used).", 645 default="auto", 646 ), 647 ], 648 manifest_path: Annotated[ 649 str | Path | None, 650 Field( 651 description="Path to a local YAML manifest file for declarative connectors.", 652 default=None, 653 ), 654 ], 655) -> str: 656 """Run a sync from a source connector to the default DuckDB cache.""" 657 source: Source = _get_mcp_source( 658 connector_name=source_connector_name, 659 override_execution_mode=override_execution_mode, 660 manifest_path=manifest_path, 661 ) 662 config_dict = resolve_connector_config( 663 config=config, 664 config_file=config_file, 665 config_secret_name=config_secret_name, 666 config_spec_jsonschema=source.config_spec, 667 ) 668 source.set_config(config_dict) 669 cache = get_default_cache() 670 671 streams = resolve_list_of_strings(streams) 672 if streams and len(streams) == 1 and streams[0] in {"*", "suggested"}: 673 # Float '*' and 'suggested' to the top-level for special processing: 674 streams = streams[0] 675 676 if isinstance(streams, str) and streams == "suggested": 677 streams = "*" # Default to all streams if 'suggested' is not otherwise specified. 678 try: 679 metadata = get_connector_metadata( 680 source_connector_name, 681 ) 682 except Exception: 683 streams = "*" # Fallback to all streams if suggested streams fail. 684 else: 685 if metadata is not None: 686 streams = metadata.suggested_streams or "*" 687 688 if isinstance(streams, str) and streams != "*": 689 streams = [streams] # Ensure streams is a list 690 691 source.read( 692 cache=cache, 693 streams=streams, 694 ) 695 del cache # Ensure the cache is closed properly 696 697 summary: str = f"Sync completed for '{source_connector_name}'!\n\n" 698 summary += "Data written to default DuckDB cache\n" 699 return summary 700 701 702class CachedDatasetInfo(BaseModel): 703 """Class to hold information about a cached dataset.""" 704 705 stream_name: str 706 """The name of the stream in the cache.""" 707 table_name: str 708 schema_name: str | None = None 709 710 711@mcp_tool( 712 read_only=True, 713 idempotent=True, 714 requires_client_filesystem=True, 715 extra_help_text=_CONFIG_HELP, 716) 717def list_cached_streams() -> list[CachedDatasetInfo]: 718 """List all streams available in the default DuckDB cache.""" 719 raise_if_untrusted_execution_context("Reading the local default cache (`list_cached_streams`)") 720 cache: DuckDBCache = get_default_cache() 721 result = [ 722 CachedDatasetInfo( 723 stream_name=stream_name, 724 table_name=(cache.table_prefix or "") + stream_name, 725 schema_name=cache.schema_name, 726 ) 727 for stream_name in cache.streams 728 ] 729 del cache # Ensure the cache is closed properly 730 return result 731 732 733@mcp_tool( 734 read_only=True, 735 idempotent=True, 736 requires_client_filesystem=True, 737 extra_help_text=_CONFIG_HELP, 738) 739def describe_default_cache() -> dict[str, Any]: 740 """Describe the currently configured default cache.""" 741 raise_if_untrusted_execution_context( 742 "Describing the local default cache (`describe_default_cache`)" 743 ) 744 cache = get_default_cache() 745 return { 746 "cache_type": type(cache).__name__, 747 "cache_dir": str(cache.cache_dir), 748 "cache_db_path": str(Path(cache.db_path).absolute()), 749 "cached_streams": list(cache.streams.keys()), 750 } 751 752 753def _is_safe_sql(sql_query: str) -> bool: 754 """Check if a SQL query is safe to execute. 755 756 For security reasons, we only allow read-only operations like SELECT, DESCRIBE, and SHOW. 757 Multi-statement queries (containing semicolons) are also disallowed for security. 758 759 Note: SQLAlchemy will also validate downstream, but this is a first-pass check. 760 761 Args: 762 sql_query: The SQL query to check 763 764 Returns: 765 True if the query is safe to execute, False otherwise 766 """ 767 # Remove leading/trailing whitespace and convert to uppercase for checking 768 normalized_query = sql_query.strip().upper() 769 770 # Disallow multi-statement queries (containing semicolons) 771 # Note: We check the original query to catch semicolons anywhere, including in comments 772 if ";" in sql_query: 773 return False 774 775 # List of allowed SQL statement prefixes (read-only operations) 776 allowed_prefixes = ( 777 "SELECT", 778 "DESCRIBE", 779 "DESC", # Short form of DESCRIBE 780 "SHOW", 781 "EXPLAIN", # Also safe - shows query execution plan 782 ) 783 784 # Check if the query starts with any allowed prefix 785 return any(normalized_query.startswith(prefix) for prefix in allowed_prefixes) 786 787 788@mcp_tool( 789 read_only=True, 790 idempotent=True, 791 requires_client_filesystem=True, 792 extra_help_text=_CONFIG_HELP, 793) 794def run_sql_query( 795 sql_query: Annotated[ 796 str, 797 Field(description="The SQL query to execute."), 798 ], 799 max_records: Annotated[ 800 int, 801 Field( 802 description="Maximum number of records to return.", 803 default=1000, 804 ), 805 ], 806) -> list[dict[str, Any]]: 807 """Run a SQL query against the default cache. 808 809 The dialect of SQL should match the dialect of the default cache. 810 Use `describe_default_cache` to see the cache type. 811 812 For DuckDB-type caches: 813 - Use `SHOW TABLES` to list all tables. 814 - Use `DESCRIBE <table_name>` to get the schema of a specific table 815 816 For security reasons, only read-only operations are allowed: SELECT, DESCRIBE, SHOW, EXPLAIN. 817 """ 818 raise_if_untrusted_execution_context("Querying the local default cache (`run_sql_query`)") 819 # Check if the query is safe to execute 820 if not _is_safe_sql(sql_query): 821 return [ 822 { 823 "ERROR": "Unsafe SQL query detected. Only read-only operations are allowed: " 824 "SELECT, DESCRIBE, SHOW, EXPLAIN", 825 "SQL_QUERY": sql_query, 826 } 827 ] 828 829 cache: DuckDBCache = get_default_cache() 830 try: 831 return cache.run_sql_query( 832 sql_query, 833 max_records=max_records, 834 ) 835 except Exception as ex: 836 tb_str = traceback.format_exc() 837 return [ 838 { 839 "ERROR": f"Error running SQL query: {ex!r}, {ex!s}", 840 "TRACEBACK": tb_str, 841 "SQL_QUERY": sql_query, 842 } 843 ] 844 finally: 845 del cache # Ensure the cache is closed properly 846 847 848@mcp_tool( 849 destructive=True, 850 requires_client_filesystem=True, 851) 852def destination_smoke_test( # noqa: PLR0913, PLR0917 853 destination_connector_name: Annotated[ 854 str, 855 Field( 856 description=( 857 "The name of the destination connector to test " 858 "(e.g. 'destination-snowflake', 'destination-motherduck')." 859 ), 860 ), 861 ], 862 config: Annotated[ 863 dict | str | None, 864 Field( 865 description=( 866 "The destination configuration as a dict object or JSON string. " 867 "Must not contain hardcoded secrets; use secret_reference::ENV_VAR_NAME instead." 868 ), 869 default=None, 870 ), 871 ], 872 config_file: Annotated[ 873 str | Path | None, 874 Field( 875 description="Path to a YAML or JSON file containing the destination configuration.", 876 default=None, 877 ), 878 ], 879 config_secret_name: Annotated[ 880 str | None, 881 Field( 882 description="The name of the secret containing the destination configuration.", 883 default=None, 884 ), 885 ], 886 scenarios: Annotated[ 887 list[str] | str, 888 Field( 889 description=( 890 "Which scenarios to run. Use 'fast' (default) for all fast predefined " 891 "scenarios (excludes large_batch_stream), 'all' for every predefined " 892 "scenario including large batch, or provide a list of scenario names " 893 "or a comma-separated string." 894 ), 895 default="fast", 896 ), 897 ], 898 custom_scenarios: Annotated[ 899 list[dict[str, Any]] | None, 900 Field( 901 description=( 902 "Additional custom test scenarios to inject. Each scenario should define " 903 "'name', 'json_schema', and optionally 'records' and 'primary_key'. " 904 "These are unioned with the predefined scenarios." 905 ), 906 default=None, 907 ), 908 ], 909 docker_image: Annotated[ 910 str | None, 911 Field( 912 description=( 913 "Optional Docker image override for the destination connector " 914 "(e.g. 'airbyte/destination-snowflake:3.14.0')." 915 ), 916 default=None, 917 ), 918 ], 919 namespace_suffix: Annotated[ 920 str | None, 921 Field( 922 description=( 923 "Optional suffix appended to the auto-generated namespace. " 924 "Defaults to 'smoke_test' (format: 'zz_deleteme_yyyymmdd_hhmm_{suffix}'). " 925 "Use this to distinguish concurrent runs." 926 ), 927 default=None, 928 ), 929 ], 930 reuse_namespace: Annotated[ 931 str | None, 932 Field( 933 description=( 934 "Exact namespace to reuse from a previous run. " 935 "When set, no new namespace is generated. " 936 "Useful for running a second test against an already-populated namespace." 937 ), 938 default=None, 939 ), 940 ], 941 skip_preflight: Annotated[ 942 bool, 943 Field( 944 description=( 945 "Skip the automatic preflight check that runs basic_types before " 946 "the requested scenarios. Set to true when you expect basic_types " 947 "itself to fail or want to save time on repeated runs." 948 ), 949 default=False, 950 ), 951 ], 952) -> DestinationSmokeTestResult: 953 """Run smoke tests against a destination connector. 954 955 Sends synthetic test data from the smoke test source to the specified 956 destination and reports success or failure. The smoke test source generates 957 data across predefined scenarios covering common destination failure patterns: 958 type variations, null handling, naming edge cases, schema variations, and 959 batch sizes. 960 961 When the destination has a compatible cache implementation (DuckDB, 962 Postgres, Snowflake, BigQuery, MotherDuck), readback introspection is 963 automatically performed after a successful write. The readback produces 964 stats on the written data: table row counts, column names/types, and 965 per-column null/non-null counts. Results are included in the response 966 as `table_statistics` and `tables_not_found`. 967 """ 968 raise_if_untrusted_execution_context("Destination smoke test (`destination_smoke_test`)") 969 # Resolve destination config 970 config_dict = resolve_connector_config( 971 config=config, 972 config_file=config_file, 973 config_secret_name=config_secret_name, 974 ) 975 976 # Set up destination 977 destination_kwargs: dict[str, Any] = { 978 "name": destination_connector_name, 979 "config": config_dict, 980 } 981 if docker_image: 982 destination_kwargs["docker_image"] = docker_image 983 elif is_docker_installed(): 984 destination_kwargs["docker_image"] = True 985 986 destination_obj = get_destination(**destination_kwargs) 987 988 # Resolve scenarios for the shared helper 989 resolved_scenarios: str | list[str] 990 if isinstance(scenarios, str): 991 resolved_scenarios = scenarios 992 else: 993 resolved_scenarios = resolve_list_of_strings(scenarios) or "fast" 994 995 return run_destination_smoke_test( 996 destination=destination_obj, 997 scenarios=resolved_scenarios, 998 namespace_suffix=namespace_suffix, 999 reuse_namespace=reuse_namespace, 1000 custom_scenarios=custom_scenarios, 1001 skip_preflight=skip_preflight, 1002 ) 1003 1004 1005def register_local_tools(app: FastMCP) -> None: 1006 """Register local tools with the FastMCP app. 1007 1008 Args: 1009 app: FastMCP application instance 1010 """ 1011 register_mcp_tools(app, mcp_module=__name__)