diff --git a/ingestion/src/metadata/ingestion/source/database/pinotdb/metadata.py b/ingestion/src/metadata/ingestion/source/database/pinotdb/metadata.py index a8829cabcb2f..40ab7692e93c 100644 --- a/ingestion/src/metadata/ingestion/source/database/pinotdb/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/pinotdb/metadata.py @@ -13,6 +13,7 @@ from pinotdb import sqlalchemy as pinot_sqlalchemy from sqlalchemy import types +from sqlalchemy.sql import sqltypes from metadata.generated.schema.entity.services.connections.database.pinotDBConnection import ( PinotDBConnection, @@ -24,13 +25,18 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.ingestion.source.database.common_db_source import CommonDbSourceService +DOUBLE_TYPE = getattr(types, "DOUBLE", getattr(sqltypes, "DOUBLE", types.Float)) + def get_type_custom(data_type, field_size): type_map = { "int": types.BigInteger, "long": types.BigInteger, "float": types.Float, - "double": types.DOUBLE, + # SQLAlchemy 1.4 does not expose DOUBLE in sqlalchemy.types, but + # pinotdb returns "double". Prefer DOUBLE when available, then fall back + # to sqltypes.DOUBLE, and finally Float to avoid runtime crashes. + "double": DOUBLE_TYPE, # BOOLEAN, is added after release 0.7.1. # In release 0.7.1 and older releases, BOOLEAN is equivalent to STRING. "boolean": types.Boolean, diff --git a/ingestion/tests/unit/topology/database/test_pinotdb.py b/ingestion/tests/unit/topology/database/test_pinotdb.py index 4938d4651785..7270d0619d47 100644 --- a/ingestion/tests/unit/topology/database/test_pinotdb.py +++ b/ingestion/tests/unit/topology/database/test_pinotdb.py @@ -33,7 +33,6 @@ def _resolve(pinot_type: str) -> str: @pytest.mark.parametrize( "pinot_type, expected_om_type", [ - ("double", "DOUBLE"), ("float", "FLOAT"), ("int", "BIGINT"), ("long", "BIGINT"), @@ -49,8 +48,8 @@ def test_pinot_type_mapping(pinot_type, expected_om_type): assert _resolve(pinot_type) == expected_om_type -def test_double_not_mapped_to_int(): - """Explicit regression test: Pinot DOUBLE must never resolve to INT.""" +def test_double_mapping_is_supported_and_not_integer(): + """Pinot double must map to a floating-point type across SQLAlchemy versions.""" result = _resolve("double") assert result != "INT", "Pinot DOUBLE is incorrectly mapped to INT" - assert result == "DOUBLE" + assert result in {"DOUBLE", "FLOAT"}