Merge branch 'main' into fix/invalidate-nonexistent-environment

Author: mday-io

sqlmesh/core/dialect.py

@@ -556,6 +556,10 @@ def _parse_if(self: Parser) -> t.Optional[exp.Expr]:
     # to parse a statement / command to support the macro @IF(condition, statement)
     index = self._index
     try:
+        if self.dialect == "tsql":
+            if not (self._index >= 2 and self._tokens[self._index - 2].text == "@"):
+                return self.__parse_if()  # type: ignore
+            return Parser.__parse_if(self)  # type: ignore
         return self.__parse_if()  # type: ignore
     except ParseError:
         self._retreat(index)
@@ -1133,8 +1137,8 @@ def extend_sqlglot() -> None:
     _override(Parser, _parse_value)
     _override(Parser, _parse_lambda)
     _override(Parser, _parse_types)
-    _override(TSQL.Parser, Parser._parse_if)
     _override(Parser, _parse_if)
+    _override(TSQL.Parser, Parser._parse_if)
     _override(Parser, _parse_id_var)
     _override(Parser, _parse_interval_span)
     _override(Parser, _warn_unsupported)

sqlmesh/core/engine_adapter/snowflake.py

@@ -407,7 +407,7 @@ def query_factory() -> Query:
             elif isinstance(df, pd.DataFrame):
                 from snowflake.connector.pandas_tools import write_pandas
 
-                ordered_df = df[list(source_columns_to_types)]
+                ordered_df = df[list(source_columns_to_types)].reset_index(drop=True)
 
                 # Workaround for https://github.com/snowflakedb/snowflake-connector-python/issues/1034
                 # The above issue has already been fixed upstream, but we keep the following

sqlmesh/core/model/seed.py

@@ -113,7 +113,7 @@ def read(self, batch_size: t.Optional[int] = None) -> t.Generator[pd.DataFrame,
         batch_size = batch_size or df.size
         batch_start = 0
         while batch_start < df.shape[0]:
-            yield df.iloc[batch_start : batch_start + batch_size, :]
+            yield df.iloc[batch_start : batch_start + batch_size, :].copy()
             batch_start += batch_size
 
     def _get_df(self) -> pd.DataFrame:

sqlmesh/dbt/source.py

@@ -82,10 +82,12 @@ def canonical_name(self, context: DbtContext) -> str:
                     f"'source' macro failed for '{self.config_name}' with exception '{e}'."
                 )
 
+            identifier = relation.identifier or ""
+            needs_identifier_quoting = "." in identifier or " " in identifier
             relation = relation.quote(
                 database=False,
                 schema=False,
-                identifier=False,
+                identifier=needs_identifier_quoting,
             )
             if relation.database == context.target.database:
                 relation = relation.include(database=False)

tests/core/engine_adapter/test_snowflake.py

@@ -469,6 +469,26 @@ def test_df_to_source_queries_use_schema(
     assert 'USE SCHEMA "other_catalog"."other_db"' in to_sql_calls(adapter)
 
 
+def test_df_to_source_queries_reset_non_default_index(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.snowflake.SnowflakeEngineAdapter.table_exists",
+        return_value=False,
+    )
+    write_pandas = mocker.patch("snowflake.connector.pandas_tools.write_pandas", return_value=None)
+    adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)
+
+    df = pd.DataFrame({"a": [2, 3], "b": [5, 6]}, index=[1, 2])
+    adapter.replace_query(
+        "other_db.test_table", df, {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}
+    )
+
+    uploaded_df = write_pandas.call_args.args[1]
+    assert uploaded_df.index.equals(pd.RangeIndex(start=0, stop=2, step=1))
+    assert uploaded_df.to_dict("list") == {"a": [2, 3], "b": [5, 6]}
+
+
 def test_create_managed_table(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
     adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)
 

tests/core/test_dialect.py

@@ -709,6 +709,20 @@ def test_conditional_statement():
     q = parse_one("@IF(cond, VACUUM ANALYZE);", read="postgres")
     assert q.sql(dialect="postgres") == "@IF(cond, VACUUM ANALYZE)"
 
