dagster-cloud 1.12.6__py3-none-any.whl → 1.12.8__py3-none-any.whl

dagster_cloud/agent/dagster_cloud_agent.py

@@ -505,6 +505,16 @@ class DagsterCloudAgent:
                         if self._instance.user_code_launcher.agent_metrics_enabled
                         else {}
                     },
+                    {
+                        "allowed_full_deployment_locations": self._instance.allowed_full_deployment_locations
+                        if self._instance.allowed_full_deployment_locations
+                        else {},
+                    },
+                    {
+                        "allowed_branch_deployment_locations": self._instance.allowed_branch_deployment_locations
+                        if self._instance.allowed_branch_deployment_locations
+                        else [],
+                    },
                 ),
                 run_worker_statuses=run_worker_statuses_dict[deployment_name],
                 code_server_heartbeats=code_server_heartbeats_dict.get(deployment_name, []),
@@ -599,6 +609,23 @@ class DagsterCloudAgent:
 
         for entry in entries:
             location_name = entry["locationName"]
+
+            # Skip locations not in the allowed list if configured
+            if not self._instance.is_location_allowed(
+                deployment_name, location_name, is_branch_deployment
+            ):
+                deployment_type = (
+                    "for branch deployments"
+                    if is_branch_deployment
+                    else f"in deployment {deployment_name}"
+                )
+                self._logger.error(
+                    f"Skipping location {location_name} {deployment_type} - not in allowed locations list. "
+                    "Either configure this location's dagster_cloud.yaml file to use a different queue, "
+                    "or redeploy this agent with this location in the list of allowed locations."
+                )
+                continue
+
             code_location_deploy_data = deserialize_value(
                 entry["serializedDeploymentMetadata"], CodeLocationDeployData
             )
@@ -730,6 +757,22 @@ class DagsterCloudAgent:
                 for entry in entries:
                     location_name = entry["locationName"]
 
+                    # Skip locations not in the allowed list if configured
+                    if not self._instance.is_location_allowed(
+                        deployment_name, location_name, is_branch_deployment
+                    ):
+                        deployment_type = (
+                            "for branch deployments"
+                            if is_branch_deployment
+                            else f"in deployment {deployment_name}"
+                        )
+                        self._logger.error(
+                            f"Skipping location {location_name} {deployment_type} - not in allowed locations list. "
+                            "Either configure this location's dagster_cloud.yaml file to use a different queue, "
+                            "or redeploy this agent with this location in the list of allowed locations."
+                        )
+                        continue
+
                     location_key = (deployment_name, location_name)
 
                     all_locations.add(location_key)
@@ -833,6 +876,22 @@ class DagsterCloudAgent:
         code_location_origin = self._get_location_origin_from_request(request)
         location_name = code_location_origin.location_name if code_location_origin else None
 
+        # Validate that the location is in the allowed list if configured
+        if location_name and not self._instance.is_location_allowed(
+            deployment_name, location_name, is_branch_deployment
+        ):
+            deployment_type = (
+                "for branch deployments"
+                if is_branch_deployment
+                else f"in deployment '{deployment_name}'"
+            )
+            raise Exception(
+                f"Agent is not allowed to serve location '{location_name}' {deployment_type}. "
+                f"Location is not in the alloweupd locations list configured for this agent. "
+                "Either configure this location's dagster_cloud.yaml file to use a different queue, "
+                "or redeploy this agent with this location in the list of allowed locations."
+            )
+
         if api_name == DagsterCloudApi.PING_LOCATION:
             # Do nothing - this request only exists to bump TTL for the location
             return DagsterCloudApiSuccess()
@@ -1233,7 +1292,9 @@ class DagsterCloudAgent:
                     (deployment_name, location_name, is_branch_deployment)
                 ] = time.time()
 
-                if not user_code_launcher.has_grpc_endpoint(deployment_name, location_name):
+                if self._instance.is_location_allowed(
+                    deployment_name, location_name, is_branch_deployment
+                ) and not user_code_launcher.has_grpc_endpoint(deployment_name, location_name):
                     # Next completed periodic workspace update will make the location up to date
                     # - keep this in the queue until then
                     invalid_requests.append(json_request)

dagster_cloud/instance/__init__.py

@@ -10,6 +10,7 @@ import yaml
 from dagster import (
     Array,
     Field,
+    Map,
     String,
     _check as check,
 )
@@ -79,6 +80,8 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
         agent_replicas=None,
         isolated_agents=None,
         agent_queues=None,
+        allowed_full_deployment_locations=None,
+        allowed_branch_deployment_locations=None,
         agent_metrics=None,
         opentelemetry=None,
         **kwargs,
@@ -145,6 +148,13 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
         )
         self.agent_queues_config = AgentQueuesConfig(**processed_agent_queues_config)  # pyright: ignore[reportCallIssue]
 
+        self._allowed_full_deployment_locations: Optional[dict[str, list[str]]] = (
+            allowed_full_deployment_locations
+        )
+        self._allowed_branch_deployment_locations: Optional[list[str]] = (
+            allowed_branch_deployment_locations
+        )
+
         self._opentelemetry_config: Optional[Mapping[str, Any]] = self._get_processed_config(
             "opentelemetry", opentelemetry, opentelemetry_config_schema()
         )
