diff --git a/ingestion/src/metadata/profiler/orm/functions/datetime.py b/ingestion/src/metadata/profiler/orm/functions/datetime.py index 2a1dfe7d5f28..33f2cb0e264f 100644 --- a/ingestion/src/metadata/profiler/orm/functions/datetime.py +++ b/ingestion/src/metadata/profiler/orm/functions/datetime.py @@ -106,6 +106,12 @@ def _(elements, compiler, **kwargs): return generic_function(elements, compiler, **kwargs) +@compiles(DatetimeAddFn, Dialects.Doris) +def _(elements, compiler, **kwargs): + """Doris date and datetime function""" + return doris_function(elements, compiler, **kwargs) + + @compiles(DatetimeAddFn, Dialects.MySQL) def _(elements, compiler, **kwargs): """MySQL date and datetime function""" @@ -185,6 +191,12 @@ def _(elements, compiler, **kwargs): # pylint: disable=unused-argument return f"DATETIME_SUB({func.current_timestamp()}, INTERVAL {interval} {interval_unit})" +@compiles(TimestampAddFn, Dialects.Doris) +def _(elements, compiler, **kwargs): + """Doris timestamp function""" + return doris_function(elements, compiler, **kwargs) + + @compiles(TimestampAddFn, Dialects.MySQL) def _(elements, compiler, **kwargs): """MySQL timestamp function""" @@ -234,6 +246,13 @@ def generic_function(elements, compiler, **kwargs): return f"CAST(CURRENT_TIMESTAMP - interval '{interval}' {interval_unit} AS TIMESTAMP)" +def doris_function(elements, compiler, **kwargs): + """Doris timestamp and datetime function""" + interval = elements.clauses.clauses[0].value + interval_unit = compiler.process(elements.clauses.clauses[1], **kwargs) + return f"CAST(CURRENT_TIMESTAMP - interval '{interval}' {interval_unit} AS DATETIME)" + + def mysql_function(elements, compiler, **kwargs): """MySQL timestamp and datetime function""" interval = elements.clauses.clauses[0].value diff --git a/ingestion/tests/unit/source/database/doris/test_connection.py b/ingestion/tests/unit/source/database/doris/test_connection.py index 3975ff3e31d9..bb5d5c36e755 100644 --- a/ingestion/tests/unit/source/database/doris/test_connection.py +++ b/ingestion/tests/unit/source/database/doris/test_connection.py @@ -11,7 +11,8 @@ """Unit tests for Doris connection handling.""" import pytest -from sqlalchemy import Column, Integer, MetaData, Table, select +from pydoris.sqlalchemy.dialect import DorisDialect +from sqlalchemy import DATETIME, TIMESTAMP, Column, Integer, MetaData, Table, select, text from metadata.generated.schema.entity.services.connections.database.dorisConnection import ( DorisConnection as DorisConnectionConfig, @@ -21,6 +22,7 @@ ) from metadata.ingestion.connections.connection import BaseConnection from metadata.ingestion.source.database.doris.connection import DorisConnection +from metadata.utils.sqa_utils import dispatch_to_date_or_datetime @pytest.fixture @@ -58,6 +60,19 @@ def test_doris_identifiers_are_always_quoted(column_name: str, doris_connection_ assert "FROM `events`" in query +@pytest.mark.parametrize( + "column_type", + [DATETIME(), TIMESTAMP()], + ids=["datetime", "timestamp"], +) +def test_doris_time_partition_filter_uses_datetime_cast(column_type: DATETIME | TIMESTAMP): + partition_boundary = dispatch_to_date_or_datetime(1, text("DAY"), column_type) + + query = str(partition_boundary.compile(dialect=DorisDialect())) + + assert query == "CAST(CURRENT_TIMESTAMP - interval '1' DAY AS DATETIME)" + + def test_basic_auth_builds_doris_engine( doris_connection_config: DorisConnectionConfig, ):