Fix the type code issue

Author: Peng Ren

pymongosql/result_set.py

@@ -430,22 +430,142 @@ class PreProcessedResultSet(ResultSet):
     """Result set for pre-formatted rows that don't need projection processing.
 
     Used when rows are already projected and formatted (e.g., from SQLite intermediate storage).
-    Skips the document projection step and goes directly to formatting as tuples.
+    Skips the document projection step but applies type conversions for projection function results.
     """
 
     def _process_and_cache_batch(self, batch: List[Dict[str, Any]]) -> None:
         """Process and cache a batch of pre-processed documents.
 
-        Unlike the base ResultSet, this skips projection processing and directly formats results.
+        Unlike the base ResultSet, this skips projection processing but still applies type conversions
+        for values that came from projection functions (e.g., converting strings back to datetime objects).
         """
         if not batch:
             return
-        # Skip projection processing - rows are already in final form
-        # Just format to output format (tuple)
-        formatted_batch = [self._format_result(doc) for doc in batch]
+        # Skip full projection processing - rows are already in final form
+        # But apply type conversions for projection function results
+        converted_batch = [self._convert_projection_types(doc) for doc in batch]
+        # Format to output format (tuple)
+        formatted_batch = [self._format_result(doc) for doc in converted_batch]
         self._cached_results.extend(formatted_batch)
         self._total_fetched += len(batch)
 
+    def _convert_projection_types(self, doc: Dict[str, Any]) -> Dict[str, Any]:
+        """Convert string values from SQLite back to proper types based on projection functions.
+
+        When projection functions produce datetime/date objects, SQLite stores them as strings.
+        This method converts them back to proper Python types.
+        """
+        if not self._execution_plan:
+            return doc
+
+        from datetime import datetime, timezone
+
+        from .sql.projection_functions import ProjectionFunctionRegistry
+
+        converted = dict(doc)
+
+        # Get projection information from execution plan
+        projection_functions = getattr(self._execution_plan, "projection_functions", {})
+        column_aliases = getattr(self._execution_plan, "column_aliases", {})
+
+        if not projection_functions:
+            return converted
+
+        _ = ProjectionFunctionRegistry()
+
+        # Iterate through projection functions to convert values
+        for field_name, func_info in projection_functions.items():
+            # The column name might be aliased, so check both the original name and the alias
+            col_names_to_check = [field_name]
+            if field_name in column_aliases:
+                col_names_to_check.append(column_aliases[field_name])
+            # Also check for the mongo_to_bracket_key format
+            col_names_to_check.append(self._mongo_to_bracket_key(field_name))
+
+            col_name = None
+            for check_name in col_names_to_check:
+                if check_name in converted:
+                    col_name = check_name
+                    break
+
+            if not col_name:
+                continue
+
+            value = converted[col_name]
+            if value is None:
+                continue
+
+            # Extract function name and format parameter
+            func_name = None
+            format_param = None
+
+            if isinstance(func_info, dict):
+                func_name = func_info.get("name")
+                format_param = func_info.get("format_param")
+            elif isinstance(func_info, (list, tuple)):
+                if len(func_info) >= 1:
+                    func_name = func_info[0]
+                if len(func_info) >= 2:
+                    format_param = func_info[1]
+
+            if not func_name:
+                continue
+
+            # Handle special cases for SQLite-stored projection function results
+            if func_name.upper() in ("DATE", "DATETIME", "TIMESTAMP"):
+                # These functions should produce datetime objects
+                try:
+                    if isinstance(value, str):
+                        # Handle BSON Timestamp string representation: "Timestamp(timestamp_int, increment)"
+                        if value.startswith("Timestamp("):
+                            match = re.match(r"Timestamp\((\d+),\s*\d+\)", value)
+                            if match:
+                                timestamp_int = int(match.group(1))
+                                # Convert UNIX timestamp to datetime
+                                converted[col_name] = datetime.fromtimestamp(timestamp_int, tz=timezone.utc).replace(
+                                    tzinfo=None
+                                )
+                            continue
+
+                        # Try to parse as ISO format date/datetime string
+                        try:
+                            # Try full datetime first
+                            converted[col_name] = datetime.fromisoformat(value)
+                        except (ValueError, TypeError):
+                            try:
+                                # Try date only (YYYY-MM-DD)
+                                converted[col_name] = datetime.strptime(value, "%Y-%m-%d")
+                            except (ValueError, TypeError):
+                                try:
+                                    # Try with custom format if provided
+                                    if format_param:
+                                        converted[col_name] = datetime.strptime(value, format_param)
+                                except (ValueError, TypeError):
+                                    # Keep original value if all conversions fail
+                                    pass
+                except Exception as e:
+                    _logger.debug(f"Error converting {col_name} using {func_name}: {e}")
+                    # Keep original value if conversion fails
+            elif func_name.upper() in ("NUMBER", "INT"):
+                # These functions should produce numeric values
+                try:
+                    if isinstance(value, str):
+                        # Try to convert to float
+                        converted[col_name] = float(value)
+                except (ValueError, TypeError):
+                    # Keep original value if conversion fails
+                    pass
+            elif func_name.upper() == "BOOL":
+                # BOOL function should produce boolean values
+                try:
+                    if isinstance(value, str):
+                        converted[col_name] = value.lower() in ("true", "1", "yes")
+                except Exception:
+                    # Keep original value if conversion fails
+                    pass
+
+        return converted
+
 
 # For backward compatibility
 MongoResultSet = ResultSet

pymongosql/superset_mongodb/executor.py

@@ -161,8 +161,16 @@ def execute(
             # This tells the cursor to use PreProcessedResultSet instead of regular ResultSet
             final_execution_plan.from_intermediate_storage = True
 
-            # Extract projection_output from MongoDB result set description
-            # This provides the correct mapping from outer query column names to type codes
+            # Extract projection_output and projection_functions from MongoDB execution plan
+            # This provides the correct mapping from outer query column names to type codes and conversion functions
+            if hasattr(mongo_execution_plan, "projection_output") and mongo_execution_plan.projection_output:
+                final_execution_plan.projection_output = mongo_execution_plan.projection_output
+
+            # Also copy projection_functions from the inner MongoDB query so PreProcessedResultSet
+            # can convert string values from SQLite back to proper types (datetime, etc.)
+            if hasattr(mongo_execution_plan, "projection_functions") and mongo_execution_plan.projection_functions:
+                final_execution_plan.projection_functions = mongo_execution_plan.projection_functions
+
             if mongo_result_set.description:
                 projection_output = []
                 for col_name, type_code, *_ in mongo_result_set.description:
@@ -177,7 +185,9 @@ def execute(
 
                     projection_output.append({"output_name": col_name, "function": func_info})
 
-                final_execution_plan.projection_output = projection_output
+                # Only update if projection_output wasn't already set from mongo_execution_plan
+                if not hasattr(final_execution_plan, "projection_output") or not final_execution_plan.projection_output:
+                    final_execution_plan.projection_output = projection_output
 
             self._execution_plan = final_execution_plan
 

tests/test_projection_functions.py

@@ -87,7 +87,7 @@ def test_convert_none(self):
     def test_get_type_code(self):
         """Test type code for date function"""
         func = DateFunction()
-        assert func.get_type_code() == "datetime"
+        assert func.get_type_code() == datetime
 
 
 class TestDatetimeFunction:
@@ -167,7 +167,7 @@ def test_convert_none(self):
     def test_get_type_code(self):
         """Test type code for datetime function"""
         func = DatetimeFunction()
-        assert func.get_type_code() == "datetime"
+        assert func.get_type_code() == datetime
 
 
 class TestTimestampFunction:
@@ -237,7 +237,7 @@ def test_convert_none(self):
     def test_get_type_code(self):
         """Test type code for timestamp function"""
         func = TimestampFunction()
-        assert func.get_type_code() == "datetime"
+        assert func.get_type_code() == datetime
 
 
 class TestNumberFunction:
@@ -281,7 +281,7 @@ def test_convert_none(self):
     def test_get_type_code(self):
         """Test type code for number function"""
         func = NumberFunction()
-        assert func.get_type_code() == "float"
+        assert func.get_type_code() == float
 
 
 class TestBoolFunction:
@@ -332,7 +332,7 @@ def test_convert_none(self):
     def test_get_type_code(self):
         """Test type code for bool function"""
         func = BoolFunction()
-        assert func.get_type_code() == "bool"
+        assert func.get_type_code() == bool
 
 
 class TestSubstrFunction:
@@ -394,7 +394,7 @@ def test_convert_substr_non_string(self):
     def test_get_type_code(self):
         """Test type code for substr function"""
         func = SubstrFunction()
-        assert func.get_type_code() == "str"
+        assert func.get_type_code() == str
 
 
 class TestReplaceFunction:
@@ -449,7 +449,7 @@ def test_convert_replace_non_string(self):
     def test_get_type_code(self):
         """Test type code for replace function"""
         func = ReplaceFunction()
-        assert func.get_type_code() == "str"
+        assert func.get_type_code() == str
 
 
 class TestTrimFunction:
@@ -503,7 +503,7 @@ def test_convert_trim_non_string(self):
     def test_get_type_code(self):
         """Test type code for trim function"""
         func = TrimFunction()
-        assert func.get_type_code() == "str"
+        assert func.get_type_code() == str
 
 
 class TestUpperFunction:
@@ -553,7 +553,7 @@ def test_convert_non_string(self):
     def test_get_type_code(self):
         """Test type code for upper function"""
         func = UpperFunction()
-        assert func.get_type_code() == "str"
+        assert func.get_type_code() == str
 
 
 class TestLowerFunction:
@@ -603,7 +603,7 @@ def test_convert_non_string(self):
     def test_get_type_code(self):
         """Test type code for lower function"""
         func = LowerFunction()
-        assert func.get_type_code() == "str"
+        assert func.get_type_code() == str
 
 
 class TestProjectionFunctionRegistry:

tests/test_superset_connection.py

@@ -425,14 +425,18 @@ def test_projection_functions_with_superset_execution(self, superset_conn):
         # All projection functions return string type codes like 'float', 'datetime', 'date', etc.
         import datetime
 
-        # NUMBER() should produce 'float' type code
-        assert type_codes["numeric_age"] == "float", f"Expected 'float' but got {type_codes['numeric_age']!r}"
+        # NUMBER() should produce float type code
+        assert type_codes["numeric_age"] == float, f"Expected float but got {type_codes['numeric_age']!r}"
 
-        # DATE() should produce 'datetime' type code (note: projection functions return 'datetime' for both DATE and DATETIME)
-        assert type_codes["creation_date"] == "datetime", f"Expected 'datetime' but got {type_codes['creation_date']!r}"
+        # DATE() should produce datetime type code (note: projection functions return datetime for both DATE and DATETIME)
+        assert (
+            type_codes["creation_date"] == datetime.datetime
+        ), f"Expected datetime but got {type_codes['creation_date']!r}"
 
-        # DATETIME() should produce 'datetime' type code
-        assert type_codes["last_updated"] == "datetime", f"Expected 'datetime' but got {type_codes['last_updated']!r}"
+        # DATETIME() should produce datetime type code
+        assert (
+            type_codes["last_updated"] == datetime.datetime
+        ), f"Expected datetime but got {type_codes['last_updated']!r}"
 
         # Verify data values are correctly converted
         for row in rows:
@@ -445,11 +449,11 @@ def test_projection_functions_with_superset_execution(self, superset_conn):
                 row[numeric_age_idx], (int, float)
             ), f"numeric_age should be numeric, got {type(row[numeric_age_idx])}"
 
-            # DATE() should convert to date object
+            # DATE() should convert to datetime object
             creation_date_val = row[creation_date_idx]
             assert creation_date_val is None or isinstance(
                 creation_date_val, datetime.date
-            ), f"creation_date should be date or None, got {type(creation_date_val)}"
+            ), f"creation_date should be datetime or None, got {type(creation_date_val)}"
 
             # DATETIME() should convert to datetime object
             last_updated_val = row[last_updated_idx]
