update airbyte server block(s) script

ddpui/ddpprefect/prefect_service.py

@@ -204,8 +204,16 @@ def create_airbyte_server_block(blockname):
 
 
 def update_airbyte_server_block(blockname):
-    """We don't update server blocks"""
-    raise Exception("not implemented")
+    """Create airbyte server block in prefect"""
+    response = prefect_put(
+        "blocks/airbyte/server/",
+        {
+            "blockName": blockname,
+            "serverHost": os.getenv("AIRBYTE_SERVER_HOST"),
+            "serverPort": os.getenv("AIRBYTE_SERVER_PORT"),
+        },
+    )
+    return (response["block_id"], response["cleaned_block_name"])
 
 
 def delete_airbyte_server_block(block_id):

ddpui/management/commands/update_airbyte_server_blocks.py

@@ -0,0 +1,33 @@
+from django.core.management.base import BaseCommand
+from ddpui.models.org import Org, OrgPrefectBlockv1
+from ddpui.ddpprefect import prefect_service
+
+
+class Command(BaseCommand):
+    """Update airbyte server blocks in prefect"""
+
+    help = "Update airbyte server blocks in prefect"
+
+    def add_arguments(self, parser):
+        """adds command line arguments"""
+        parser.add_argument("org", type=str, help="Org slug, use 'all' to update for all orgs")
+
+    def handle(self, *args, **options):
+        """Use airbyte host and port from the .env"""
+        orgs = Org.objects.all()
+        if options["org"] != "all":
+            orgs = orgs.filter(slug=options["org"])
+
+        print(f"Updating airbyte server blocks in prefect for {len(orgs)} orgs")
+
+        for org in orgs:
+            server_block = OrgPrefectBlockv1.objects.filter(org=org).first()
+
+            if not server_block:
+                print(f"Org {org.slug} does not have a server block")
+                continue
+
+            # update host and port
+            prefect_service.update_airbyte_server_block(server_block.block_name)
+
+            print(f"Updated airbyte server block for org {org.slug}")