Fix(mysql): use 8-character random ID for temp table names (#1589)

* Use 8-character random ID for temp table names

* Tidy tests

* Convert temp_table_name test function to fixture

* PR feedback

Author: treysp

Makefile

@@ -4,7 +4,7 @@ install-dev:
 	pip3 install -e ".[dev,web,slack]"
 
 install-engine-integration:
-	pip3 install -e ".[dev,web,slack,mysql,postgres,databricks,redshift,bigquery,snowflake]"
+	pip3 install -e ".[dev,web,slack,mysql,postgres,databricks,redshift,bigquery,snowflake,trino]"
 
 install-pre-commit:
 	pre-commit install

sqlmesh/core/engine_adapter/base.py

@@ -14,7 +14,6 @@
 import sys
 import types
 import typing as t
-import uuid
 from datetime import datetime, timezone
 from enum import Enum
 from functools import partial
@@ -34,7 +33,7 @@
 from sqlmesh.core.engine_adapter.shared import DataObject, set_catalog
 from sqlmesh.core.model.kind import TimeColumn
 from sqlmesh.core.schema_diff import SchemaDiffer
-from sqlmesh.utils import double_escape
+from sqlmesh.utils import double_escape, random_id
 from sqlmesh.utils.connection_pool import create_connection_pool
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_ts
 from sqlmesh.utils.errors import SQLMeshError, UnsupportedCatalogOperationError
@@ -1431,7 +1430,7 @@ def _get_temp_table(
         Returns the name of the temp table that should be used for the given table name.
         """
         table = t.cast(exp.Table, exp.to_table(table).copy())
-        table.set("this", exp.to_identifier(f"__temp_{table.name}_{uuid.uuid4().hex}"))
+        table.set("this", exp.to_identifier(f"__temp_{table.name}_{random_id(short=True)}"))
 
         if table_only:
             table.set("db", None)

sqlmesh/utils/__init__.py

@@ -3,7 +3,9 @@
 import copy
 import importlib
 import os
+import random
 import re
+import string
 import sys
 import time
 import traceback
@@ -22,6 +24,8 @@
 VALUE = t.TypeVar("VALUE")
 DECORATOR_RETURN_TYPE = t.TypeVar("DECORATOR_RETURN_TYPE")
 
+ALPHANUMERIC = string.ascii_lowercase + string.digits
+
 
 def optional_import(name: str) -> t.Optional[types.ModuleType]:
     """Optionally import a module.
@@ -47,7 +51,10 @@ def unique(iterable: t.Iterable[T], by: t.Callable[[T], t.Any] = lambda i: i) ->
     return list({by(i): None for i in iterable})
 
 
-def random_id() -> str:
+def random_id(short: bool = False) -> str:
+    if short:
+        return "".join(random.choices(ALPHANUMERIC, k=8))
+
     return uuid.uuid4().hex
 
 

tests/conftest.py

@@ -260,3 +260,13 @@ def delete_cache(project_paths: str | t.List[str]) -> None:
             rmtree(path + "/.cache")
         except FileNotFoundError:
             pass
+
+
+@pytest.fixture
+def make_temp_table_name(mocker: MockerFixture) -> t.Callable:
+    def _make_function(table_name: str, random_id: str) -> exp.Table:
+        temp_table = exp.to_table(table_name)
+        temp_table.set("this", exp.to_identifier(f"__temp_{temp_table.name}_{random_id}"))
+        return temp_table
+
+    return _make_function

tests/core/engine_adapter/test_bigquery.py

@@ -1,7 +1,6 @@
 # type: ignore
 import sys
 import typing as t
-import uuid
 
 import pandas as pd
 import pytest
@@ -44,19 +43,20 @@ def test_insert_overwrite_by_time_partition_query(
 
 
 def test_insert_overwrite_by_partition_query(
-    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
 ):
     adapter = make_mocked_engine_adapter(BigQueryEngineAdapter)
     execute_mock = mocker.patch(
         "sqlmesh.core.engine_adapter.bigquery.BigQueryEngineAdapter.execute"
     )
 
-    temp_table_uuid = uuid.uuid4()
-    uuid4_mock = mocker.patch("uuid.uuid4")
-    uuid4_mock.return_value = temp_table_uuid
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
+    table_name = "test_schema.test_table"
+    temp_table_id = "abcdefgh"
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
 
     adapter.insert_overwrite_by_partition(
-        "test_schema.test_table",
+        table_name,
         parse_one("SELECT a, ds FROM tbl"),
         partitioned_by=[
             d.parse_one("DATETIME_TRUNC(ds, MONTH)"),
@@ -70,15 +70,15 @@ def test_insert_overwrite_by_partition_query(
     sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == [
         "CREATE SCHEMA IF NOT EXISTS `test_schema`",
-        f"CREATE TABLE IF NOT EXISTS `test_schema`.`__temp_test_table_{temp_table_uuid.hex}` PARTITION BY DATETIME_TRUNC(`ds`, MONTH) AS SELECT `a`, `ds` FROM `tbl`",
-        f"DECLARE _sqlmesh_target_partitions_ ARRAY<DATETIME> DEFAULT (SELECT ARRAY_AGG(PARSE_DATETIME('%Y%m', partition_id)) FROM `test_schema`.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '__temp_test_table_{temp_table_uuid.hex}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__');",
-        f"MERGE INTO `test_schema`.`test_table` AS `__MERGE_TARGET__` USING (SELECT `a`, `ds` FROM (SELECT * FROM `test_schema`.`__temp_test_table_{temp_table_uuid.hex}`) AS `_subquery` WHERE DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`)) AS `__MERGE_SOURCE__` ON FALSE WHEN NOT MATCHED BY SOURCE AND DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`) THEN DELETE WHEN NOT MATCHED THEN INSERT (`a`, `ds`) VALUES (`a`, `ds`)",
-        f"DROP TABLE IF EXISTS `test_schema`.`__temp_test_table_{temp_table_uuid.hex}`",
+        f"CREATE TABLE IF NOT EXISTS `test_schema`.`__temp_test_table_{temp_table_id}` PARTITION BY DATETIME_TRUNC(`ds`, MONTH) AS SELECT `a`, `ds` FROM `tbl`",
+        f"DECLARE _sqlmesh_target_partitions_ ARRAY<DATETIME> DEFAULT (SELECT ARRAY_AGG(PARSE_DATETIME('%Y%m', partition_id)) FROM `test_schema`.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '__temp_test_table_{temp_table_id}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__');",
+        f"MERGE INTO `test_schema`.`test_table` AS `__MERGE_TARGET__` USING (SELECT `a`, `ds` FROM (SELECT * FROM `test_schema`.`__temp_test_table_{temp_table_id}`) AS `_subquery` WHERE DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`)) AS `__MERGE_SOURCE__` ON FALSE WHEN NOT MATCHED BY SOURCE AND DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`) THEN DELETE WHEN NOT MATCHED THEN INSERT (`a`, `ds`) VALUES (`a`, `ds`)",
+        f"DROP TABLE IF EXISTS `test_schema`.`__temp_test_table_{temp_table_id}`",
     ]
 
 
 def test_insert_overwrite_by_partition_query_unknown_column_types(
-    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
 ):
     adapter = make_mocked_engine_adapter(BigQueryEngineAdapter)
     execute_mock = mocker.patch(
@@ -93,12 +93,13 @@ def test_insert_overwrite_by_partition_query_unknown_column_types(
         "ds": exp.DataType.build("DATETIME"),
     }
 
-    temp_table_uuid = uuid.uuid4()
-    uuid4_mock = mocker.patch("uuid.uuid4")
-    uuid4_mock.return_value = temp_table_uuid
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
+    table_name = "test_schema.test_table"
+    temp_table_id = "abcdefgh"
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
 
     adapter.insert_overwrite_by_partition(
-        "test_schema.test_table",
+        table_name,
         parse_one("SELECT a, ds FROM tbl"),
         partitioned_by=[
             d.parse_one("DATETIME_TRUNC(ds, MONTH)"),
@@ -110,16 +111,16 @@ def test_insert_overwrite_by_partition_query_unknown_column_types(
     )
 
     columns_mock.assert_called_once_with(
-        exp.to_table(f"test_schema.__temp_test_table_{temp_table_uuid.hex}")
+        exp.to_table(f"test_schema.__temp_test_table_{temp_table_id}")
     )
 
     sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == [
         "CREATE SCHEMA IF NOT EXISTS `test_schema`",
-        f"CREATE TABLE IF NOT EXISTS `test_schema`.`__temp_test_table_{temp_table_uuid.hex}` PARTITION BY DATETIME_TRUNC(`ds`, MONTH) AS SELECT `a`, `ds` FROM `tbl`",
-        f"DECLARE _sqlmesh_target_partitions_ ARRAY<DATETIME> DEFAULT (SELECT ARRAY_AGG(PARSE_DATETIME('%Y%m', partition_id)) FROM `test_schema`.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '__temp_test_table_{temp_table_uuid.hex}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__');",
-        f"MERGE INTO `test_schema`.`test_table` AS `__MERGE_TARGET__` USING (SELECT `a`, `ds` FROM (SELECT * FROM `test_schema`.`__temp_test_table_{temp_table_uuid.hex}`) AS `_subquery` WHERE DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`)) AS `__MERGE_SOURCE__` ON FALSE WHEN NOT MATCHED BY SOURCE AND DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`) THEN DELETE WHEN NOT MATCHED THEN INSERT (`a`, `ds`) VALUES (`a`, `ds`)",
-        f"DROP TABLE IF EXISTS `test_schema`.`__temp_test_table_{temp_table_uuid.hex}`",
+        f"CREATE TABLE IF NOT EXISTS `test_schema`.`__temp_test_table_{temp_table_id}` PARTITION BY DATETIME_TRUNC(`ds`, MONTH) AS SELECT `a`, `ds` FROM `tbl`",
+        f"DECLARE _sqlmesh_target_partitions_ ARRAY<DATETIME> DEFAULT (SELECT ARRAY_AGG(PARSE_DATETIME('%Y%m', partition_id)) FROM `test_schema`.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '__temp_test_table_{temp_table_id}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__');",
+        f"MERGE INTO `test_schema`.`test_table` AS `__MERGE_TARGET__` USING (SELECT `a`, `ds` FROM (SELECT * FROM `test_schema`.`__temp_test_table_{temp_table_id}`) AS `_subquery` WHERE DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`)) AS `__MERGE_SOURCE__` ON FALSE WHEN NOT MATCHED BY SOURCE AND DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`) THEN DELETE WHEN NOT MATCHED THEN INSERT (`a`, `ds`) VALUES (`a`, `ds`)",
+        f"DROP TABLE IF EXISTS `test_schema`.`__temp_test_table_{temp_table_id}`",
     ]
 
 
@@ -232,10 +233,6 @@ def test_replace_query_pandas(make_mocked_engine_adapter: t.Callable, mocker: Mo
         "sqlmesh.core.engine_adapter.bigquery.BigQueryEngineAdapter.execute"
     )
 
-    temp_table_uuid = uuid.uuid4()
-    uuid4_mock = mocker.patch("uuid.uuid4")
-    uuid4_mock.return_value = temp_table_uuid
-
     df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
     adapter.replace_query(
         "test_table", df, {"a": exp.DataType.build("int"), "b": exp.DataType.build("int")}

tests/core/engine_adapter/test_mixins.py

@@ -1,6 +1,5 @@
 # type: ignore
 import typing as t
-import uuid
 from unittest.mock import call
 
 from pytest_mock.plugin import MockerFixture
@@ -53,13 +52,15 @@ def test_logical_replace_query_does_not_exist(
 
 
 def test_logical_replace_self_reference(
-    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
 ):
     adapter = make_mocked_engine_adapter(LogicalReplaceQueryMixin, "postgres")
     adapter.cursor.fetchone.return_value = (1,)
-    temp_table_uuid = uuid.uuid4()
-    uuid4_mock = mocker.patch("uuid.uuid4")
-    uuid4_mock.return_value = temp_table_uuid
+
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
+    table_name = "db.table"
+    temp_table_id = "abcdefgh"
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
 
     mocker.patch(
         "sqlmesh.core.engine_adapter.postgres.LogicalReplaceQueryMixin.table_exists",
@@ -70,14 +71,14 @@ def test_logical_replace_self_reference(
         return_value={"col": exp.DataType(this=exp.DataType.Type.INT)},
     )
 
-    adapter.replace_query("db.table", parse_one("SELECT col + 1 AS col FROM db.table"))
+    adapter.replace_query(table_name, parse_one(f"SELECT col + 1 AS col FROM {table_name}"))
 
     assert to_sql_calls(adapter) == [
         f'CREATE SCHEMA IF NOT EXISTS "db"',
-        f'CREATE TABLE IF NOT EXISTS "db"."__temp_table_{temp_table_uuid.hex}" AS SELECT "col" FROM "db"."table"',
+        f'CREATE TABLE IF NOT EXISTS "db"."__temp_table_{temp_table_id}" AS SELECT "col" FROM "db"."table"',
         'TRUNCATE "db"."table"',
-        f'INSERT INTO "db"."table" ("col") SELECT "col" + 1 AS "col" FROM "db"."__temp_table_{temp_table_uuid.hex}"',
-        f'DROP TABLE IF EXISTS "db"."__temp_table_{temp_table_uuid.hex}"',
+        f'INSERT INTO "db"."table" ("col") SELECT "col" + 1 AS "col" FROM "db"."__temp_table_{temp_table_id}"',
+        f'DROP TABLE IF EXISTS "db"."__temp_table_{temp_table_id}"',
     ]