Merge branch 'main' into fix/sqlmesh-stale-cache-file-not-found

Author: Pavan-Rana

docs/guides/vscode.md

@@ -6,7 +6,7 @@
 
     The SQLMesh Visual Studio Code extension is in preview and undergoing active development. You may encounter bugs or API incompatibilities with the SQLMesh version you are running.
 
-    We encourage you to try the extension and [create Github issues](https://github.com/tobikodata/sqlmesh/issues) for any problems you encounter.
+    We encourage you to try the extension and [create Github issues](https://github.com/SQLMesh/sqlmesh/issues) for any problems you encounter.
 
 In this guide, you'll set up the SQLMesh extension in the Visual Studio Code IDE software (which we refer to as "VSCode").
 
@@ -187,7 +187,7 @@ The most common problem is the extension not using the correct Python interprete
 
 Follow the [setup process described above](#vscode-python-interpreter) to ensure that the extension is using the correct Python interpreter.
 
-If you have checked the VSCode `sqlmesh` output channel and the extension is still not using the correct Python interpreter, please raise an issue [here](https://github.com/tobikodata/sqlmesh/issues).
+If you have checked the VSCode `sqlmesh` output channel and the extension is still not using the correct Python interpreter, please raise an issue [here](https://github.com/SQLMesh/sqlmesh/issues).
 
 ### Missing Python dependencies
 
@@ -205,4 +205,4 @@ If you are using Tobiko Cloud, make sure `lsp` is included in the list of extras
 
 While the SQLMesh VSCode extension is in preview and the APIs to the underlying SQLMesh version are not stable, we do not guarantee compatibility between the extension and the SQLMesh version you are using.
 
-If you encounter a problem, please raise an issue [here](https://github.com/tobikodata/sqlmesh/issues).
+If you encounter a problem, please raise an issue [here](https://github.com/SQLMesh/sqlmesh/issues).

sqlmesh/core/config/connection.py

@@ -1062,6 +1062,7 @@ class BigQueryConnectionConfig(ConnectionConfig):
     job_retry_deadline_seconds: t.Optional[int] = None
     priority: t.Optional[BigQueryPriority] = None
     maximum_bytes_billed: t.Optional[int] = None
+    reservation: t.Optional[str] = None
 
     concurrent_tasks: int = 1
     register_comments: bool = True
@@ -1171,6 +1172,7 @@ def _extra_engine_config(self) -> t.Dict[str, t.Any]:
                 "job_retry_deadline_seconds",
                 "priority",
                 "maximum_bytes_billed",
+                "reservation",
             }
         }
 

sqlmesh/core/engine_adapter/bigquery.py

@@ -140,8 +140,10 @@ def _job_params(self) -> t.Dict[str, t.Any]:
                 "priority", BigQueryPriority.INTERACTIVE.bigquery_constant
             ),
         }
-        if self._extra_config.get("maximum_bytes_billed"):
+        if self._extra_config.get("maximum_bytes_billed") is not None:
             params["maximum_bytes_billed"] = self._extra_config.get("maximum_bytes_billed")
+        if self._extra_config.get("reservation") is not None:
+            params["reservation"] = self._extra_config.get("reservation")
         if self.correlation_id:
             # BigQuery label keys must be lowercase
             key = self.correlation_id.job_type.value.lower()
@@ -1106,7 +1108,9 @@ def _execute(
             else []
         )
 
+        # Create job config
         job_config = QueryJobConfig(**self._job_params, connection_properties=connection_properties)
+
         self._query_job = self._db_call(
             self.client.query,
             query=sql,

tests/cli/test_cli.py

@@ -1949,7 +1949,7 @@ def test_init_dbt_template(runner: CliRunner, tmp_path: Path):
 def test_init_project_engine_configs(tmp_path):
     engine_type_to_config = {
         "redshift": "# concurrent_tasks: 4\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # user: \n      # password: \n      # database: \n      # host: \n      # port: \n      # source_address: \n      # unix_sock: \n      # ssl: \n      # sslmode: \n      # timeout: \n      # tcp_keepalive: \n      # application_name: \n      # preferred_role: \n      # principal_arn: \n      # credentials_provider: \n      # region: \n      # cluster_identifier: \n      # iam: \n      # is_serverless: \n      # serverless_acct_id: \n      # serverless_work_group: \n      # enable_merge: ",
-        "bigquery": "# concurrent_tasks: 1\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # method: oauth\n      # project: \n      # execution_project: \n      # quota_project: \n      # location: \n      # keyfile: \n      # keyfile_json: \n      # token: \n      # refresh_token: \n      # client_id: \n      # client_secret: \n      # token_uri: \n      # scopes: \n      # impersonated_service_account: \n      # job_creation_timeout_seconds: \n      # job_execution_timeout_seconds: \n      # job_retries: 1\n      # job_retry_deadline_seconds: \n      # priority: \n      # maximum_bytes_billed: ",
+        "bigquery": "# concurrent_tasks: 1\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # method: oauth\n      # project: \n      # execution_project: \n      # quota_project: \n      # location: \n      # keyfile: \n      # keyfile_json: \n      # token: \n      # refresh_token: \n      # client_id: \n      # client_secret: \n      # token_uri: \n      # scopes: \n      # impersonated_service_account: \n      # job_creation_timeout_seconds: \n      # job_execution_timeout_seconds: \n      # job_retries: 1\n      # job_retry_deadline_seconds: \n      # priority: \n      # maximum_bytes_billed: \n      # reservation: ",
         "snowflake": "account: \n      # concurrent_tasks: 4\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # user: \n      # password: \n      # warehouse: \n      # database: \n      # role: \n      # authenticator: \n      # token: \n      # host: \n      # port: \n      # application: Tobiko_SQLMesh\n      # private_key: \n      # private_key_path: \n      # private_key_passphrase: \n      # session_parameters: ",
         "databricks": "# concurrent_tasks: 1\n      # register_comments: True\n      # pre_ping: False\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # server_hostname: \n      # http_path: \n      # access_token: \n      # auth_type: \n      # oauth_client_id: \n      # oauth_client_secret: \n      # catalog: \n      # http_headers: \n      # session_configuration: \n      # databricks_connect_server_hostname: \n      # databricks_connect_access_token: \n      # databricks_connect_cluster_id: \n      # databricks_connect_use_serverless: False\n      # force_databricks_connect: False\n      # disable_databricks_connect: False\n      # disable_spark_session: False",
         "postgres": "host: \n      user: \n      password: \n      port: \n      database: \n      # concurrent_tasks: 4\n      # register_comments: True\n      # pre_ping: True\n      # pretty_sql: False\n      # schema_differ_overrides: \n      # catalog_type_overrides: \n      # keepalives_idle: \n      # connect_timeout: 10\n      # role: \n      # sslmode: \n      # application_name: ",

tests/core/test_connection_config.py

@@ -1131,6 +1131,27 @@ def test_bigquery(make_config):
     assert config.get_catalog() == "project"
     assert config.is_recommended_for_state_sync is False
 
+    # Test reservation
+    config_with_reservation = make_config(
+        type="bigquery",
+        project="project",
+        reservation="projects/my-project/locations/us-central1/reservations/my-reservation",
+        check_import=False,
+    )
+    assert isinstance(config_with_reservation, BigQueryConnectionConfig)
+    assert (
+        config_with_reservation.reservation
+        == "projects/my-project/locations/us-central1/reservations/my-reservation"
+    )
+
+    # Test that reservation is included in _extra_engine_config
+    extra_config = config_with_reservation._extra_engine_config
+    assert "reservation" in extra_config
+    assert (
+        extra_config["reservation"]
+        == "projects/my-project/locations/us-central1/reservations/my-reservation"
+    )
+
     with pytest.raises(ConfigError, match="you must also specify the `project` field"):
         make_config(type="bigquery", execution_project="execution_project", check_import=False)
 

vscode/extension/README.md

@@ -77,8 +77,8 @@ If you encounter issues, please refer to the [VSCode Extension Guide](https://sq
 
 We welcome contributions! Please:
 
-1. [Report bugs](https://github.com/tobikodata/sqlmesh/issues) you encounter
-2. [Request features](https://github.com/tobikodata/sqlmesh/issues) you'd like to see
+1. [Report bugs](https://github.com/SQLMesh/sqlmesh/issues) you encounter
+2. [Request features](https://github.com/SQLMesh/sqlmesh/issues) you'd like to see
 3. Share feedback on your experience
 
 ## 📄 License
@@ -87,7 +87,7 @@ This extension is licensed under the Apache License 2.0. See [LICENSE](LICENSE)
 
 ## 🔗 Links
 
-- [SQLMesh GitHub Repository](https://github.com/tobikodata/sqlmesh)
+- [SQLMesh GitHub Repository](https://github.com/SQLMesh/sqlmesh)
 - [Tobiko Data Website](https://tobikodata.com)
 - [Extension Marketplace Page](https://marketplace.visualstudio.com/items?itemName=tobikodata.sqlmesh)
 

vscode/extension/package.json

@@ -6,7 +6,7 @@
   "version": "0.0.7",
   "repository": {
     "type": "git",
-    "url": "https://github.com/tobikodata/sqlmesh"
+    "url": "https://github.com/SQLMesh/sqlmesh"
   },
   "main": "./dist/extension.js",
   "icon": "assets/logo.png",