Add team_id to variable APIs

Author: vincbeck

airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables.py

@@ -19,6 +19,7 @@
 
 import json
 from collections.abc import Iterable
+from uuid import UUID
 
 from pydantic import Field, JsonValue, model_validator
 
@@ -35,6 +36,7 @@ class VariableResponse(BaseModel):
     val: str = Field(alias="value")
     description: str | None
     is_encrypted: bool
+    team_id: UUID | None
 
     @model_validator(mode="after")
     def redact_val(self) -> Self:
@@ -57,6 +59,7 @@ class VariableBody(StrictBaseModel):
     key: str = Field(max_length=ID_LEN)
     value: JsonValue = Field(serialization_alias="val")
     description: str | None = Field(default=None)
+    team_id: UUID | None = Field(default=None)
 
 
 class VariableCollectionResponse(BaseModel):

airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml

@@ -13068,6 +13068,12 @@ components:
           - type: string
           - type: 'null'
           title: Description
+        team_id:
+          anyOf:
+          - type: string
+            format: uuid
+          - type: 'null'
+          title: Team Id
       additionalProperties: false
       type: object
       required:
@@ -13107,12 +13113,19 @@ components:
         is_encrypted:
           type: boolean
           title: Is Encrypted
+        team_id:
+          anyOf:
+          - type: string
+            format: uuid
+          - type: 'null'
+          title: Team Id
       type: object
       required:
       - key
       - value
       - description
       - is_encrypted
+      - team_id
       title: VariableResponse
       description: Variable serializer for responses.
     VersionInfo:

airflow-core/src/airflow/models/variable.py

@@ -30,7 +30,7 @@
 from sqlalchemy_utils import UUIDType
 
 from airflow._shared.secrets_masker import mask_secret
-from airflow.configuration import ensure_secrets_loaded
+from airflow.configuration import conf, ensure_secrets_loaded
 from airflow.models.base import ID_LEN, Base
 from airflow.models.crypto import get_fernet
 from airflow.models.team import Team
@@ -149,7 +149,7 @@ def get(
         # means SQLA etc is loaded, but we can't avoid that unless/until we add import shims as a big
         # back-compat layer
 
-        # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
+        # If this is set it means we are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
         # and should use the Task SDK API server path
         if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
             warnings.warn(
@@ -185,6 +185,7 @@ def set(
         value: Any,
         description: str | None = None,
         serialize_json: bool = False,
+        team_id: str | None = None,
         session: Session | None = None,
     ) -> None:
         """
@@ -196,13 +197,14 @@ def set(
         :param value: Value to set for the Variable
         :param description: Description of the Variable
         :param serialize_json: Serialize the value to a JSON string
+        :param team_id: ID of the team associated to the variable (if any)
         :param session: optional session, use if provided or create a new one
         """
         # TODO: This is not the best way of having compat, but it's "better than erroring" for now. This still
         # means SQLA etc is loaded, but we can't avoid that unless/until we add import shims as a big
         # back-compat layer
 
-        # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
+        # If this is set it means we are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
         # and should use the Task SDK API server path
         if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
             warnings.warn(
@@ -221,6 +223,11 @@ def set(
             )
             return
 
+        if team_id and not conf.getboolean("core", "multi_team"):
+            raise ValueError(
+                "Multi-team mode is not configured in the Airflow environment. To assign a team to a variable, multi-mode must be enabled."
+            )
+
         # check if the secret exists in the custom secrets' backend.
         Variable.check_for_write_conflict(key=key)
         if serialize_json:
@@ -235,7 +242,7 @@ def set(
             ctx = create_session()
 
         with ctx as session:
-            new_variable = Variable(key=key, val=stored_value, description=description)
+            new_variable = Variable(key=key, val=stored_value, description=description, team_id=team_id)
 
             val = new_variable._val
             is_encrypted = new_variable.is_encrypted
@@ -252,13 +259,15 @@ def set(
                     val=val,
                     description=description,
                     is_encrypted=is_encrypted,
+                    team_id=team_id,
                 )
                 stmt = pg_stmt.on_conflict_do_update(
                     index_elements=["key"],
                     set_=dict(
                         val=val,
                         description=description,
                         is_encrypted=is_encrypted,
+                        team_id=team_id,
                     ),
                 )
             elif dialect_name == "mysql":
@@ -269,11 +278,13 @@ def set(
                     val=val,
                     description=description,
                     is_encrypted=is_encrypted,
+                    team_id=team_id,
                 )
                 stmt = mysql_stmt.on_duplicate_key_update(
                     val=val,
                     description=description,
                     is_encrypted=is_encrypted,
+                    team_id=team_id,
                 )
             else:
                 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
@@ -283,13 +294,15 @@ def set(
                     val=val,
                     description=description,
                     is_encrypted=is_encrypted,
+                    team_id=team_id,
                 )
                 stmt = sqlite_stmt.on_conflict_do_update(
                     index_elements=["key"],
                     set_=dict(
                         val=val,
                         description=description,
                         is_encrypted=is_encrypted,
+                        team_id=team_id,
                     ),
                 )
 

airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts

@@ -6559,6 +6559,18 @@ export const $VariableBody = {
                 }
             ],
             title: 'Description'
+        },
+        team_id: {
+            anyOf: [
+                {
+                    type: 'string',
+                    format: 'uuid'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Team Id'
         }
     },
     additionalProperties: false,
@@ -6612,10 +6624,22 @@ export const $VariableResponse = {
         is_encrypted: {
             type: 'boolean',
             title: 'Is Encrypted'
+        },
+        team_id: {
+            anyOf: [
+                {
+                    type: 'string',
+                    format: 'uuid'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Team Id'
         }
     },
     type: 'object',
-    required: ['key', 'value', 'description', 'is_encrypted'],
+    required: ['key', 'value', 'description', 'is_encrypted', 'team_id'],
     title: 'VariableResponse',
     description: 'Variable serializer for responses.'
 } as const;

airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts

@@ -1594,6 +1594,7 @@ export type VariableBody = {
     key: string;
     value: JsonValue;
     description?: string | null;
+    team_id?: string | null;
 };
 
 /**
@@ -1612,6 +1613,7 @@ export type VariableResponse = {
     value: string;
     description: string | null;
     is_encrypted: boolean;
+    team_id: string | null;
 };
 
 /**