feat: track rows processed during model evaluation

[treysp]

Captures the number of rows processed (and potentially other information) during a model evaluation for display to the user.

Example running sushi on Bigquery (includes bytes processed for each model batch):

Context

• Python client libraries conformant to the DB-API 2.0 spec provide a cursor.row_count property that contains the number of rows affected by a DML command execution (e.g., INSERT, CTAS)
• Variations
  • Some clients provide the property, but they don't actually track rows so always return -1
  • Snowflake conforms other than CTAS
  • Bigquery python API necessary to handle all scenarios

Implementation context

• SQLMesh makes many execute calls that should not be tracked
• In most cases, SQLMesh DML actions take place via the engine adapter _execute method, so we read the row_count property there immediately after cursor.execute
• Snapshot evaluation may take place concurrently, where a single snapshot evaluation occurs within one thread

Determining what to track

• We should only track specific operations, so execute does NOT track rows by default
• Some engine adapter methods should always track, so they pass track_row_count=True in the execute call within them
• Other engine methods vary in their need to track so expose a track_row_count arg
  • Example: engine_adapter.insert_append
    • Should almost always be tracked during snapshot evaluation
    • Should never be tracked when called from state sync

Tracker implementation

• Row count values are collected in the QueryExecutionContext dataclass
  • Accumulates row counts if called multiple times
  • Ignores row counts of -1 (default value when cursor.row_count not populated)
• Trackers implemented as context manager
• Engine adapter _execute calls single record_execution method, which determines which tracker is active and should be used (if any)

[erindru]

Nice start, I think I follow the logic. I've left some comments for your consideration.

I'm not sure about the track_row_count parameter on the engine adapter methods. I'd like to understand more about the undesirable queries that were getting picked up - perhaps that indicates we were being too coarse with the statistics gathering / wrapping in the wrong place

[eakmanrq]

Building on @erindru's feedback, I wonder if in general this tracking logic could get pushed to the SnapshotEvaluator layer. Right now the engine adapter methods return nothing but if they returned row count then it seems like the SnapshotEvaluator could then determine what to do with that information.

I don't think though I deeply understand this problem and this comes from a quick review and my general intuition.

[erindru]

Right now the engine adapter methods return nothing but if they returned row count then it seems like the SnapshotEvaluator could then determine what to do with that information.

If we pass row counts around through the entire call stack between the scheduler and the engine adapter, I believe it starts getting very messy with the number of method signature changes.

The scheduler is where we want to display the information, but it's typically on a different thread and reasonably far removed from EngineAdapter._execute where the relevant queries are executed. In addition, the EngineAdapter layer is a driver that deliberately has no idea what Snapshot it's executing because that's a higher level concept.

This design that Trey is proposing means that we don't have to change the return type of the engine adapter methods or have to worry about passing things back up the stack / rely on every component in the chain between the Scheduler and EngineAdapter._execute to pass things back correctly.

It essentially uses a combination of thread-local + global state to bypass the call stack. It emits row counts in EngineAdapter._execute, assigns meaning to them in SnapshotEvaluator and allows the information to be queried back in Scheduler so it can be outputted to the console for the user to see - without making every method in between return some kind of Stats object.

imo this is a much simpler / less invasive implementation