SG-43167 Add flow_am_project_id to Toolkit Context

python/tank/authentication/flow_auth/_client.py

@@ -27,6 +27,7 @@ def get_authentication_token(self) -> str:
         return get_flow_access_token()
 
 
+# TODO: replace with proper client entry point
 def get_flow_client(endpoint_url=None):
     """Return a ready-to-use Flow GQL SDK client with authentication wired in.
 

python/tank/bootstrap/manager.py

@@ -15,7 +15,7 @@
 from .errors import TankBootstrapError
 from .configuration import Configuration
 from .resolver import ConfigurationResolver
-from ..authentication import ShotgunAuthenticator
+from ..authentication import ShotgunAuthenticator, flow_auth
 from ..pipelineconfig import PipelineConfiguration
 from .. import LogManager
 from ..errors import TankError
@@ -33,12 +33,12 @@ class ToolkitManager(object):
     # Constants used to make the manager bootstrapping:
     # - download and cache the config dependencies needed to run the engine being started in a specific environment.
     # - download and cache all the config dependencies needed to run the engine in any environment.
-    (CACHE_SPARSE, CACHE_FULL) = range(2)
+    CACHE_SPARSE, CACHE_FULL = range(2)
 
     # Constants used to indicate that the manager is:
     # - bootstrapping the toolkit (with method bootstrap_toolkit),
     # - starting up the engine (with method _start_engine).
-    (TOOLKIT_BOOTSTRAP_PHASE, ENGINE_STARTUP_PHASE) = range(2)
+    TOOLKIT_BOOTSTRAP_PHASE, ENGINE_STARTUP_PHASE = range(2)
 
     # List of constants representing the status of the progress bar when these event occurs during bootstrap.
     _RESOLVING_PROJECT_RATE = 0.0
@@ -962,8 +962,6 @@ def _check_and_trigger_am_auth(self, entity, progress_callback):
                                   Set to ``None`` to use the default callback function.
         :rtype: None
         """
-        from ..authentication import flow_auth
-
         if entity is None:
             return
 
@@ -1090,7 +1088,9 @@ def _get_configuration(self, entity, progress_callback):
 
         elif self._do_shotgun_config_lookup:
             # do the full resolve where we connect to shotgun etc.
-            log.debug("Checking for pipeline configuration overrides in Flow Production Tracking.")
+            log.debug(
+                "Checking for pipeline configuration overrides in Flow Production Tracking."
+            )
             log.debug(
                 "In order to turn this off, set do_shotgun_config_lookup to False"
             )

python/tank/context.py

@@ -19,6 +19,7 @@
 
 from tank_vendor import yaml
 from . import authentication
+from .authentication import flow_auth
 
 from .util import login
 from .util import shotgun_entity
@@ -403,6 +404,22 @@ def additional_entities(self):
         """
         return self.__additional_entities
 
+    @property
+    def flow_am_project_id(self):
+        """
+        The FlowAM project ID for this context, or ``None`` if the project is
+        not FlowAM-enabled or there is no project in this context.
+
+        The value is read from the ``sg_flow_am_id`` field on the project dict.
+        It is populated by the bootstrap manager after it queries ShotGrid.
+
+        :returns: A string containing the FlowAM project ID, or ``None``.
+        :rtype: str or None
+        """
+        if self.project:
+            return self.project.get(flow_auth.AM_READY_PROJECT_FIELD)
+        return None
+
     @property
     def entity_locations(self):
         """
@@ -815,8 +832,8 @@ def deserialize(cls, context_str):
     def to_dict(self):
         """
         Converts the context into a dictionary with keys ``project``,
-        ``entity``, ``user``, ``step``, ``task``, ``additional_entities`` and
-        ``source_entity``.
+        ``entity``, ``user``, ``step``, ``task``, ``additional_entities``,
+        ``source_entity``, and ``flow_am_project_id``.
 
         .. note ::
             Contrary to :meth:`Context.serialize`, this method discards information
@@ -835,6 +852,7 @@ def to_dict(self):
                 self._cleanup_entity(entity) for entity in self.additional_entities
             ],
             "source_entity": self._cleanup_entity(self.source_entity),
+            "flow_am_project_id": self.flow_am_project_id,
         }
 
     def _cleanup_entity(self, entity):
@@ -894,7 +912,7 @@ def _from_dict(cls, data):
 
         :returns: :class:`Context`
         """
