Feature Suggestion: Plugin system for interacting with loaded SQLMesh Contexts

[Harmuth94]

Feature Suggestion: Plugin system for interacting with loaded SQLMesh Contexts

Background

In our team's use of SQLMesh, we've encountered scenarios where we needed to interact with or manipulate the loaded context dynamically. To address these needs, we implemented a lightweight extension system that introduces a plugin mechanism for custom context manipulation.

This issue outlines our approach and proposes introducing a first-class plugin system within SQLMesh to streamline similar use cases for other teams.

What are plugins?

Plugins are modular components that extend the functionality of an application without requiring modifications to its core. In the context of SQLMesh, plugins could allow users to:

• Perform custom validations or transformations on the context.
• Inject additional metadata or configurations.
• Automate repetitive or domain-specific tasks.

Our Implementation

We created a thin wrapper around the SQLMesh CLI to enable plugin functionality. Here's the directory structure of our implementation:

sqlmeshcliextended/
├── README.md
├── __init__.py
├── cli/
│   ├── __init__.py
│   └── main.py
└── plugins/
    ├── __init__.py
    └── <plugin_name>.py

Key File: cli/main.py

We extended the SQLMesh CLI by wrapping it in our custom CLI with additional context manipulation capabilities.

from __future__ import annotations
import logging
import click
from sqlmesh import __version__
from sqlmesh.cli import error_handler
from sqlmesh.cli.main import cli as sqlmesh_cli
from plugins.plugins import plugins

logger = logging.getLogger(__name__)

@click.group(no_args_is_help=True)
@click.version_option(version=__version__, message="%(version)s")
@click.pass_context
@error_handler
def cli(ctx: click.Context, **kwargs) -> None:
    """Extended SQLMesh command line tool."""
    sqlmesh_cli.callback(**kwargs)
    context = ctx.obj

    # Apply plugins to the context
    extended_context = plugins(context)
    ctx.obj = extended_context

# Add all existing SQLMesh commands and options to the extended CLI
for name, cmd in sqlmesh_cli.commands.items():
    cli.add_command(cmd, name=name)

for opt in sqlmesh_cli.params:
    cli.params.append(opt)

Plugins Functionality

The core of the our plugin system resides in the plugins.plugins function. This function enables us to validate and manipulate the loaded SQLMesh context before it is passed downstream.

def plugins(context: Context) -> Context:
    """Apply plugins to enhance the SQLMesh context."""
    # Example validations
    context = static_schema_validation(context)
    context = assert_expectations_have_tags(context)

    # Example manipulations
    context = store_model_tags(context)
    context = ensure_rw_has_materialized_view(context)
    return context

Proposal

We propose introducing a native plugin system into SQLMesh. This system would:

1. Allow users to register custom plugins easily.
2. Provide hooks for validations, context transformations, or other pre-defined events.
3. Simplify the extension of SQLMesh functionality without requiring custom wrappers.

Benefits

• Flexibility: Teams can tailor SQLMesh to their specific workflows and domain requirements.
• Reusability: Plugins can be modular and shared across projects or teams.
• Ease of Use: A native plugin system reduces the need for custom CLI wrappers, ensuring compatibility with future SQLMesh updates.

Suggested Implementation

1. Add a plugins folder to the SQLMesh project to house custom plugins
2. Apply plugins in the GenericContext.load function after context has been loaded