feat(61156998): date_trunc and T-SQL function translation for Superset (release 3.2.0)

.github/workflows/qa.yml

@@ -58,3 +58,15 @@ jobs:
 
       - name: Run unit tests
         run: make unit
+
+      # The emulator runs the real Kusto query engine locally, so the integration
+      # tests need no cluster and no credentials.
+      - name: Start the Kusto emulator
+        run: make emulator
+
+      - name: Run integration tests
+        run: make integration
+
+      - name: Stop the Kusto emulator
+        if: always()
+        run: make emulator-stop

Makefile

@@ -1,4 +1,4 @@
-.PHONY: venv install install-dev build clean check test unit integration release pypi
+.PHONY: venv install install-dev build clean check test unit integration emulator emulator-stop release pypi
 -include .env
 
 ##############################################################################
@@ -58,11 +58,28 @@ unit: # Run unit tests
 	$(PYTHON) -m pytest -v tests/unit/
 	@echo "Done.\n"
 
-integration: # Run integration tests
+integration: # Run integration tests (against the emulator unless KUSTO_URL is set)
 	@echo "Running integration tests..."
 	$(PYTHON) -m pytest -v tests/integration/
 	@echo "Done.\n"
 
+##############################################################################
+# Kusto emulator (local target for the integration tests)
+##############################################################################
+EMULATOR_IMAGE = mcr.microsoft.com/azuredataexplorer/kustainer-linux:latest
+EMULATOR_NAME = kustainer
+
+emulator: # Start the Kusto emulator and wait until it answers
+	@echo "Starting the Kusto emulator..."
+	docker run -e ACCEPT_EULA=Y -m 4G -d -p 8080:8080 --name $(EMULATOR_NAME) $(EMULATOR_IMAGE)
+	@echo "Waiting for the query engine..."
+	@timeout 300 sh -c "until curl -s -m 3 -o /dev/null -X POST -H 'Content-Type: application/json' \
+		-d '{\"csl\":\".show cluster\"}' http://localhost:8080/v1/rest/mgmt; do sleep 2; done"
+	@echo "Ready on http://localhost:8080 (database NetDefaultDB).\n"
+
+emulator-stop: # Remove the emulator container
+	-docker rm -f $(EMULATOR_NAME)
+
 ##############################################################################
 # Build and cleanup
 ##############################################################################

README.md

@@ -105,6 +105,33 @@ kustokql+https://<CLUSTER_URL>/<DATABASE>?azure_ad_client_id=<CLIENT_ID>&azure_a
 > 
 > Current `master` branch of the `apache/superset` dependent on `sqlalchemy==1.4.36`. If you want to use `sqlalchemy-kusto` with the latest unstable version of `apache/superset`, you need to install version `2.*` of the package.
 