@@ -494,24 +498,22 @@ def test_projection_functions_with_custom_format(self, superset_conn):
 
         import datetime
 
-        from bson import Timestamp
-
         # Verify expected type codes
-        # All projection functions return string type codes
-        # DATE() should produce 'datetime' type code
+        # All projection functions return type objects, not strings
+        # DATE() should produce datetime type code
         assert (
-            type_codes["formatted_date"] == "datetime"
-        ), f"Expected 'datetime' but got {type_codes['formatted_date']!r}"
+            type_codes["formatted_date"] == datetime.datetime
+        ), f"Expected datetime but got {type_codes['formatted_date']!r}"
 
-        # DATETIME() should produce 'datetime' type code
+        # DATETIME() should produce datetime type code
         assert (
-            type_codes["formatted_datetime"] == "datetime"
-        ), f"Expected 'datetime' but got {type_codes['formatted_datetime']!r}"
+            type_codes["formatted_datetime"] == datetime.datetime
+        ), f"Expected datetime but got {type_codes['formatted_datetime']!r}"
 
-        # TIMESTAMP() should produce 'datetime' type code
+        # TIMESTAMP() should produce datetime type code
         assert (
-            type_codes["timestamp_value"] == "datetime"
-        ), f"Expected 'datetime' but got {type_codes['timestamp_value']!r}"
+            type_codes["timestamp_value"] == datetime.datetime
+        ), f"Expected datetime but got {type_codes['timestamp_value']!r}"
 
         # Verify data values are correctly converted with custom formats
         for row in rows:
@@ -528,14 +530,14 @@ def test_projection_functions_with_custom_format(self, superset_conn):
             # DATETIME() with format should convert to datetime object
             datetime_value = row[formatted_datetime_idx]
             assert datetime_value is None or isinstance(
-                datetime_value, datetime.datetime
+                datetime_value, datetime.date
             ), f"formatted_datetime should be datetime or None, got {type(datetime_value)}"
 
-            # TIMESTAMP() with format should convert to Timestamp object
+            # TIMESTAMP() with format should convert to datetime object
             timestamp_value = row[timestamp_value_idx]
             assert timestamp_value is None or isinstance(
-                timestamp_value, Timestamp
-            ), f"timestamp_value should be Timestamp or None, got {type(timestamp_value)}"
+                timestamp_value, datetime.datetime
+            ), f"timestamp_value should be datetime or None, got {type(timestamp_value)}"
 
     def test_empty_result_with_valid_description(self, superset_conn):
         """Test that description is available for result sets, even if empty after filtering"""