From 1ccf8aa20ac953071fbadc32cdce88482b1dd808 Mon Sep 17 00:00:00 2001
From: Maxime Armstrong <maxime@retain.ai>
Date: Fri, 8 Nov 2024 14:32:00 -0500
Subject: [PATCH] Update fetch_fivetran_workspace_data

---
 .../dagster_fivetran/resources.py                | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/python_modules/libraries/dagster-fivetran/dagster_fivetran/resources.py b/python_modules/libraries/dagster-fivetran/dagster_fivetran/resources.py
index 8cd31ae02718a..0987ec770054a 100644
--- a/python_modules/libraries/dagster-fivetran/dagster_fivetran/resources.py
+++ b/python_modules/libraries/dagster-fivetran/dagster_fivetran/resources.py
@@ -3,6 +3,7 @@
 import logging
 import os
 import time
+from enum import Enum
 from typing import Any, Mapping, Optional, Sequence, Tuple
 from urllib.parse import urljoin
 
@@ -43,6 +44,14 @@
 DEFAULT_POLL_INTERVAL = 10
 
 
+class FivetranConnectorSetupStateType(Enum):
+    """Enum representing each setup state for a connector in Fivetran's ontology."""
+
+    INCOMPLETE = "incomplete"
+    CONNECTED = "connected"
+    BROKEN = "broken"
+
+
 class FivetranResource(ConfigurableResource):
     """This class exposes methods on top of the Fivetran REST API."""
 
@@ -643,8 +652,11 @@ def fetch_fivetran_workspace_data(
             for connector_details in connectors_details:
                 connector_id = connector_details["id"]
 
-                setup_state = connector_details.get("status", {}).get("setup_state")
-                if setup_state and setup_state in ("incomplete", "broken"):
+                setup_state = connector_details["status"]["setup_state"]
+                if setup_state in (
+                    FivetranConnectorSetupStateType.INCOMPLETE,
+                    FivetranConnectorSetupStateType.BROKEN,
+                ):
                     continue
 
                 schema_config = client.get_schema_config_for_connector(connector_id=connector_id)