@@ -462,6 +472,32 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
     def instance_uuid(self) -> str:
         return self._instance_uuid
 
+    @property
+    def allowed_full_deployment_locations(self) -> Optional[dict[str, list[str]]]:
+        return self._allowed_full_deployment_locations
+
+    @property
+    def allowed_branch_deployment_locations(self) -> Optional[list[str]]:
+        return self._allowed_branch_deployment_locations
+
+    def is_location_allowed(
+        self, deployment_name: str, location_name: str, is_branch_deployment: bool
+    ) -> bool:
+        """Check if a location should be allowed based on the allowed locations configuration."""
+        if is_branch_deployment:
+            # For branch deployments, check the branch deployment locations list
+            if self._allowed_branch_deployment_locations is not None:
+                return location_name in self._allowed_branch_deployment_locations
+        else:
+            # For full deployments, check the deployment-specific locations map
+            if self._allowed_full_deployment_locations is not None:
+                allowed_locations = self._allowed_full_deployment_locations.get(deployment_name)
+                if allowed_locations is not None:
+                    return location_name in allowed_locations
+
+        # If no restrictions are configured, allow all locations
+        return True
+
     @property
     def agent_display_name(self) -> str:
         if self.dagster_cloud_api_agent_label:
@@ -536,6 +572,16 @@ instance_class:
                 cls._isolated_agents_config_schema(), is_required=False
             ),  # deprecated in favor of isolated_agents
             "agent_queues": Field(cls._agent_queues_config_schema(), is_required=False),
+            "allowed_full_deployment_locations": Field(
+                Map(String, Array(String)),
+                is_required=False,
+                description="Mapping of full deployment names to allowed location names",
+            ),
+            "allowed_branch_deployment_locations": Field(
+                Array(String),
+                is_required=False,
+                description="List of allowed location names for branch deployments",
+            ),
             "opentelemetry": Field(
                 opentelemetry_config_schema(), is_required=False, default_value={"enabled": False}
             ),

dagster_cloud/version.py

@@ -1 +1 @@
-__version__ = "1.12.6"
+__version__ = "1.12.8"

dagster_cloud/workspace/user_code_launcher/utils.py

@@ -46,7 +46,14 @@ def get_instance_ref_for_user_code(instance_ref: InstanceRef) -> InstanceRef:
     if custom_instance_class_data:
         config_dict = custom_instance_class_data.config_dict
         new_config_dict = {
-            key: val for key, val in config_dict.items() if key not in {"agent_queues"}
+            key: val
+            for key, val in config_dict.items()
+            if key
+            not in {
+                "agent_queues",
+                "allowed_full_deployment_locations",
+                "allowed_branch_deployment_locations",
+            }
         }
 
         user_code_launcher_config = config_dict.get("user_code_launcher", {}).get("config")

