Add team_id to variable APIs

Author: vincbeck

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

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import json
+from uuid import UUID
 
 from pydantic import Field, JsonValue, model_validator
 
@@ -34,6 +35,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:
@@ -56,6 +58,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

@@ -12961,6 +12961,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:
@@ -13000,12 +13006,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/api_fastapi/execution_api/datamodels/variable.py

@@ -34,3 +34,4 @@ class VariablePostBody(StrictBaseModel):
 
     value: str | None = Field(alias="val")
     description: str | None = Field(default=None)
+    team_id: str | None = Field(default=None)

airflow-core/src/airflow/dag_processing/processor.py

@@ -549,7 +549,7 @@ def _handle_request(self, msg: ToManager, log: FilteringBoundLogger, req_id: int
             else:
                 resp = var
         elif isinstance(msg, PutVariable):
-            self.client.variables.set(msg.key, msg.value, msg.description)
+            self.client.variables.set(msg.key, msg.value, msg.description, msg.team_id)
         elif isinstance(msg, DeleteVariable):
             resp = self.client.variables.delete(msg.key)
         elif isinstance(msg, GetPreviousDagRun):

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

@@ -146,7 +146,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(
@@ -186,6 +186,7 @@ def set(
         value: Any,
         description: str | None = None,
         serialize_json: bool = False,
+        team_id: str | None = None,
         session: Session | None = None,
     ) -> None:
         """
@@ -197,13 +198,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(
@@ -219,6 +221,7 @@ def set(
                 value=value,
                 description=description,
                 serialize_json=serialize_json,
+                team_id=team_id,
             )
             return
 
@@ -236,7 +239,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
@@ -255,6 +258,7 @@ def set(
                 val=val,
                 description=description,
                 is_encrypted=is_encrypted,
+                team_id=team_id,
             )
 
             # Apply dialect-specific upsert
@@ -264,6 +268,7 @@ def set(
                     val=val,
                     description=description,
                     is_encrypted=is_encrypted,
+                    team_id=team_id,
                 )
             else:
                 # PostgreSQL and SQLite: ON CONFLICT DO UPDATE
@@ -273,6 +278,7 @@ def set(
                         val=val,
                         description=description,
                         is_encrypted=is_encrypted,
+                        team_id=team_id,
                     ),
                 )
 

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

@@ -6502,6 +6502,18 @@ export const $VariableBody = {
                 }
             ],
             title: 'Description'
+        },
+        team_id: {
+            anyOf: [
+                {
+                    type: 'string',
+                    format: 'uuid'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Team Id'
         }
     },
     additionalProperties: false,
@@ -6555,10 +6567,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

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