Add Type Map and Fix Tests

Author: bh2smith

config.yaml

@@ -36,7 +36,7 @@ jobs:
       query_string: "SELECT 1 as number, '\\x1234'::bytea as my_bytes;"
     destination:
       ref: Dune
-      table_name: dune_sync_test_table
+      table_name: bh2smith.dune_sync_test
 
   - name: cow-solvers
     source:

src/destinations/dune.py

@@ -1,5 +1,7 @@
 """Destination logic for Dune Analytics."""
 
+import io
+
 from dune_client.client import DuneClient
 from dune_client.models import DuneError
 
@@ -59,30 +61,34 @@ def save(self, data: TypedDataFrame) -> int:
 
         """
         try:
-            # TODO: Determine user name from DuneAPI key?
-            namespace = "username"
-            table_name = self.table_name
             log.debug("Uploading DF to Dune...")
+            namespace, table_name = self._get_namespace_and_table_name()
             # TODO check first if table exists? Or warn if it did...
             self.client.create_table(
                 namespace,
                 table_name,
                 schema=[
-                    {"name": name, "type": dtype, "nullable": "true"}
-                    for name, dtype in data.types.items()
+                    {"name": name, "type": dtype} for name, dtype in data.types.items()
                 ],
             )
             result = self.client.insert_table(
                 namespace,
                 table_name,
-                # TODO - bytes -> IO[bytes]
-                data=data.dataframe.to_csv(index=False),  # type: ignore
+                data=io.BytesIO(data.dataframe.to_csv(index=False).encode()),
                 content_type="text/csv",
             )
             if not result:
                 raise RuntimeError("Dune Upload Failed")
+            log.debug("Inserted DF to Dune, %s", result)
         except DuneError as dune_e:
             log.error("Dune did not accept our upload: %s", dune_e)
         except (ValueError, RuntimeError) as e:
             log.error("Data processing error: %s", e)
         return len(data)
+
+    def _get_namespace_and_table_name(self) -> tuple[str, str]:
+        """Split the namespace, table name from the provided table name."""
+        if "." not in self.table_name:
+            raise ValueError("Table name must be in the format namespace.table_name")
+        namespace, table_name = self.table_name.split(".")
+        return namespace, table_name

src/sources/postgres.py

@@ -11,6 +11,7 @@
 
 from src.interfaces import Source, TypedDataFrame
 from src.logger import log
+from src.sources.type_maps import PG_TO_DUNE
 
 
 def _convert_bytea_to_hex(df: DataFrame) -> DataFrame:
@@ -118,11 +119,17 @@ async def fetch(self) -> TypedDataFrame:
         # of SQLAlchemy's synchronous interface.
         # The current solution using run_in_executor is a workaround
         # that moves the blocking operation to a thread pool.
+        # First get the column types
+        with self.engine.connect() as conn:
+            result = conn.execute(text(self.query_string))
+            types = {
+                col.name: PG_TO_DUNE[col.type_code] for col in result.cursor.description
+            }
         df = await loop.run_in_executor(
             None, lambda: pd.read_sql_query(self.query_string, con=self.engine)
         )
         # TODO - extract types and return TypedDataFrame.
-        return TypedDataFrame(dataframe=_convert_bytea_to_hex(df), types={})
+        return TypedDataFrame(dataframe=_convert_bytea_to_hex(df), types=types)
 
     def is_empty(self, data: TypedDataFrame) -> bool:
         """Check if the provided TypedDataFrame is empty.

src/sources/type_maps.py

@@ -33,14 +33,19 @@
 # 1. Notice `DOUBLE_PRECISION` has two pre-images: we chose double
 # 2. timestamp with time zone not aligned with timestamp
 # 3. Apparently no JSONB support here.
-PG_TO_DUNE: dict[type[Any] | NUMERIC, str] = {
-    BIGINT: "bigint",
-    INTEGER: "integer",
-    BYTEA: "varbinary",
-    DATE: "date",
-    BOOLEAN: "boolean",
-    VARCHAR: "varchar",
-    DOUBLE_PRECISION: "double",
-    TIMESTAMP: "timestamp",  # This doesn't match with above
-    NUMERIC: "uint256",
+PG_TO_DUNE: dict[int, str] = {
+    16: "boolean",
+    17: "varbinary",
+    20: "bigint",
+    21: "bigint",  # smallint
+    23: "integer",
+    25: "varchar",
+    # 26: "oid",
+    700: "double",
+    701: "double",
+    1042: "varchar",
+    1043: "varchar",
+    1082: "timestamp",
+    1114: "timestamp",  # This doesn't match with above
+    1700: "uint256",
 }

tests/unit/destinations_test.py

@@ -1,6 +1,6 @@
 import os
 import unittest
-from logging import ERROR, WARNING
+from logging import DEBUG, ERROR, WARNING
 from unittest.mock import patch
 
 import pandas as pd
@@ -26,12 +26,10 @@ def setUpClass(cls):
         )
         cls.env_patcher.start()
 
