updated streams for old configs with granularity

airbytehq · May 23, 2024 · c6130f3 · c6130f3
1 parent 7b8687f
commit c6130f3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/airbyte-integrations/connectors/source-tiktok-marketing/source_tiktok_marketing/source.py b/airbyte-integrations/connectors/source-tiktok-marketing/source_tiktok_marketing/source.py
@@ -39,6 +39,7 @@
     "creative_assets_portfolios",
     "creative_assets_videos",
 ]
+REPORT_GRANULARITY = {"DAY": "daily", "HOUR": "hourly", "LIFETIME": "lifetime"}
 
 
 class SourceTiktokMarketing(YamlDeclarativeSource):
@@ -57,11 +58,19 @@ def _is_sandbox(config: Mapping[str, Any]) -> bool:
         return is_sandbox
 
     def streams(self, config: Mapping[str, Any]) -> List[Stream]:
-        is_production = not self._is_sandbox(config)
+        granularity = config.get("report_granularity")
         streams = super().streams(config)
 
-        if is_production:
-            return streams
+        if self._is_sandbox(config):
+            streams = [stream for stream in streams if stream.name in SANDBOX_STREAM_NAMES]
 
-        sandbox_streams = [stream for stream in streams if stream.name in SANDBOX_STREAM_NAMES]
-        return sandbox_streams
+        if granularity:
+            granularity_streams = []
+            for stream in streams:
+                if "report" not in stream.name or REPORT_GRANULARITY[granularity] in stream.name:
+                    # old configs with provided report granularity don't have granularity in stream name
+                    stream.name = stream.name.replace(f"_{REPORT_GRANULARITY[granularity]}", "")
+                    granularity_streams.append(stream)
+            return granularity_streams
+
+        return streams
diff --git a/airbyte-integrations/connectors/source-tiktok-marketing/unit_tests/test_source.py b/airbyte-integrations/connectors/source-tiktok-marketing/unit_tests/test_source.py
@@ -11,13 +11,17 @@
 
 SANDBOX_CONFIG_FILE = "secrets/sandbox_config.json"
 PROD_CONFIG_FILE = "secrets/prod_config.json"
+PROD_LIFETIME_CONFIG_FILE = "secrets/prod_config_with_lifetime_granularity.json"
+PROD_DAILY_CONFIG_FILE = "secrets/prod_config_with_day_granularity.json"
 
 
 @pytest.mark.parametrize(
     "config, stream_len",
     [
         (PROD_CONFIG_FILE, 36),
         (SANDBOX_CONFIG_FILE, 28),
+        (PROD_LIFETIME_CONFIG_FILE, 15),
+        (PROD_DAILY_CONFIG_FILE, 27),
     ],
 )
 def test_source_streams(config, stream_len):