Add support for Dremio integration

apollo/agent/proxy_client_factory.py

@@ -235,6 +235,16 @@ def _get_proxy_client_msk_kafka(
     return MskKafkaProxyClient(credentials=credentials, platform=platform)
 
 
+def _get_proxy_client_dremio(
+    credentials: Optional[Dict], platform: str, **kwargs  # type: ignore
+) -> BaseProxyClient:
+    from apollo.integrations.db.dremio_proxy_client import (
+        DremioProxyClient,
+    )
+
+    return DremioProxyClient(credentials=credentials, platform=platform)
+
+
 @dataclass
 class ProxyClientCacheEntry:
     created_time: datetime
@@ -267,6 +277,7 @@ class ProxyClientCacheEntry:
     "hive": _get_proxy_client_hive,
     "msk-connect": _get_proxy_client_msk_connect,
     "msk-kafka": _get_proxy_client_msk_kafka,
+    "dremio": _get_proxy_client_dremio,
 }
 
 

apollo/integrations/db/dremio_proxy_client.py

@@ -0,0 +1,91 @@
+from typing import Optional, Dict, Any, List, Union, Tuple
+
+from pyarrow.flight import FlightDescriptor, FlightCallOptions
+
+from apollo.integrations.db.base_db_proxy_client import BaseDbProxyClient
+
+from pyarrow import flight, Schema, types as pyarrow_types
+
+_ATTR_CONNECT_ARGS = "connect_args"
+
+
+class DremioProxyClient(BaseDbProxyClient):
+    """
+    Proxy client for Dremio. Credentials are expected to be supplied under "connect_args" and
+    will be passed directly to `pyarrow.flight.connect`, so only attributes supported as parameters by
+    `flight.connect` should be passed.
+    """
+
+    def __init__(self, credentials: Optional[Dict], **kwargs: Any):
+        super().__init__(connection_type="dremio")
+        if not credentials or _ATTR_CONNECT_ARGS not in credentials:
+            raise ValueError(
+                f"Dremio agent client requires {_ATTR_CONNECT_ARGS} in credentials"
+            )
+        self._connection = flight.connect(**credentials[_ATTR_CONNECT_ARGS])  # type: ignore
+        self._headers = [
+            (
+                b"authorization",
+                f"bearer {credentials.get('token')}".encode("utf-8"),
+            )
+        ]
+
+    @property
+    def wrapped_client(self):
+        return self._connection
+
+    def execute(self, query: str) -> Dict:
+        # Send the query to the server and save the ticket identifier
+        flight_info = self._connection.get_flight_info(
+            FlightDescriptor.for_command(command=query),
+            FlightCallOptions(headers=self._headers),
+        )
+        ticket = flight_info.endpoints[0].ticket
+
+        # create a FlightStreamReader from the ticket
+        reader = self._connection.do_get(
+            ticket, FlightCallOptions(headers=self._headers)
+        )
+
+        # read all results of the query
+        results = reader.read_all()
+
+        # turn results into a serializable list
+        records = [
+            list(row)
+            for row in zip(*[column.to_pylist() for column in results.itercolumns()])
+        ]
+
+        return {
+            "records": records,
+            "description": self._get_dbapi_description(reader.schema),
+            "rowcount": len(records),
+        }
+
+    def _get_dbapi_description(self, schema: Schema) -> List[Tuple]:
+        """
+        Create a DB API compliant description object from the FlightStreamReader schema
+        """
+
+        def arrow_type_to_cursor_type(arrow_type: Any) -> Union[int, str]:
+            if pyarrow_types.is_string(arrow_type):
+                return 1
+            elif pyarrow_types.is_int32(arrow_type):
+                return 2
+            elif pyarrow_types.is_float32(arrow_type):
+                return 3
+            # Add more type mappings as needed
+            return "unknown"
+
+        return [
+            (
+                field.name,
+                arrow_type_to_cursor_type(field.type),
+                None,
+                None,
+                None,
+                None,
+                None,
+            )
+            for field in schema
+        ]

tests/test_dremio_client.py

@@ -0,0 +1,82 @@
+from unittest import TestCase
+from unittest.mock import patch, MagicMock
+
+from apollo.integrations.db.dremio_proxy_client import DremioProxyClient
+
+import pyarrow
+
+
+class DremioClientTests(TestCase):
+    def setUp(self):
+        self.credentials = {
+            "connect_args": {"location": "localhost"},
+            "token": "dummy_token",
+        }
+        self.query = "SELECT * FROM test_table"
+
+    @patch("apollo.integrations.db.dremio_proxy_client.flight.connect")
+    def test_init_success(self, mock_connect):
+        mock_connection = MagicMock()
+        mock_connect.return_value = mock_connection
+
+        client = DremioProxyClient(credentials=self.credentials)
+
+        self.assertEqual(client._connection, mock_connection)
+        self.assertEqual(client._headers, [(b"authorization", b"bearer dummy_token")])
+        mock_connect.assert_called_once_with(location="localhost")
+
+    def test_init_missing_connect_args(self):
+        with self.assertRaises(ValueError):
+            DremioProxyClient(credentials={})
+
+    @patch("apollo.integrations.db.dremio_proxy_client.flight.connect")
+    @patch("apollo.integrations.db.dremio_proxy_client.FlightCallOptions")
+    def test_execute(self, mock_flight_call_options, mock_connect):
+        mock_connection = MagicMock()
+        mock_flight_info = MagicMock()
+        mock_flight_call_options_instance = MagicMock()
+
+        mock_flight_call_options.return_value = mock_flight_call_options_instance
+
+        mock_connect.return_value = mock_connection
+        mock_connection.get_flight_info.return_value = mock_flight_info
+        mock_flight_info.endpoints = [MagicMock(ticket="dummy_ticket")]
+
+        mock_reader = MagicMock()
+        mock_reader.read_all.return_value = pyarrow.table(
+            {"col1": [1, 2, 3], "col2": ["a", "b", "c"]}
+        )
+
+        mock_reader.schema = pyarrow.schema(
+            [("col1", pyarrow.int32()), ("col2", pyarrow.string())]
+        )
+        mock_connection.do_get.return_value = mock_reader
+
+        client = DremioProxyClient(credentials=self.credentials)
+
+        result = client.execute(self.query)
+
+        self.assertEqual(result["records"], [[1, "a"], [2, "b"], [3, "c"]])
+        self.assertEqual(result["rowcount"], 3)
+        self.assertEqual(
+            result["description"],
+            [
+                ("col1", 2, None, None, None, None, None),
+                ("col2", 1, None, None, None, None, None),
+            ],
+        )
+
+    @patch("apollo.integrations.db.dremio_proxy_client.flight.connect")
+    def test_get_description(self, mock_connect):
+        schema = pyarrow.schema([("col1", pyarrow.int32()), ("col2", pyarrow.string())])
+
+        client = DremioProxyClient(credentials=self.credentials)
+        description = client._get_dbapi_description(schema)
+
+        self.assertEqual(
+            description,
+            [
+                ("col1", 2, None, None, None, None, None),
+                ("col2", 1, None, None, None, None, None),
+            ],
+        )