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
@@ -941,16 +941,13 @@ def _resolve_project_id(self, entity):
             raise TankBootstrapError("Cannot resolve project for %s" % entity)
         return data["project"]["id"]
 
-    def _check_and_trigger_am_auth(self, entity, progress_callback):
+    def _trigger_am_auth(self, entity, progress_callback):
         """
-        If the resolved project is AM-ready, proactively obtain a Flow/MEDM
-        access token. Silent path (file store -> refresh) is tried first; falls
+        Proactively obtain a Flow/MEDM access token.
+        Silent path (file store -> refresh) is tried first; falls
         back to opening a browser for PKCE if no usable cached/refresh token
         exists.
 
-        No-op for non-AM projects or when ``entity`` is None. The project and
-        its AM-ready field are resolved in a single ShotGrid request.
-
         Configuration errors raise ``TankBootstrapError`` (deployment bug).
         Runtime auth failures are logged and swallowed unless the
         ``TK_FLOW_AUTH_REQUIRED`` env var is set to ``"1"``, in which case
@@ -962,45 +959,9 @@ 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
 
-        am_field = flow_auth.AM_READY_PROJECT_FIELD
-
-        if entity.get("type") == "Project":
-            project_id = entity["id"]
-            sg_project = self._sg_connection.find_one(
-                "Project", [["id", "is", project_id]], [am_field]
-            )
-        elif "project" in entity and entity["project"].get("type") == "Project":
-            project_id = entity["project"]["id"]
-            sg_project = self._sg_connection.find_one(
-                "Project", [["id", "is", project_id]], [am_field]
-            )
-        else:
-            # Fetch the project link and AM-ready field in a single request
-            # using ShotGrid's deep-field notation, saving one API round-trip.
-            data = self._sg_connection.find_one(
-                entity["type"],
-                [["id", "is", entity["id"]]],
-                ["project", "project.Project.%s" % am_field],
-            )
-            if not data or not data.get("project"):
-                return
-            project_id = data["project"]["id"]
-            sg_project = {am_field: data.get("project.Project.%s" % am_field)}
-
-        if not sg_project or not sg_project.get(am_field):
-            return
-
-        log.info("Project %s is AM-ready; triggering MEDM auth.", project_id)
-        self._report_progress(
-            progress_callback,
-            self._UPDATING_CONFIGURATION_RATE,
-            "Authenticating with Autodesk identity...",
-        )
         try:
             settings = flow_auth.resolve_flow_auth_settings()
             flow_auth.init_authentication(settings)
@@ -1009,14 +970,12 @@ def _check_and_trigger_am_auth(self, entity, progress_callback):
             flow_auth.get_access_token()
         except flow_auth.FlowAuthConfigurationError as e:
             raise TankBootstrapError(
-                "MEDM auth misconfigured for AM-ready project %s: %s"
-                % (project_id, e)
+                "MEDM auth misconfigured for AM-ready project: %s" % e
             )
         except Exception as e:
             if os.environ.get("TK_FLOW_AUTH_REQUIRED") == "1":
                 raise TankBootstrapError(
-                    "MEDM auth failed for AM-ready project %s: %s"
-                    % (project_id, e)
+                    "MEDM auth failed for AM-ready project: %s" % e
                 )
             log.warning(
                 "MEDM auth failed; bootstrap will continue without a "
@@ -1090,7 +1049,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"
             )
@@ -1156,7 +1117,22 @@ def _get_updated_configuration(self, entity, progress_callback):
         else:
             raise TankBootstrapError("Unknown configuration update status!")
 
-        self._check_and_trigger_am_auth(entity, progress_callback)
+        if entity is not None:
+            project_id = self._resolve_project_id(entity)
+            if project_id:
+                sg_project = self._sg_connection.find_one(
+                    "Project",
+                    [["id", "is", project_id]],
+                    [flow_auth.AM_READY_PROJECT_FIELD],
+                )
+                if sg_project and sg_project.get(flow_auth.AM_READY_PROJECT_FIELD):
+                    # Retrieve and cache the flow am project id on the context object
+                    flow_project_id = sg_project.get(flow_auth.AM_READY_PROJECT_FIELD)
+                    tk, _ = config.get_tk_instance(self._sg_user)
+                    ctx = tk.context_from_entity_dictionary(entity)
+                    ctx.project[flow_auth.AM_READY_PROJECT_FIELD] = flow_project_id
+                    # Authenticate into Flow AM
+                    self._trigger_am_auth(entity, progress_callback)
 
         return config
 

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,13 @@ 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
+            and flow_auth.AM_READY_PROJECT_FIELD not in ctx.project
+        ):
+            ctx.project[flow_auth.AM_READY_PROJECT_FIELD] = data["flow_am_project_id"]
+        return ctx
 
     ################################################################################################
     # private methods
@@ -1305,7 +1330,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 +1745,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 +1776,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 +1833,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 +1895,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/bootstrap_tests/test_manager_flow_auth.py

@@ -25,79 +25,33 @@
     "tank.authentication.ShotgunAuthenticator.get_user",
     return_value=mock.Mock(),
 )
-class FlowAuthHookTests(ShotgunTestBase):
-    """Coverage for ToolkitManager._check_and_trigger_am_auth."""
+class TriggerAmAuthTests(ShotgunTestBase):
+    """Coverage for ToolkitManager._trigger_am_auth."""
 
     PROJECT_ID = 42
 
-    def _build_manager_with_sg(self, sg_project_payload):
-        """Create a ToolkitManager whose _sg_connection.find_one returns
-        ``sg_project_payload`` for a Project query."""
-        mgr = ToolkitManager()
-        mgr._sg_connection = mock.Mock()
-        mgr._sg_connection.find_one.return_value = sg_project_payload
-        return mgr
-
     @mock.patch("tank.authentication.flow_auth.get_access_token")
     @mock.patch("tank.authentication.flow_auth.init_authentication")
     @mock.patch("tank.authentication.flow_auth.resolve_flow_auth_settings")
-    def test_am_ready_project_triggers_auth(self, mock_resolve, mock_init, mock_get, _):
+    def test_triggers_auth(self, mock_resolve, mock_init, mock_get, _):
         mock_resolve.return_value = mock.Mock()
-        mgr = self._build_manager_with_sg({flow_auth.AM_READY_PROJECT_FIELD: "abc-123"})
+        mgr = ToolkitManager()
 
-        mgr._check_and_trigger_am_auth(
+        mgr._trigger_am_auth(
             {"type": "Project", "id": self.PROJECT_ID}, progress_callback=mock.Mock()
         )
 
         mock_resolve.assert_called_once()
         mock_init.assert_called_once_with(mock_resolve.return_value)
         mock_get.assert_called_once()
 
-    @mock.patch("tank.authentication.flow_auth.get_access_token")
-    @mock.patch("tank.authentication.flow_auth.init_authentication")
-    def test_non_am_ready_project_skips_auth(self, mock_init, mock_get, _):
-        mgr = self._build_manager_with_sg({flow_auth.AM_READY_PROJECT_FIELD: None})
-
-        mgr._check_and_trigger_am_auth(
-            {"type": "Project", "id": self.PROJECT_ID}, progress_callback=None
-        )
-
-        mock_init.assert_not_called()
-        mock_get.assert_not_called()
-
-    @mock.patch("tank.authentication.flow_auth.get_access_token")
-    @mock.patch("tank.authentication.flow_auth.init_authentication")
-    def test_missing_project_entity_skips_auth(self, mock_init, mock_get, _):
-        mgr = self._build_manager_with_sg(None)
-
-        mgr._check_and_trigger_am_auth(
-            {"type": "Project", "id": self.PROJECT_ID}, progress_callback=None
-        )
-
-        mock_init.assert_not_called()
-        mock_get.assert_not_called()
-
-    @mock.patch("tank.authentication.flow_auth.get_access_token")
-    @mock.patch("tank.authentication.flow_auth.init_authentication")
-    def test_empty_project_entity_skips_auth(self, mock_init, mock_get, _):
-        mgr = self._build_manager_with_sg({})
-
-        mgr._check_and_trigger_am_auth(
-            {"type": "Project", "id": self.PROJECT_ID}, progress_callback=None
-        )
-
-        mock_init.assert_not_called()
-        mock_get.assert_not_called()
-
     @mock.patch("tank.authentication.flow_auth.get_access_token")
     @mock.patch("tank.authentication.flow_auth.init_authentication")
     def test_none_entity_skips_auth(self, mock_init, mock_get, _):
-        mgr = self._build_manager_with_sg({flow_auth.AM_READY_PROJECT_FIELD: "abc-123"})
+        mgr = ToolkitManager()
 
-        mgr._check_and_trigger_am_auth(None, progress_callback=None)
+        mgr._trigger_am_auth(None, progress_callback=None)
 
-        # find_one should not even be called when entity is None
-        mgr._sg_connection.find_one.assert_not_called()
         mock_init.assert_not_called()
         mock_get.assert_not_called()
 
@@ -108,10 +62,10 @@ def test_configuration_error_raises_TankBootstrapError(
     ):
         mock_resolve.return_value = mock.Mock()
         mock_init.side_effect = flow_auth.FlowAuthConfigurationError("missing app id")
-        mgr = self._build_manager_with_sg({flow_auth.AM_READY_PROJECT_FIELD: "abc-123"})
+        mgr = ToolkitManager()
 
         with self.assertRaises(TankBootstrapError):
-            mgr._check_and_trigger_am_auth(
+            mgr._trigger_am_auth(
                 {"type": "Project", "id": self.PROJECT_ID},
                 progress_callback=mock.Mock(),
             )
@@ -124,10 +78,10 @@ def test_runtime_error_soft_fails_by_default(
     ):
         mock_resolve.return_value = mock.Mock()
         mock_get.side_effect = RuntimeError("network down")
-        mgr = self._build_manager_with_sg({flow_auth.AM_READY_PROJECT_FIELD: "abc-123"})
+        mgr = ToolkitManager()
 
         # Should not raise.
-        mgr._check_and_trigger_am_auth(
+        mgr._trigger_am_auth(
             {"type": "Project", "id": self.PROJECT_ID}, progress_callback=mock.Mock()
         )
 
@@ -139,11 +93,11 @@ def test_runtime_error_hard_fails_with_env_var(
     ):
         mock_resolve.return_value = mock.Mock()
         mock_get.side_effect = RuntimeError("network down")
-        mgr = self._build_manager_with_sg({flow_auth.AM_READY_PROJECT_FIELD: "abc-123"})
+        mgr = ToolkitManager()
 
         with temp_env_var(TK_FLOW_AUTH_REQUIRED="1"):
             with self.assertRaises(TankBootstrapError):
-                mgr._check_and_trigger_am_auth(
+                mgr._trigger_am_auth(
                     {"type": "Project", "id": self.PROJECT_ID},
                     progress_callback=mock.Mock(),
                 )