-        return Context(
+        ctx = Context(
             tk=data.get("tk"),
             project=data.get("project"),
             entity=data.get("entity"),
@@ -904,6 +922,9 @@ def _from_dict(cls, data):
             additional_entities=data.get("additional_entities"),
             source_entity=data.get("source_entity"),
         )
+        if ctx.project is not None and "flow_am_project_id" in data:
+            ctx.project[flow_auth.AM_READY_PROJECT_FIELD] = data["flow_am_project_id"]
+        return ctx
 
     ################################################################################################
     # private methods
@@ -1305,7 +1326,8 @@ def _from_entity_type_and_id(tk, entity, source_entity=None):
 
         if sg_entity is None:
             raise TankError(
-                "Entity %s with id %s not found in Flow Production Tracking!" % (entity_type, entity_id)
+                "Entity %s with id %s not found in Flow Production Tracking!"
+                % (entity_type, entity_id)
             )
 
         if sg_entity.get("task"):
@@ -1719,7 +1741,7 @@ def context_yaml_representer(dumper, context):
     # pipeline config path as part of the dict
     context_dict["_pc_path"] = context.tank.pipeline_configuration.get_path()
 
-    return dumper.represent_mapping(u"!TankContext", context_dict)
+    return dumper.represent_mapping("!TankContext", context_dict)
 
 
 def context_yaml_constructor(loader, node):
@@ -1750,7 +1772,7 @@ def context_yaml_constructor(loader, node):
 
 
 yaml.add_representer(Context, context_yaml_representer)
-yaml.add_constructor(u"!TankContext", context_yaml_constructor)
+yaml.add_constructor("!TankContext", context_yaml_constructor)
 
 ################################################################################################
 # utility methods
@@ -1807,7 +1829,9 @@ def _task_from_sg(tk, task_id, additional_fields=None):
         "Task", [["id", "is", task_id]], standard_fields + additional_fields
     )
     if not task:
-        raise TankError("Unable to locate Task with id %s in Flow Production Tracking" % task_id)
+        raise TankError(
+            "Unable to locate Task with id %s in Flow Production Tracking" % task_id
+        )
 
     # add task so it can be processed with other shotgun entities
     task["task"] = {"type": "Task", "id": task_id, "name": task["content"]}
@@ -1867,7 +1891,8 @@ def _entity_from_sg(tk, entity_type, entity_id):
 
     if not data:
         raise TankError(
-            "Unable to locate %s with id %s in Flow Production Tracking" % (entity_type, entity_id)
+            "Unable to locate %s with id %s in Flow Production Tracking"
+            % (entity_type, entity_id)
         )
 
     # create context

tests/core_tests/test_context.py

@@ -1261,6 +1261,7 @@ def test_dict_cleanup(self):
                 {"type": "Sequence", "name": "seq_name", "id": self.seq["id"]}
             ],
             "user": {"type": "HumanUser", "id": self.user["id"], "name": USER_NAME},
+            "flow_am_project_id": None,
         }
 
         ctx = context.Context(**self.kws)
@@ -1438,3 +1439,56 @@ def test_non_primary_path(self, get_current_user):
 
         self.assertIsNone(result.step)
         self.assertIsNone(result.task)
+
+
+class TestContextFlowAmProjectId(TankTestBase):
+    """Tests for Context.flow_am_project_id."""
+
+    def setUp(self):
+        super().setUp()
+        self.setup_fixtures()
+
+    def test_returns_none_for_non_am_project(self):
+        """flow_am_project_id is None when sg_flow_am_id is not set on the project."""
+        ctx = context.Context(self.tk, project=self.project)
+        self.assertIsNone(ctx.flow_am_project_id)
+
+    def test_returns_none_for_empty_context(self):
+        """flow_am_project_id is None when the context has no project."""
+        ctx = context.create_empty(self.tk)
+        self.assertIsNone(ctx.flow_am_project_id)
+
+    def test_returns_value_when_set(self):
+        """flow_am_project_id returns sg_flow_am_id when set on the project dict."""
+        project = dict(self.project, sg_flow_am_id="am-project-abc")
+        ctx = context.Context(self.tk, project=project)
+        self.assertEqual(ctx.flow_am_project_id, "am-project-abc")
+
+    def test_survives_serialization_roundtrip(self):
+        """flow_am_project_id is preserved through serialize/deserialize."""
+        project = dict(self.project, sg_flow_am_id="am-project-abc")
+        ctx = context.Context(self.tk, project=project)
+        serialized = ctx.serialize()
+        restored = context.deserialize(serialized)
+        self.assertEqual(restored.flow_am_project_id, "am-project-abc")
+
+    def test_set_on_project_dict(self):
+        """flow_am_project_id can be set by writing sg_flow_am_id onto the project dict."""
+        ctx = context.Context(self.tk, project=dict(self.project))
+        ctx.project["sg_flow_am_id"] = "explicit-am-id"
+        self.assertEqual(ctx.flow_am_project_id, "explicit-am-id")
+
+    def test_included_in_to_dict(self):
+        """flow_am_project_id appears in the dict returned by to_dict()."""
+        project = dict(self.project, sg_flow_am_id="explicit-am-id")
+        ctx = context.Context(self.tk, project=project)
+        d = ctx.to_dict()
+        self.assertIn("flow_am_project_id", d)
+        self.assertEqual(d["flow_am_project_id"], "explicit-am-id")
+
+    def test_deepcopy_preserves_value(self):
+        """Deepcopying a context preserves flow_am_project_id."""
+        project = dict(self.project, sg_flow_am_id="am-project-abc")
+        ctx = context.Context(self.tk, project=project)
+        ctx_copy = copy.deepcopy(ctx)
+        self.assertEqual(ctx_copy.flow_am_project_id, "am-project-abc")