+## Running the tests
+
+Unit tests need nothing but the package itself:
+
+```shell
+make unit
+```
+
+Integration tests run against a local [Kusto emulator](https://learn.microsoft.com/en-us/azure/data-explorer/kusto-emulator-overview),
+so they need neither a cluster nor credentials:
+
+```shell
+make emulator      # starts the kustainer-linux container on http://localhost:8080
+make integration
+make emulator-stop
+```
+
+The emulator has no TLS and no authentication, hence the `kustosql+http://` URL:
+
+```
+kustosql+http://localhost:8080/NetDefaultDB
+```
+
+To run the same tests against a real cluster instead, set `KUSTO_URL`, `DATABASE` and
+the `AZURE_AD_*` variables (see `.env.sample`). When neither the emulator nor a cluster
+is reachable, the integration tests are skipped rather than failed.
+
 ## Contributing
 
 Please see the [CONTRIBUTING.md](.github/CONTRIBUTING.md) for development setup and contributing process guidelines.

pyproject.toml

@@ -67,6 +67,7 @@ lint.ignore = [
     "ARG002",  # Unused method argument
 
     "PLR0913", # Too many arguments in function definition
+    "PLR0917", # Too many positional arguments
 ]
 
 [tool.ruff.lint.pycodestyle]

setup.py

@@ -3,7 +3,7 @@
 
 NAME = "sqlalchemy-kusto"
 DESCRIPTION = "Azure Data Explorer (Kusto) dialect for SQLAlchemy"
-VERSION = "3.1.2"
+VERSION = "3.2.0"
 
 
 REQUIREMENTS = [
@@ -37,6 +37,7 @@
     entry_points={
         "sqlalchemy.dialects": [
             "kustosql.https = sqlalchemy_kusto.dialect_sql:KustoSqlHttpsDialect",
+            "kustosql.http = sqlalchemy_kusto.dialect_sql:KustoSqlHttpDialect",
             "kustokql.https = sqlalchemy_kusto.dialect_kql:KustoKqlHttpsDialect",
         ]
     },

sqlalchemy_kusto/dbapi.py

@@ -12,6 +12,252 @@
 
 from sqlalchemy_kusto import errors
 
+# PostgreSQL date_trunc grain → (T-SQL datepart, DATEDIFF base epoch).
+# Epoch 0 is 1900-01-01, a Monday. Sub-day grains use '2000-01-01' instead because
+# DATEDIFF(second, 0, <today>) overflows a 32-bit int.
+# week_sun uses epoch -1 = 1899-12-31, a Sunday, keeping weeks Sunday-aligned.
+# Grains finer than a second are deliberately absent: DATEDIFF(millisecond, ...)
+# overflows after 24 days from any fixed epoch, so there is no usable form.
+_GRAIN_TO_TSQL: dict[str, tuple[str, str]] = {
+    "second": ("second", "'2000-01-01'"),
+    "minute": ("minute", "'2000-01-01'"),
+    "hour": ("hour", "'2000-01-01'"),
+    "day": ("day", "0"),
+    "week": ("week", "0"),
+    "week_sun": ("week", "-1"),  # extension: Sunday-based week (not in PostgreSQL)
+    "month": ("month", "0"),
+    "quarter": ("quarter", "0"),
+    "year": ("year", "0"),
+}
+
+
+def _date_trunc_expr(grain: str, expr: str) -> str | None:
+    """Build the T-SQL equivalent of date_trunc(grain, expr), or None if the grain is unknown.
+
+    Grain matching is case-insensitive and tolerates surrounding quotes.
+    """
+    key = grain.strip().strip("'\"").lower()
+    part_and_epoch = _GRAIN_TO_TSQL.get(key)
+    if part_and_epoch is None:
+        return None
+
+    part, epoch = part_and_epoch
+    if key == "week":
+        # DATEDIFF(week, ...) counts Sunday boundaries while epoch 0 is a Monday.
+        # Shifting the input back one day keeps Sunday inside the preceding Mon-Sun
+        # week, which is PostgreSQL's ISO (Monday-based) truncation.
+        expr = f"DATEADD(day, -1, {expr})"
+    return f"DATEADD({part}, DATEDIFF({part}, {epoch}, {expr}), {epoch})"
+
+
+def _skip_literal(sql: str, start: int) -> int:
+    """Return the index just past the string literal opening at `start`."""
+    quote = sql[start]
+    i = start + 1
+    while i < len(sql):
+        if sql[i] == quote:
+            if sql[i + 1 : i + 2] == quote:  # doubled quote escapes itself
+                i += 2
+                continue
+            return i + 1
+        i += 1
+    return len(sql)  # unterminated literal: treat the remainder as opaque
+
+
+def _skip_opaque(sql: str, i: int) -> int | None:
+    """Index just past the literal or comment starting at `i`, else None.
+
+    Comments must be skipped, not just left in place: an apostrophe inside one
+    ("-- don't ask") would otherwise be read as the start of a string literal and
+    swallow the rest of the query, so a real call after it never got translated.
+    """
+    char = sql[i]
+    if char in "'\"":
+        return _skip_literal(sql, i)
+    if sql.startswith("--", i):
+        end = sql.find("\n", i)
+        return len(sql) if end == -1 else end + 1
+    if sql.startswith("/*", i):
+        end = sql.find("*/", i)
+        return len(sql) if end == -1 else end + 2
+    return None
+
+
+def _split_call_args(sql: str, start: int) -> tuple[list[str], int] | None:
+    """Split the argument list of a call whose '(' ends just before `start`.
+
+    Returns (arguments, index just past the closing paren), or None when the call is
+    unbalanced, so the caller can leave such SQL for Kusto to reject. Commas inside
+    nested calls and inside string literals do not split.
+    """
+    args: list[str] = []
+    depth = 1
+    arg_start = i = start
+    while i < len(sql):
+        char = sql[i]
+        skipped = _skip_opaque(sql, i)
+        if skipped is not None:
+            i = skipped
+            continue
+        if char == "(":
+            depth += 1
+        elif char == ")":
+            depth -= 1
+            if depth == 0:
+                args.append(sql[arg_start:i])
+                return args, i + 1
+        elif char == "," and depth == 1:
+            args.append(sql[arg_start:i])
+            arg_start = i + 1
+        i += 1
+    return None
+
+
+def _fold_lower_literal(arg: str) -> str | None:
+    """lower('LITERAL') → 'literal'; anything but a single string literal stays put.
+
+    Kusto requires a LIKE pattern to be a string literal and rejects lower('%x%')
+    with "OTR0001: LIKE pattern is not a string" — yet SQLAlchemy's stock ILIKE
+    compiles to exactly lower(x) LIKE lower(y). Folding the constant keeps the
+    pattern a literal without losing the case-insensitive match; lower() on the
+    left-hand column is valid Kusto and passes through untouched.
+    """
+    if arg.startswith("'") and _skip_literal(arg, 0) == len(arg):
+        return arg.lower()
+    return None
+
+
+def _date_part_expr(unit: str, expr: str) -> str | None:
+    """date_part('unit', x) → DATEPART(unit, x); T-SQL spells the unit bare."""
+    part = unit.strip().strip("'\"").lower()
+    if not part.isalpha():
+        return None
+    return f"DATEPART({part}, {expr})"
+
+
+# Functions Kusto's T-SQL emulation does not implement, mapped to the form it does.
+# Every replacement here was executed against the Kusto emulator; anything without an
+# exact equivalent (FORMAT, DATE_FORMAT, TRY_CAST, JSON_VALUE, OPENJSON) is left alone
+# on purpose, so Kusto reports it instead of us guessing at the semantics.
+_CALL_TRANSLATIONS: dict[str, tuple[int, Any]] = {
+    # name: (argument count, builder)
+    "date_trunc": (2, _date_trunc_expr),
+    "datetrunc": (2, _date_trunc_expr),  # T-SQL 2022 spelling, same arguments
+    "date_part": (2, _date_part_expr),
+    "startofday": (1, lambda expr: _date_trunc_expr("day", expr)),
+    "startofweek": (1, lambda expr: _date_trunc_expr("week_sun", expr)),
+    "startofmonth": (1, lambda expr: _date_trunc_expr("month", expr)),
+    "startofyear": (1, lambda expr: _date_trunc_expr("year", expr)),
+    "ifnull": (2, lambda a, b: f"COALESCE({a}, {b})"),
+    # Kusto has a single datetime type, so CAST/CONVERT to DATE keeps the time of day.
+    # DATE(x) and to_date(x) are expected to drop it, which is a truncation to the day.
+    "to_date": (1, lambda expr: _date_trunc_expr("day", expr)),
+    "date": (1, lambda expr: _date_trunc_expr("day", expr)),
+    "lower": (1, _fold_lower_literal),
+}
+
+# Longest first, so "date" cannot shadow "date_trunc".
+_TRANSLATION_NAMES_LONGEST_FIRST = sorted(_CALL_TRANSLATIONS, key=len, reverse=True)
+# First letters of those names: lets the scanner skip most characters outright.
+_TRANSLATION_FIRST_CHARS = frozenset(name[0] for name in _CALL_TRANSLATIONS)
+_IDENTIFIER_CHARS = frozenset('_.[]"')
+
+
+def _translate_raw_functions(sql: str) -> str:
+    """Rewrite calls Kusto cannot parse into their T-SQL equivalents.
+
+    Applied in Cursor.execute() so it covers every execution path including Superset
+    SQLLab, which talks to the DBAPI cursor directly and bypasses the dialect hooks.
+    Names match case-insensitively; string literals, unknown grains and unknown call
+    shapes are left untouched.
+    """
+    result: list[str] = []
+    sql_lower = sql.lower()
+    i = 0
+
+    while i < len(sql):
+        end = _skip_opaque(sql, i)  # never rewrite inside a literal or a comment
+        if end is not None:
+            result.append(sql[i:end])
+            i = end
+            continue
+
+        match = _match_call_name(sql, sql_lower, i)
+        if match is None:
+            result.append(sql[i])
+            i += 1
+            continue
+
+        name, args_start = match
+        arity, builder = _CALL_TRANSLATIONS[name]
+        parsed = _split_call_args(sql, args_start)
+        if parsed is None:  # unbalanced parentheses: leave the remainder as it is
+            result.append(sql[i:])
+            break
+
+        args, end = parsed
+        if len(args) != arity or any(not arg.strip() for arg in args):
+            result.append(sql[i:end])
+            i = end
+            continue
+
+        translated = builder(*(_translate_raw_functions(arg.strip()) for arg in args))
+        result.append(translated if translated is not None else sql[i:end])
+        i = end
+
+    return "".join(result)
+
+
+def _match_call_name(sql: str, sql_lower: str, i: int) -> tuple[str, int] | None:
+    """Return (function name, index just past its '(') if a known call starts at `i`."""
+    if sql_lower[i] not in _TRANSLATION_FIRST_CHARS:
+        return None  # cheap gate: no translated name starts with this character
+    previous = sql[i - 1] if i else ""
+    if previous.isalnum() or previous in _IDENTIFIER_CHARS:
+        return None  # part of a longer identifier, e.g. my_date(
+    for name in _TRANSLATION_NAMES_LONGEST_FIRST:
+        if not sql_lower.startswith(name, i):
+            continue
+        after = i + len(name)
+        while after < len(sql) and sql[after] in " \t\r\n":
+            after += 1  # `date (x)` is the same call as `date(x)`
+        if sql[after : after + 1] == "(":
+            return name, after + 1
+    return None
+
+
+def is_tsql_query(operation: str) -> bool:
+    """Whether to send `operation` as T-SQL rather than KQL.
+
+    Kusto needs the query language up front and the only hint available is the text.
+    A leading WITH counts: Kusto executes CTEs, and treating those as KQL used to make
+    every CTE query fail before it was even sent. Leading whitespace, semicolons and
+    comments are skipped; scanning them by hand keeps this linear, which a regex with
+    a repeated alternation would not be.
+    """
+    i, length = 0, len(operation)
+    while i < length:
+        char = operation[i]
+        if char.isspace() or char == ";":
+            i += 1
+        elif operation.startswith("--", i):
+            end = operation.find("\n", i)
+            i = length if end == -1 else end + 1
+        elif operation.startswith("/*", i):
+            end = operation.find("*/", i)
+            i = length if end == -1 else end + 2
+        else:
+            break
+    for keyword in ("select", "with"):
+        end = i + len(keyword)
+        if operation[i:end].lower() != keyword:
+            continue
+        # must be a whole word: `withdrawals | count` is a KQL table, not a CTE
+        following = operation[end : end + 1]
+        if not (following.isalnum() or following == "_"):
+            return True
+    return False
+
 
 def check_closed(func):
     """Decorator that checks if connection/cursor is closed."""
@@ -82,7 +328,10 @@ def __init__(
         self.cursors: list[Cursor] = []
         kcsb = None
 
-        if azure_ad_client_id and azure_ad_client_secret and azure_ad_tenant_id:
+        if cluster.startswith("http://"):
+            # Plain HTTP means the Kusto emulator: it has no TLS and no authentication.
+            kcsb = KustoConnectionStringBuilder(cluster)
+        elif azure_ad_client_id and azure_ad_client_secret and azure_ad_tenant_id:
             # Service Principal auth
             kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication(
                 connection_string=cluster,
@@ -198,13 +447,16 @@ def close(self):
     @check_closed
     def execute(self, operation, parameters=None) -> "Cursor":
         """Executes query. Supports only SELECT statements."""
-        if operation.lower().startswith("select"):
-            self.properties.set_option("query_language", "sql")
-        else:
-            self.properties.set_option("query_language", "kql")
+        tsql = is_tsql_query(operation)
+        self.properties.set_option("query_language", "sql" if tsql else "kql")
 
         query = Cursor._apply_parameters(operation, parameters)
         query = query.rstrip()
+        if tsql:
+            # After parameter substitution, so a bound ILIKE pattern inside
+            # lower(%(p)s) can fold into the string literal Kusto requires.
+            # T-SQL only: a KQL query must never be rewritten.
+            query = _translate_raw_functions(query)
         try:
             server_response = self.kusto_client.execute(
                 self.database, query, self.properties