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

[izeigerman]

Hey @Harmuth94 ! Thanks a lot for the proposal!

It seems you're fine with using the Context API and aren't looking to extend the inner workings of the context.

What seems to be missing for you today is the ability to manipulate the loaded Context instance while still using the existing SQLMesh CLI.

As far as I can tell, all the manipulations that you're looking to do concern the loaded model objects. We already have an extension mechanism within SQLMesh that can let you do things like that.

Unfortunately, it's not documented today, but here's a brief explanation of how it works.

Users can create a custom loader class by inheriting the SqlMeshLoader class. They can then override the _load_models method to inject any logic they need after the models were loaded eg:

class MyCustomLoaderClass(SqlMeshLoader):

    def _load_models(
        self,
        macros: MacroRegistry,
        jinja_macros: JinjaMacroRegistry,
        gateway: t.Optional[str],
        audits: UniqueKeyDict[str, ModelAudit],
        signals: UniqueKeyDict[str, signal],
    ) -> UniqueKeyDict[str, Model]:
        models = self._load_models(
            macros=maros,
            jinja_macros=jinja_macros,
            gateway=gateway,
            audits=audits,
            signals=signals,
       )
       # Modify models here
       return models

The last step is to pass the custom loader class into the SQLMesh configuration object:

from sqlmesh.config import Config

config = Config(
    ...,
    loader=MyCustomLoaderClass,
)

Presently, this is only supported when using the config.py configuration file and is not supported with YAML config. At the same time, the custom loader class can be defined directly in the config.py file.

Additionally, other SQLMesh objects like audits, macros, etc can be modified in a similar manner, by overriding the appropriate methods.

What do you think, @Harmuth94? Is this something that would work for your use cases?