Include Request Timeout in Dune Source

Author: bh2smith

src/config.py

@@ -13,14 +13,13 @@
 import requests
 import yaml
 from dotenv import load_dotenv
-from dune_client.query import QueryBase
 
 from src.destinations.dune import DuneDestination
 from src.destinations.postgres import PostgresDestination
 from src.interfaces import Destination, Source
 from src.job import Database, Job
 from src.logger import log
-from src.sources.dune import DuneSource, parse_query_parameters
+from src.sources.dune import DuneSource, DuneSourceConfig
 from src.sources.postgres import PostgresSource
 
 
@@ -259,15 +258,7 @@ def _build_source(
         match source.type:
             case Database.DUNE:
                 return DuneSource(
-                    api_key=source.key,
-                    query=QueryBase(
-                        query_id=int(source_config["query_id"]),
-                        params=parse_query_parameters(
-                            source_config.get("parameters", [])
-                        ),
-                    ),
-                    poll_frequency=source_config.get("poll_frequency", 1),
-                    query_engine=source_config.get("query_engine", "medium"),
+                    dune_config=DuneSourceConfig(source.key, source_config)
                 )
 
             case Database.POSTGRES:

src/destinations/dune.py

@@ -14,16 +14,11 @@ class DuneDestination(Destination[TypedDataFrame]):
 
     Attributes
     ----------
-    api_key : str
-        The API key used for accessing the Dune Analytics API.
+    client : DuneClient
+        The API for accessing the Dune Analytics.
     table_name : str
         The name of the table where the query results will be stored.
 
-    [optional]
-    request_timeout : int
-        Default: 10
-        The request timeout for the dune client.
-
     """
 
     def __init__(self, api_key: str, table_name: str, request_timeout: int):

src/sources/dune.py

@@ -5,6 +5,7 @@
 import json
 import re
 from abc import ABC
+from dataclasses import dataclass
 from typing import Any, Literal
 
 import pandas as pd
@@ -42,6 +43,28 @@
 }
 
 
+@dataclass
+class DuneSourceConfig:
+    """Represents Parameters used by Dune Client and Query execution constructors."""
+
+    api_key: str
+    query: QueryBase
+    request_timeout: int
+    poll_frequency: int
+    query_engine: Literal["medium", "large"]
+
+    def __init__(self, api_key: str, data: dict[str, Any]):
+        self.api_key = api_key
+        if data["query_id"] is not None:
+            self.query = QueryBase(
+                query_id=int(data["query_id"]),
+                params=parse_query_parameters(data.get("parameters", [])),
+            )
+        self.request_timeout = int(data.get("request_timeout", 10))
+        self.poll_frequency = int(data.get("poll_frequency", 1))
+        self.query_engine = data.get("query_engine", "medium")
+
+
 def _parse_varchar_type(type_str: str) -> int | None:
     """Extract the length from Dune's varchar type string like varchar(255).
 
@@ -240,14 +263,15 @@ class DuneSource(Source[TypedDataFrame], ABC):
 
     def __init__(
         self,
-        api_key: str,
-        query: QueryBase,
-        poll_frequency: int = 1,
-        query_engine: Literal["medium", "large"] = "medium",
+        dune_config: DuneSourceConfig,
     ) -> None:
-        self.query = query
-        self.poll_frequency = poll_frequency
-        self.client = AsyncDuneClient(api_key, performance=query_engine)
+        self.query = dune_config.query
+        self.poll_frequency = dune_config.poll_frequency
+        self.client = AsyncDuneClient(
+            dune_config.api_key,
+            performance=dune_config.query_engine,
+            request_timeout=dune_config.request_timeout,
+        )
         super().__init__()
 
     def validate(self) -> bool:

tests/unit/job_test.py

@@ -2,11 +2,10 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-from dune_client.query import QueryBase
 
 from src.destinations.postgres import PostgresDestination
 from src.job import Database, Job
-from src.sources.dune import DuneSource
+from src.sources.dune import DuneSource, DuneSourceConfig
 
 
 class DatabaseTests(unittest.IsolatedAsyncioTestCase):
@@ -20,7 +19,7 @@ def test_database_resolution(self):
         self.assertEqual("Unknown Database type: redis", exc.exception.args[0])
 
     def test_job_name_formatting(self):
-        src = DuneSource(api_key="f00b4r", query=QueryBase(query_id=1234))
+        src = DuneSource(DuneSourceConfig(api_key="f00b4r", data={"query_id": 1234}))
         dest = PostgresDestination(
             db_url="postgresql://postgres:postgres@localhost:5432/postgres",
             table_name="some_table",
@@ -31,7 +30,7 @@ def test_job_name_formatting(self):
     @pytest.mark.asyncio
     @patch("src.metrics.push_to_gateway")
     async def test_metrics_collection(self, mock_metrics_push, *_):
-        src = DuneSource(api_key="f00b4r", query=QueryBase(query_id=1234))
+        src = DuneSource(DuneSourceConfig(api_key="f00b4r", data={"query_id": 1234}))
         dest = PostgresDestination(
             db_url="postgresql://postgres:postgres@localhost:5432/postgres",
             table_name="some_table",