goblet · anovis · Mar 18, 2024 · Mar 10, 2024 · Mar 10, 2024 · Mar 10, 2024
diff --git a/goblet/__version__.py b/goblet/__version__.py
@@ -1,4 +1,4 @@
-VERSION = (0, 13, 0)
+VERSION = (0, 13, 1)
 
 
 __version__ = ".".join(map(str, VERSION))
diff --git a/goblet/decorators.py b/goblet/decorators.py
@@ -18,6 +18,8 @@
     BigQuerySparkStoredProcedure,
 )
 
+from functools import partial
+
 log = logging.getLogger(__name__)
 log.setLevel(logging.getLevelName(os.getenv("GOBLET_LOG_LEVEL", "INFO")))
 
@@ -140,6 +142,8 @@ def pubsub_subscription(self, topic, **kwargs):
         dlq = kwargs.pop("dlq", False)
         dlq_topic_config = kwargs.pop("dlq_topic_config", {})
         dlq_alerts = kwargs.pop("dlq_alerts", [])
+
+        extra_registrations = []
         if dlq:
             log.info(f"DLQ enabled use of subscription will be forced to topic {topic}")
             kwargs["use_subscription"] = True
@@ -157,25 +161,26 @@ def pubsub_subscription(self, topic, **kwargs):
                 else dlq_pull_subscription_config.pop("name")
             )
             # Create DLQ topic