+    # Verify that the original error case from issue #5823 (Required keyword: 'true' missing) is resolved.
+    # It must be parsed as a macro function containing an Anonymous expression rather than exp.If.
+    q = parse_one("@IF(1 = 1, ALTER TABLE x ADD y INT);", read="tsql")
+    assert q.sql(dialect="tsql") == "@IF(1 = 1, ALTER TABLE x ADD y INTEGER)"
+    assert isinstance(q.this, exp.Anonymous)
+    assert q.this.name == "IF"
+
+    # Note: SQLGlot's fallback Command parser strips quotes from string literal tokens when parsing unparsed commands
+    q = parse_one("@IF(cond, PRINT 'hello');", read="tsql")
+    assert q.sql(dialect="tsql") == "@IF(cond, PRINT hello)"
+
+    q = parse_one("@IF(@runtime_stage = 'evaluating', SELECT 1);", read="tsql")
+    assert q.sql(dialect="tsql") == "@IF(@runtime_stage = 'evaluating', SELECT 1)"
+
 
 def test_model_name_cannot_be_string():
     with pytest.raises(ParseError) as parse_error:

tests/core/test_seed.py

@@ -58,6 +58,21 @@ def test_read_custom_settings():
     pd.testing.assert_frame_equal(next(dfs), expected_df)
 
 
+def test_read_returns_independent_batches():
+    content = """key,value
+1,one
+2,two
+"""
+    seed = Seed(content=content)
+    seed_reader = seed.reader()
+
+    batches = list(seed_reader.read(batch_size=1))
+    batches[0].at[0, "value"] = "changed"
+
+    assert [df["value"].tolist() for df in batches] == [["changed"], ["two"]]
+    assert next(seed_reader.read())["value"].tolist() == ["one", "two"]
+
+
 def test_column_hashes():
     content = """key,value,ds
 1,one,2022-01-01

tests/dbt/test_config.py

@@ -521,6 +521,71 @@ def test_quoting():
     assert str(BaseRelation.create(**source.relation_info)) == 'foo."bar"'
 
 
+def test_source_canonical_name_with_dots_and_spaces(mocker):
+    from dbt.adapters.base import BaseRelation
+
+    mock_context = mocker.Mock()
+    mock_context.target.database = "target_db"
+
+    def mock_source_macro(source_name, table_name):
+        if table_name == "my_table_dot":
+            identifier = "FILENAME.CSV"
+        elif table_name == "my_table_space":
+            identifier = "my table space"
+        else:
+            identifier = "my_table_std"
+        return BaseRelation.create(
+            database="RAW_DEV",
+            schema="raw_schema",
+            identifier=identifier,
+        )
+
+    mock_context.get_callable_macro.return_value = mock_source_macro
+
+    # 1. Identifier with a dot
+    source_dot = SourceConfig(
+        name="my_table_dot",
+        source_name="my_source",
+        identifier="FILENAME.CSV",
+    )
+    assert source_dot.canonical_name(mock_context) == 'RAW_DEV.raw_schema."FILENAME.CSV"'
+
+    # 2. Identifier with a space
+    source_space = SourceConfig(
+        name="my_table_space",
+        source_name="my_source",
+        identifier="my table space",
+    )
+    assert source_space.canonical_name(mock_context) == 'RAW_DEV.raw_schema."my table space"'
+
+    # 3. Standard identifier (without dots or spaces) should not be quoted
+    source_std = SourceConfig(
+        name="my_table_std",
+        source_name="my_source",
+        identifier="my_table_std",
+    )
+    assert source_std.canonical_name(mock_context) == "RAW_DEV.raw_schema.my_table_std"
+
+    # 4. Standard identifier, but with database matching target database (to test database omission)
+    mock_context_target_db = mocker.Mock()
+    mock_context_target_db.target.database = "RAW_DEV"
+    mock_context_target_db.get_callable_macro.return_value = mock_source_macro
+
+    source_dot_target = SourceConfig(
+        name="my_table_dot",
+        source_name="my_source",
+        identifier="FILENAME.CSV",
+    )
+    source_std_target = SourceConfig(
+        name="my_table_std",
+        source_name="my_source",
+        identifier="my_table_std",
+    )
+
+    assert source_dot_target.canonical_name(mock_context_target_db) == 'raw_schema."FILENAME.CSV"'
+    assert source_std_target.canonical_name(mock_context_target_db) == "raw_schema.my_table_std"
+
+
 def _test_warehouse_config(
     config_yaml: str, target_class: t.Type[TargetConfig], *params_path: str
 ) -> TargetConfig: