Add support for ignoring queries that violate an index requirement based on callstack

Author: muglug

src/AsyncMysql/AsyncMysqlConnection.php

@@ -43,6 +43,7 @@ public function __construct(private string $host, private int $port, private str
 		dict<string, string> $_query_attributes = dict[],
 	): Awaitable<AsyncMysqlQueryResult> {
 		Logger::log(Verbosity::QUERIES, "SQLFake [verbose]: $query");
+		QueryContext::$query = $query;
 
 		$config = $this->server->config;
 		$strict_sql_before = QueryContext::$strictSQLMode;
@@ -62,11 +63,13 @@ public function __construct(private string $host, private int $port, private str
 		try {
 			list($results, $rows_affected) = SQLCommandProcessor::execute($query, $this);
 		} catch (\Exception $e) {
-			// this makes debugging a failing unit test easier, show the actual query that failed parsing along with the parser error
 			QueryContext::$strictSQLMode = $strict_sql_before;
 			QueryContext::$strictSchemaMode = $strict_schema_before;
+			// Make debugging a failing unit test locally easier,
+			// by showing the actual query that failed parsing along with the parser error
 			$msg = $e->getMessage();
 			$type = \get_class($e);
+			Logger::log(Verbosity::QUIET, $e->getFile().' '.$e->getLine());
 			Logger::log(Verbosity::QUIET, "SQL Fake $type: $msg in SQL query: $query");
 			throw $e;
 		}
@@ -82,7 +85,10 @@ public function __construct(private string $host, private int $port, private str
 	}
 
 	<<__Override>>
-	public function queryf(\HH\FormatString<\HH\SQLFormatter> $query, mixed ...$args): Awaitable<AsyncMysqlQueryResult> {
+	public function queryf(
+		\HH\FormatString<\HH\SQLFormatter> $query,
+		mixed ...$args
+	): Awaitable<AsyncMysqlQueryResult> {
 		invariant($query is string, '\\HH\\FormatString<_> is opaque, but we need it to be a string here.');
 		return $this->query($this->queryStringifier->formatString($query, vec($args)));
 	}

src/Query/QueryPlanner.hack

@@ -161,7 +161,26 @@ abstract class QueryPlanner {
 					false,
 				);
 			} else if (QueryContext::$requireIndexes) {
-				throw new \Exception('Query without index: '.(QueryContext::$query ?? ''));
+				$stack = \debug_backtrace(\DEBUG_BACKTRACE_IGNORE_ARGS);
+				C\pop_front(inout $stack);
+
+				$grandfathered = false;
+				foreach ($stack as $stack_frame) {
+					$class = $stack_frame['class'] ?? null;
+					$function = $stack_frame['function'];
+
+					foreach (QueryContext::$allowed_index_violation_traces as $allowed_trace) {
+						$allowed_trace_class = $allowed_trace['class'] ?? null;
+
+						if ($class === $allowed_trace_class && $allowed_trace['function'] === $function) {
+							$grandfathered = true;
+						}
+					}
+				}
+
+				if (!$grandfathered) {
+					throw new \Exception('Query without index: '.(QueryContext::$query ?? ''));
+				}
 			}
 
 			return null;

src/QueryContext.php

@@ -29,6 +29,12 @@
 	 */
 	public static bool $requireIndexes = false;
 
+	/**
+	 * Allow violations of the must-use-index policy when the stack trace matches any of these
+	 * class-function pairs.
+	 */
+	public static vec<shape(?'class' => string, 'function' => string)> $allowed_index_violation_traces = vec[];
+
 	/**
 	 * Require the presence of the table's Vitess sharding key
 	 */