fix: results for handlers

Author: inishchith

application_sdk/clients/sql.py

@@ -404,7 +404,7 @@ async def run_query(self, query: str, batch_size: int = 100000):
         logger.info("Query execution completed")
 
     def _execute_pandas_query(
-        self, conn, query
+        self, conn, query, chunksize: Optional[int]
     ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
         """Helper function to execute SQL query using pandas.
            The function is responsible for using import_optional_dependency method of the pandas library to import sqlalchemy
@@ -424,13 +424,13 @@ def _execute_pandas_query(
         from sqlalchemy import text
 
         if import_optional_dependency("sqlalchemy", errors="ignore"):
-            return pd.read_sql_query(text(query), conn, chunksize=self.chunk_size)
+            return pd.read_sql_query(text(query), conn, chunksize=chunksize)
         else:
             dbapi_conn = getattr(conn, "connection", None)
-            return pd.read_sql_query(query, dbapi_conn, chunksize=self.chunk_size)
+            return pd.read_sql_query(query, dbapi_conn, chunksize=chunksize)
 
     def _read_sql_query(
-        self, session: "Session", query: str
+        self, session: "Session", query: str, chunksize: Optional[int]
     ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
         """Execute SQL query using the provided session.
 
@@ -442,10 +442,10 @@ def _read_sql_query(
                 or iterator of DataFrames if chunked.
         """
         conn = session.connection()
-        return self._execute_pandas_query(conn, query)
+        return self._execute_pandas_query(conn, query, chunksize=chunksize)
 
     def _execute_query_daft(
-        self, query: str
+        self, query: str, chunksize: Optional[int]
     ) -> Union["daft.DataFrame", Iterator["daft.DataFrame"]]:
         """Execute SQL query using the provided engine and daft.
 
@@ -462,15 +462,11 @@ def _execute_query_daft(
             raise ValueError("Engine is not initialized. Call load() first.")
 
         if isinstance(self.engine, str):
-            return daft.read_sql(
-                query, self.engine, infer_schema_length=self.chunk_size
-            )
-        return daft.read_sql(
-            query, self.engine.connect, infer_schema_length=self.chunk_size
-        )
+            return daft.read_sql(query, self.engine, infer_schema_length=chunksize)
+        return daft.read_sql(query, self.engine.connect, infer_schema_length=chunksize)
 
     def _execute_query(
-        self, query: str
+        self, query: str, chunksize: Optional[int]
     ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
         """Execute SQL query using the provided engine and pandas.
 
@@ -482,7 +478,7 @@ def _execute_query(
             raise ValueError("Engine is not initialized. Call load() first.")
 
         with self.engine.connect() as conn:
-            return self._execute_pandas_query(conn, query)
+            return self._execute_pandas_query(conn, query, chunksize)
 
     async def get_batched_results(
         self,
@@ -513,12 +509,14 @@ async def get_batched_results(
 
             if async_session:
                 async with async_session() as session:
-                    return await session.run_sync(self._read_sql_query, query)
+                    return await session.run_sync(
+                        self._read_sql_query, query, chunksize=self.chunk_size
+                    )
             else:
                 # Run the blocking operation in a thread pool
                 with concurrent.futures.ThreadPoolExecutor() as executor:
                     return await asyncio.get_event_loop().run_in_executor(  # type: ignore
-                        executor, self._execute_query, query
+                        executor, self._execute_query, query, self.chunk_size
                     )
         except Exception as e:
             logger.error(f"Error reading batched data(pandas) from SQL: {str(e)}")
@@ -549,12 +547,14 @@ async def get_results(self, query: str) -> "pd.DataFrame":
 
             if async_session:
                 async with async_session() as session:
-                    return await session.run_sync(self._read_sql_query, query)
+                    return await session.run_sync(
+                        self._read_sql_query, query, chunksize=None
+                    )
             else:
                 # Run the blocking operation in a thread pool
                 with concurrent.futures.ThreadPoolExecutor() as executor:
                     result = await asyncio.get_event_loop().run_in_executor(
-                        executor, self._execute_query, query
+                        executor, self._execute_query, query, None
                     )
                     import pandas as pd
 

tests/unit/activities/metadata_extraction/test_sql.py

@@ -1,5 +1,6 @@
 """Unit tests for SQL metadata extraction activities (context-free)."""
 
+import os
 from typing import Any, Dict
 from unittest.mock import AsyncMock, Mock, patch
 
@@ -370,7 +371,9 @@ async def test_transform_data_success(
         assert result is not None
         assert isinstance(result, ActivityStatistics)
         assert result.total_record_count == 20
-        mock_download_files.assert_called_once_with("/test/path/raw", ".parquet", None)
+        # Normalize path for cross-platform compatibility
+        expected_path = os.path.join("/test/path", "raw")
+        mock_download_files.assert_called_once_with(expected_path, ".parquet", None)
         mock_transform_metadata.assert_called_once()
         mock_write_daft_dataframe.assert_called_once()
 

tests/unit/io/readers/test_json_reader.py

@@ -63,7 +63,12 @@ async def test_download_file_invoked_for_missing_files() -> None:
 
     def mock_isfile(path):
         # Return False for initial local check, True for downloaded files
-        if path in ["./local/tmp/local/a.json", "./local/tmp/local/b.json"]:
+        # Normalize paths for cross-platform comparison
+        expected_paths = [
+            os.path.join("./local/tmp/local", "a.json"),
+            os.path.join("./local/tmp/local", "b.json"),
+        ]
+        if path in expected_paths:
             return True
         return False
 
@@ -80,12 +85,24 @@ def mock_isfile(path):
         result = await download_files(json_input.path, ".json", json_input.file_names)
 
         # Each file should be attempted to be downloaded - using correct signature (with destination)
+        # Normalize paths for cross-platform compatibility
         expected_calls = [
-            call(source="local/a.json", destination="./local/tmp/local/a.json"),
-            call(source="local/b.json", destination="./local/tmp/local/b.json"),
+            call(
+                source="local/a.json",
+                destination=os.path.join("./local/tmp/local", "a.json"),
+            ),
+            call(
+                source="local/b.json",
+                destination=os.path.join("./local/tmp/local", "b.json"),
+            ),
         ]
         mock_download.assert_has_calls(expected_calls, any_order=True)
-        assert result == ["./local/tmp/local/a.json", "./local/tmp/local/b.json"]
+        # Normalize result paths for comparison
+        expected_result = [
+            os.path.join("./local/tmp/local", "a.json"),
+            os.path.join("./local/tmp/local", "b.json"),
+        ]
+        assert result == expected_result
 
 
 @pytest.mark.asyncio

tests/unit/io/test_base_io.py

@@ -231,17 +231,20 @@ async def test_download_files_download_specific_files_success(self):
         file_names = ["file1.parquet", "file2.parquet"]
         input_instance = MockReader(path, file_names)
         # Expected files will be in temporary directory after download
+        # Normalize paths for cross-platform compatibility
         expected_files = [
-            "./local/tmp/data/file1.parquet",
-            "./local/tmp/data/file2.parquet",
+            os.path.join("./local/tmp/data", "file1.parquet"),
+            os.path.join("./local/tmp/data", "file2.parquet"),
         ]
 
         def mock_isfile(path):
             # Return False for initial local check, True for downloaded files
-            if path in [
-                "./local/tmp/data/file1.parquet",
-                "./local/tmp/data/file2.parquet",
-            ]:
+            # Normalize paths for cross-platform comparison
+            expected_paths = [
+                os.path.join("./local/tmp/data", "file1.parquet"),
+                os.path.join("./local/tmp/data", "file2.parquet"),
+            ]
+            if path in expected_paths:
                 return True
             return False
 
@@ -262,14 +265,15 @@ def mock_isfile(path):
             )
 
             # Should download each specific file
+            # Normalize paths for cross-platform compatibility
             assert mock_download.call_count == 2
             mock_download.assert_any_call(
                 source="data/file1.parquet",
-                destination="./local/tmp/data/file1.parquet",
+                destination=os.path.join("./local/tmp/data", "file1.parquet"),
             )
             mock_download.assert_any_call(
                 source="data/file2.parquet",
-                destination="./local/tmp/data/file2.parquet",
+                destination=os.path.join("./local/tmp/data", "file2.parquet"),
             )
             assert result == expected_files