{dagster_cloud-1.12.6.dist-info → dagster_cloud-1.12.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dagster-cloud
-Version: 1.12.6
+Version: 1.12.8
 Author-email: Elementl <support@elementl.com>
 License: Apache-2.0
 Project-URL: Homepage, https://dagster.io/cloud
@@ -28,9 +28,9 @@ Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
 Classifier: Operating System :: OS Independent
 Requires-Python: <3.14,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: dagster==1.12.6
-Requires-Dist: dagster-shared==1.12.6
-Requires-Dist: dagster-cloud-cli==1.12.6
+Requires-Dist: dagster==1.12.8
+Requires-Dist: dagster-shared==1.12.8
+Requires-Dist: dagster-cloud-cli==1.12.8
 Requires-Dist: opentelemetry-api<2,>=1.27.0
 Requires-Dist: opentelemetry-sdk<2,>=1.27.0
 Requires-Dist: opentelemetry-exporter-otlp-proto-grpc<2,>=1.27.0
@@ -65,12 +65,12 @@ Provides-Extra: insights
 Requires-Dist: pyarrow; extra == "insights"
 Provides-Extra: docker
 Requires-Dist: docker; extra == "docker"
-Requires-Dist: dagster-docker==0.28.6; extra == "docker"
+Requires-Dist: dagster-docker==0.28.8; extra == "docker"
 Provides-Extra: kubernetes
 Requires-Dist: kubernetes; extra == "kubernetes"
-Requires-Dist: dagster-k8s==0.28.6; extra == "kubernetes"
+Requires-Dist: dagster-k8s==0.28.8; extra == "kubernetes"
 Provides-Extra: ecs
-Requires-Dist: dagster-aws==0.28.6; extra == "ecs"
+Requires-Dist: dagster-aws==0.28.8; extra == "ecs"
 Requires-Dist: boto3; extra == "ecs"
 Provides-Extra: sandbox
 Requires-Dist: supervisor; extra == "sandbox"

{dagster_cloud-1.12.6.dist-info → dagster_cloud-1.12.8.dist-info}/RECORD

@@ -1,9 +1,9 @@
 dagster_cloud/__init__.py,sha256=zyM9bqyJFxtEClv_5X4VRldrj0UniKgZzEl0pPJJ_Ts,355
 dagster_cloud/constants.py,sha256=CPAqXJ99SWGMviksdIA2A9894FEvHChNk8UcP4TluYM,455
 dagster_cloud/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dagster_cloud/version.py,sha256=K-O_Zmi1DuK_mEBcDOT_b5zNp9ONlW-73xBIaMAX9Fo,23
+dagster_cloud/version.py,sha256=MukaistZHiQ_HvCd7oJvFM6KoAwAzuUNKkDHr_pCzwk,23
 dagster_cloud/agent/__init__.py,sha256=_erVyIrxuHUiyNerwX8vNZcKZN8NAloTEkPq8vPZ3MI,811
-dagster_cloud/agent/dagster_cloud_agent.py,sha256=bNMPSJBURWYbwAPhTbYq34o5qxPrqVWI2JA9b6vKHiw,58341
+dagster_cloud/agent/dagster_cloud_agent.py,sha256=3W4wdz61LEGHYiG0cY8oDCm8oVgCalM5QTU8AZJjCgA,61566
 dagster_cloud/agent/queries.py,sha256=iI84GQ1Zxt5ryo6M1ELIaIae-gwUY14QPPMUeiFK97o,1837
 dagster_cloud/agent/cli/__init__.py,sha256=rGbeQJ2Ap95wPYQuk5XbyHAfP9cs-XPUSmuVM_k278k,9084
 dagster_cloud/agent/instrumentation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -46,7 +46,7 @@ dagster_cloud/execution/cloud_run_launcher/process.py,sha256=eYxvzcaJbqYKjHcXMUU
 dagster_cloud/execution/monitoring/__init__.py,sha256=2_o2UjeEPvfSpnOb0zky3ZQs-G9VHABAXzjvfMnAh8A,17782
 dagster_cloud/execution/utils/__init__.py,sha256=EfU-tK5-ScZYLF_zmWqUgIrmLsb9DDHWFN-CSg3BSXY,254
 dagster_cloud/execution/utils/process.py,sha256=Er6FKMdlGBjpOHDEIkukZf4-TiZ3Xi_iqurfxkOozGM,911
-dagster_cloud/instance/__init__.py,sha256=ecxGScVMeqQJqegEUOzMA-LRkovHSB5LgkL4zcFyz8c,26135
+dagster_cloud/instance/__init__.py,sha256=v_9zjhcf-nXYpK9KoOBax4T1vb8rp_05U1bAi3Tgot0,28228
 dagster_cloud/instrumentation/__init__.py,sha256=OKaT4qkwQNmCFdKveVrgB7JIo5sZD1F1rutQKdozK_4,886
 dagster_cloud/metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dagster_cloud/metadata/source_code.py,sha256=yCUfD6V-wR2s4AUrmVpX_Dd7Gln7SYbtpuEpAH71J1I,5952
@@ -127,8 +127,8 @@ dagster_cloud/workspace/kubernetes/utils.py,sha256=2PrxCXew-KZmluVyW9jC9z_kuPcch
 dagster_cloud/workspace/user_code_launcher/__init__.py,sha256=E-Izs69AHPAXD9pqd3UH46l4uKxM4Lbz7y2G4KeWzEQ,880
 dagster_cloud/workspace/user_code_launcher/process.py,sha256=6TwjlXZzBlzyQEvmGzuVAUgxe_vCye0Q3aYkPPDAshs,14315
 dagster_cloud/workspace/user_code_launcher/user_code_launcher.py,sha256=BWkm4stk35HrFGydB5Cp61v-OI3z1eTZsFUec1hpGSI,103700
-dagster_cloud/workspace/user_code_launcher/utils.py,sha256=t8Epee9MrXtRhWL-b_3avXxgMGrjLScUNWtBUUGpMCg,5285
-dagster_cloud-1.12.6.dist-info/METADATA,sha256=jbYPAjaRt_dKqF0q-f8HksyqFny7utaQk0VlQMeX8-M,6582
-dagster_cloud-1.12.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dagster_cloud-1.12.6.dist-info/top_level.txt,sha256=2hMt-U33jyCgnywNrDB9Ih0EpaVmiO6dFkYcJ7Iwx4I,14
-dagster_cloud-1.12.6.dist-info/RECORD,,
+dagster_cloud/workspace/user_code_launcher/utils.py,sha256=173BNpm1ylt6euwLSnK1kpvp_cVqKJxpitetuWnXBI0,5460
+dagster_cloud-1.12.8.dist-info/METADATA,sha256=cZ3pt_nweBDoJ5EDMoN_aQDc5bvyKoQhaiQuacWXxqc,6582
+dagster_cloud-1.12.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dagster_cloud-1.12.8.dist-info/top_level.txt,sha256=2hMt-U33jyCgnywNrDB9Ih0EpaVmiO6dFkYcJ7Iwx4I,14
+dagster_cloud-1.12.8.dist-info/RECORD,,