-    @patch("pandas.core.generic.NDFrame.to_csv", name="Fake csv writer")
     @patch("dune_client.api.table.TableAPI.create_table", name="Fake Table Creator")
     @patch("dune_client.api.table.TableAPI.insert_table", name="Fake Table Inserter")
-    @patch("requests.sessions.Session.post")
     def test_ensure_index_disabled_when_uploading(
-        self, mock_to_csv, mock_create_table, mock_insert_table, *_
+        self, mock_create_table, mock_insert_table, *_
     ):
         mock_create_table.return_value = {
             "namespace": "my_user",
@@ -41,24 +39,28 @@ def test_ensure_index_disabled_when_uploading(
             "already_existed": False,
             "message": "Table created successfully",
         }
+
         mock_insert_table.return_value = {"rows_written": 9000, "bytes_written": 90}
 
         dummy_df = TypedDataFrame(
             dataframe=pd.DataFrame(
                 [
-                    {"foo": "bar"},
-                    {"baz": "daz"},
+                    {"foo": "bar", "baz": "one"},
+                    {"foo": "two", "baz": "two"},
                 ]
             ),
             types={"foo": "varchar", "baz": "varchar"},
         )
         destination = DuneDestination(
             api_key=os.getenv("DUNE_API_KEY"),
-            table_name="foo",
+            table_name="foo.bar",
             request_timeout=10,
         )
-        destination.save(dummy_df)
-        mock_to_csv.assert_called_once_with(index=False)
+        with self.assertLogs(level=DEBUG) as logs:
+            destination.save(dummy_df)
+
+        self.assertIn("Uploading DF to Dune", logs.output[0])
+        self.assertIn("Inserted DF to Dune,", logs.output[1])
 
     @patch("pandas.core.generic.NDFrame.to_csv", name="Fake csv writer")
     def test_duneclient_sets_timeout(self, mock_to_csv, *_):
@@ -70,13 +72,12 @@ def test_duneclient_sets_timeout(self, mock_to_csv, *_):
             )
             assert destination.client.request_timeout == timeout
 
-    @patch("dune_client.api.table.TableAPI.upload_csv", name="Fake CSV uploader")
     @patch("dune_client.api.table.TableAPI.create_table", name="Fake Table Creator")
     @patch("dune_client.api.table.TableAPI.insert_table", name="Fake Table Inserter")
-    def test_dune_error_handling(
-        self, mock_dune_upload_csv, mock_create_table, mock_insert_table
-    ):
-        dest = DuneDestination(api_key="f00b4r", table_name="foo", request_timeout=10)
+    def test_dune_error_handling(self, mock_create_table, mock_insert_table):
+        dest = DuneDestination(
+            api_key="f00b4r", table_name="foo.bar", request_timeout=10
+        )
         df = pd.DataFrame([{"foo": "bar"}])
 
         mock_create_table.return_value = {
@@ -96,13 +97,13 @@ def test_dune_error_handling(
         val_err = ValueError("Oops")
         runtime_err = RuntimeError("Big Oops")
 
-        mock_dune_upload_csv.side_effect = dune_err
+        mock_create_table.side_effect = dune_err
 
         data = TypedDataFrame(df, {})
         with self.assertLogs(level=ERROR) as logs:
             dest.save(data)
 
-        mock_dune_upload_csv.assert_called_once()
+        mock_create_table.assert_called_once()
 
         # does this shit really look better just because it's < 88 characters long?
         exmsg = (
@@ -111,36 +112,36 @@ def test_dune_error_handling(
         )
         self.assertIn(exmsg, logs.output[0])
 
-        mock_dune_upload_csv.reset_mock()
-        mock_dune_upload_csv.side_effect = val_err
+        mock_create_table.reset_mock()
+        mock_create_table.side_effect = val_err
 
         with self.assertLogs(level=ERROR) as logs:
             dest.save(data)
 
-        mock_dune_upload_csv.assert_called_once()
+        mock_create_table.assert_called_once()
         expected_message = "Data processing error: Oops"
         self.assertIn(expected_message, logs.output[0])
 
-        mock_dune_upload_csv.reset_mock()
-        mock_dune_upload_csv.side_effect = runtime_err
+        mock_create_table.reset_mock()
+        mock_create_table.side_effect = runtime_err
         with self.assertLogs(level=ERROR) as logs:
             dest.save(data)
 
-        mock_dune_upload_csv.assert_called_once()
+        mock_create_table.assert_called_once()
         expected_message = "Data processing error: Big Oops"
         self.assertIn(expected_message, logs.output[0])
 
-        mock_dune_upload_csv.reset_mock()
+        mock_create_table.reset_mock()
 
         # TIL: reset_mock() doesn't clear side effects....
-        mock_dune_upload_csv.side_effect = None
+        mock_create_table.side_effect = None
 
-        mock_dune_upload_csv.return_value = None
+        mock_create_table.return_value = None
 
         with self.assertLogs(level=ERROR) as logs:
             dest.save(data)
 
-        mock_dune_upload_csv.assert_called_once()
+        mock_create_table.assert_called_once()
         self.assertIn("Dune Upload Failed", logs.output[0])