Add Vitess query validation support for update changes primary vindex…

… column (#49)

* Add Vitess query validation support for `update changes primary vindex column`

* linters

src/Query/UpdateQuery.php

@@ -9,14 +9,14 @@ public function __construct(public from_table $updateClause, public string $sql,
   public vec<BinaryOperatorExpression> $setClause = vec[];
 
   public function execute(AsyncMysqlConnection $conn): int {
-    list($table_name, $database, $data) = $this->processUpdateClause($conn);
-    Metrics::trackQuery(QueryType::UPDATE, $conn->getServer()->name, $table_name, $this->sql);
-    $schema = QueryContext::getSchema($database, $table_name);
+    list($tableName, $database, $data) = $this->processUpdateClause($conn);
+    Metrics::trackQuery(QueryType::UPDATE, $conn->getServer()->name, $tableName, $this->sql);
+    $schema = QueryContext::getSchema($database, $tableName);
 
     list($rows_affected, $_) = $this->applyWhere($conn, $data)
       |> $this->applyOrderBy($conn, $$)
       |> $this->applyLimit($$)
-      |> $this->applySet($conn, $database, $table_name, $$, $data, $this->setClause, $schema);
+      |> $this->applySet($conn, $database, $tableName, $$, $data, $this->setClause, $schema);
 
     return $rows_affected;
   }
@@ -26,8 +26,8 @@ public function execute(AsyncMysqlConnection $conn): int {
    * add a row identifier to each element in the result which we can later use to update the underlying table
    */
   protected function processUpdateClause(AsyncMysqlConnection $conn): (string, string, dataset) {
-    list($database, $table_name) = Query::parseTableName($conn, $this->updateClause['name']);
-    $table = $conn->getServer()->getTable($database, $table_name) ?? vec[];
-    return tuple($table_name, $database, $table);
+    list($database, $tableName) = Query::parseTableName($conn, $this->updateClause['name']);
+    $table = $conn->getServer()->getTable($database, $tableName) ?? vec[];
+    return tuple($tableName, $database, $table);
   }
 }

src/VitessQueryValidator.php

@@ -14,6 +14,7 @@ enum UnsupportedCases: string as string {
     GROUP_BY_COLUMNS = 'unsupported: in scatter query: group by column must reference column in SELECT list';
     ORDER_BY_COLUMNS = 'unsupported: in scatter query: order by column must reference column in SELECT list';
     UNIONS = 'unsupported: UNION cannot be executed as a single route';
+	PRIMARY_VINDEX_COLUMN = 'unsupported: update changes primary vindex column';
 }
 
 abstract class VitessQueryValidator {
@@ -36,20 +37,56 @@ abstract protected function getHandlers(): dict<string, (function(): Awaitable<v
     public static function validate(Query $query, AsyncMysqlConnection $conn): void {
         if ($query is SelectQuery) {
             \HH\Asio\join((new SelectQueryValidator($query, $conn))->processHandlers());
-        }
+		} else if ($query is UpdateQuery) {
+            /*HHAST_FIXME[DontUseAsioJoin]*/
+            \HH\Asio\join((new UpdateQueryValidator($query, $conn))->processHandlers());
+		}
     }
 
     public static function extractColumnExprNames(vec<Expression> $selectExpressions): keyset<string> {
         $exprNames = keyset[];
         foreach ($selectExpressions as $expr) {
             if ($expr is ColumnExpression) {
                 $exprNames[] = $expr->name;
-            }
+			} else if ($expr is BinaryOperatorExpression) {
+				$exprNames[] = $expr->left->name;
+		   	}
         }
         return $exprNames;
     }
 }
 
+final class UpdateQueryValidator extends VitessQueryValidator {
+    public function __construct(public UpdateQuery $query, public AsyncMysqlConnection $conn) {}
+
+    <<__Override>>
+    public function getHandlers(): dict<string, (function(): Awaitable<void>)> {
+		return dict[
+			UnsupportedCases::PRIMARY_VINDEX_COLUMN => inst_meth($this, 'updateChangesPrimaryVindexColumn')
+		];
+	}
+
+	public async function updateChangesPrimaryVindexColumn(): Awaitable<void> {
+    	$set = $this->query->setClause;
+
+    	list($database, $table_name) = Query::parseTableName($this->conn, $this->query->updateClause['name']);
+        $table_schema = QueryContext::getSchema($database, $table_name);
+        $vitess_sharding = $table_schema['vitess_sharding'] ?? null;
+
+		if ($vitess_sharding === null) {
+			throw new SQLFakeVitessQueryViolation(Str\format('Missing Vitess sharding information for: %s', $table_name));
+		}
+
+        $columns = VitessQueryValidator::extractColumnExprNames($set);
+
+        if (C\contains_key($columns, $vitess_sharding['sharding_key'])) {
+            throw new SQLFakeVitessQueryViolation(
+                Str\format('Vitess query validation error: %s', UnsupportedCases::PRIMARY_VINDEX_COLUMN),
+            );
+		}
+	}
+}
+
 final class SelectQueryValidator extends VitessQueryValidator {
 
     public function __construct(public SelectQuery $query, public AsyncMysqlConnection $conn) {}

tests/UpdateQueryValidatorTest.php

@@ -0,0 +1,50 @@
+<?hh // strict
+
+namespace Slack\SQLFake;
+
+use function Facebook\FBExpect\expect;
+use type Facebook\HackTest\HackTest;
+
+final class UpdateQueryValidatorTest extends HackTest {
+	private static ?AsyncMysqlConnection $conn;
+
+    <<__Override>>
+    public static async function beforeFirstTestAsync(): Awaitable<void> {
+        static::$conn = await SharedSetup::initVitessAsync();
+		// block hole logging
+		// ? copied from SelectQueryValidatorTest.php, not sure what that means.
+        Logger::setHandle(new \Facebook\CLILib\TestLib\StringOutput());
+    }
+
+    <<__Override>>
+    public async function beforeEachTestAsync(): Awaitable<void> {
+        restore('vitess_setup');
+        QueryContext::$strictSchemaMode = false;
+        QueryContext::$strictSQLMode = false;
+    }
+
+	public async function testUpdateChangesPrimaryVindex(): Awaitable<void> {
+        $conn = static::$conn as nonnull;
+
+		$unsupported_test_cases = vec[
+			'update vt_table1 set id=1 where id=1',
+			'update vt_table2 set vt_table1_id=1 where id=1'
+		];
+
+		foreach($unsupported_test_cases as $sql) {
+            expect(() ==> $conn->query($sql))->toThrow(
+                SQLFakeVitessQueryViolation::class,
+                'Vitess query validation error: unsupported: update changes primary vindex column',
+            );
+		}
+
+		$supported_test_cases = vec[
+			"update vt_table1 set name='foo' where id = 1"
+		];
+
+        foreach ($supported_test_cases as $sql) {
+            expect(() ==> $conn->query($sql))->notToThrow(SQLFakeVitessQueryViolation::class);
+        }
+
+	}
+}