-            self._register_infrastructure(
-                handler_type="pubsub_topic",
-                kwargs={
-                    "name": dlq_topic_name,
-                    "kwargs": {
-                        "config": dlq_topic_config,
-                        "dlq": True,
-                        "dlq_pull_subscription": {
-                            "name": dlq_pull_subscription_name,
-                            "config": dlq_pull_subscription_config,
+            extra_registrations.append(
+                partial(
+                    self._register_infrastructure,
+                    handler_type="pubsub_topic",
+                    kwargs={
+                        "name": dlq_topic_name,
+                        "kwargs": {
+                            "config": dlq_topic_config,
+                            "dlq": True,
+                            "dlq_pull_subscription": {
+                                "name": dlq_pull_subscription_name,
+                                "config": dlq_pull_subscription_config,
+                            },
                         },
                     },
-                },
+                )
             )
             dlq_policy = {
                 "deadLetterPolicy": {
-                    "deadLetterTopic": self.infrastructure["pubsub_topic"].resources[
-                        dlq_topic_name
-                    ]["name"],
+                    "deadLetterTopic": f"projects/goblet/subscriptions/topics/{dlq_topic_name}"
                 }
             }
             if "config" in kwargs:
@@ -185,11 +190,12 @@ def pubsub_subscription(self, topic, **kwargs):
 
             for dlq_alert in dlq_alerts:
                 dlq_alert.update_extras({"topic": dlq_topic_name})
-                self.alerts.register(dlq_alert)
+                extra_registrations.append(partial(self.alerts.register, dlq_alert))
 
         return self._create_registration_function(
             handler_type="pubsub",
             registration_kwargs={"topic": topic, "kwargs": kwargs},
+            extra_registrations=extra_registrations,
         )
 
     def topic(self, topic, **kwargs):
@@ -261,29 +267,34 @@ def http(self, headers={}):
 
     def job(self, name, task_id=0, schedule=None, timezone="UTC", **kwargs):
         """Cloudrun Job"""
+        extra_registrations = []
         if schedule and task_id != 0:
             raise ValueError("Schedule can only be added to task_id with value 0")
         if kwargs and task_id != 0:
             raise ValueError("Arguments can only be added to task_id with value 0")
         if schedule:
-            self._register_handler(
-                "schedule",
-                f"schedule-job-{name}",
-                None,
-                kwargs={
-                    "schedule": schedule,
-                    "timezone": timezone,
-                    "kwargs": {
-                        "uri": f"https://{get_default_location()}-run.googleapis.com/apis/run.googleapis.com/v1/namespaces/{get_default_project()}/jobs/{self.function_name}-{name}:run",
-                        "httpMethod": "POST",
-                        "authMethod": "oauthToken",
-                        **kwargs,
+            extra_registrations.append(
+                partial(
+                    self._register_handler,
+                    "schedule",
+                    f"schedule-job-{name}",
+                    None,
+                    kwargs={
+                        "schedule": schedule,
+                        "timezone": timezone,
+                        "kwargs": {
+                            "uri": f"https://{get_default_location()}-run.googleapis.com/apis/run.googleapis.com/v1/namespaces/{get_default_project()}/jobs/{self.function_name}-{name}:run",
+                            "httpMethod": "POST",
+                            "authMethod": "oauthToken",
+                            **kwargs,
+                        },
                     },
-                },
+                )
             )
         return self._create_registration_function(
             handler_type="jobs",
             registration_kwargs={"name": name, "task_id": task_id, "kwargs": kwargs},
+            extra_registrations=extra_registrations,
         )
 
     def apigateway(self, name, backend_url, filename=None, openapi_dict=None, **kwargs):
@@ -414,12 +425,17 @@ def _register_stage(func):
 
         return _register_stage
 
-    def _create_registration_function(self, handler_type, registration_kwargs=None):
+    def _create_registration_function(
+        self, handler_type, extra_registrations=[], registration_kwargs=None
+    ):
         def _register_handler(user_handler):
             if user_handler:
                 handler_name = user_handler.__name__
                 kwargs = registration_kwargs or {}
                 self._register_handler(handler_type, handler_name, user_handler, kwargs)
+                # Deploy extra registrations correctly with stage decorator
+                for extra in extra_registrations:
+                    extra()
             return user_handler
 
         return _register_handler

diff --git a/goblet/tests/test_bqsparkstoredprocedure.py b/goblet/tests/test_bqsparkstoredprocedure.py
@@ -101,7 +101,7 @@ def spark_handler():
             routine_response["body"]["sparkOptions"]["connection"]
             == connection_response["body"]["name"]
         )
-        assert 3 == get_replay_count()
+        assert 2 == get_replay_count()
 
     def test_destroy_bqsparkstoredprocedure(self, monkeypatch):
         test_deploy_name = "bqsparkstoredprocedure-destroy"
@@ -128,7 +128,7 @@ def spark_handler():
         responses = get_responses(test_deploy_name)
 
         assert len(responses) != 0
-        assert 3 == get_replay_count()
+        assert 2 == get_replay_count()
 
     def test_deploy_bqsparkstoredprocedure_remote_code(self, monkeypatch):
         test_name = "bqsparkstoredprocedure-remote-deploy"
@@ -186,7 +186,7 @@ def main():
             routine_response["body"]["sparkOptions"]["mainFileUri"]
             == f"gs://{test_name}/spark.py"
         )
-        assert 5 == get_replay_count()
+        assert 4 == get_replay_count()
         os.remove("spark.py")
 
     def test_destroy_bqsparkstoredprocedure_remote_code(self, monkeypatch):
@@ -212,4 +212,4 @@ def test_destroy_bqsparkstoredprocedure_remote_code(self, monkeypatch):
         responses = get_responses(test_deploy_name)
 
         assert len(responses) != 0
-        assert 6 == get_replay_count()
+        assert 5 == get_replay_count()
diff --git a/goblet/tests/test_jobs.py b/goblet/tests/test_jobs.py
@@ -224,3 +224,41 @@ def test_deploy_jobs_from_artifact_tag(self, monkeypatch):
                 0
             ]["image"]
         )
+
+    def test_schedule_with_stages(self, monkeypatch):
+        monkeypatch.setenv("STAGE", "TEST")
+
+        app = Goblet("test", backend="cloudrun", config={"stages": {"TEST": {}}})
+
+        @app.job("testjob1", schedule="* * * * *")
+        @app.stage("TEST")
+        def dummy_function():
+            return "test"
+
+        @app.job("testjob2", schedule="* * * * *")
+        @app.stage("TEST2")
+        def dummy_function2():
+            return "test"
+
+        @app.job("testjob3", schedule="* * * * *")
+        @app.stage(stages=["TEST", "TEST2"])
+        def dummy_function3():
+            return "test"
+
+        assert list(app.handlers["schedule"].resources.keys()) == [
+            "schedule-job-testjob1",
+            "schedule-job-testjob3",
+        ]
+
+    def test_schedule_without_stages(self, monkeypatch):
+        monkeypatch.setenv("STAGE", "TEST")
+
+        app = Goblet("test", backend="cloudrun", config={"stages": {"TEST": {}}})
+
+        @app.job("testjob1", schedule="* * * * *")
+        def dummy_function():
+            return "test"
+
+        assert list(app.handlers["schedule"].resources.keys()) == [
+            "schedule-job-testjob1"
+        ]
diff --git a/goblet/tests/test_pubsub_subscription.py b/goblet/tests/test_pubsub_subscription.py
@@ -570,6 +570,42 @@ def test_deploy_pubsub_subscription_with_dlq(self, monkeypatch):
 
         assert get_replay_count() == 15
 
+    def test_deploy_pubsub_subscription_with_dlq_with_stage(self, monkeypatch):
+        monkeypatch.setenv("GOOGLE_PROJECT", "goblet")
+        monkeypatch.setenv("GOOGLE_LOCATION", "us-central1")
+        monkeypatch.setenv("G_TEST_NAME", "pubsub-deploy-subscription-dlq")
+        monkeypatch.setenv("G_HTTP_TEST", "REPLAY")
+        service_account = "[email protected]"
+
+        app = Goblet(
+            function_name="goblet-topic-subscription",
+            config={"pubsub": {"serviceAccountEmail": service_account}},
+        )
+        setattr(app, "entrypoint", "app")
+
+        @app.pubsub_subscription(
+            "test",
+            dlq=True,
+            dlq_alerts=[
+                PubSubDLQAlert(
+                    "pubsubdlq",
+                    conditions=[
+                        PubSubDLQCondition(
+                            "pubsubdlq",
+                            subscription_id="pubsub-deploy-subscription",
+                        )
+                    ],
+                )
+            ],
+        )
+        @app.stage("TEST")
+        def dummy():
+            return
+
+        assert len(app.alerts.resources) == 0
+        assert len(app.infrastructure["pubsub_topic"].resources) == 0
+        assert len(app.handlers["pubsub"].resources) == 0
+
     def test_deploy_local(self, monkeypatch):
         monkeypatch.setenv("GOOGLE_PROJECT", "goblet")
         monkeypatch.setenv("GOOGLE_LOCATION", "us-central1")

diff --git a/utils/schema/build.py b/utils/schema/build.py
@@ -11,7 +11,7 @@ def get_references(d):
         if isinstance(v, dict):
             get_references(v)
         else:
-            if k == "$ref" and v.startswith("https://") and 'googleapis.com' in v:
+            if k == "$ref" and v.startswith("https://") and "googleapis.com" in v:
                 references.add(v.split("#")[0])