From 6f444ab7ca60f66f0e4402720cd5615e391d5d85 Mon Sep 17 00:00:00 2001 From: Anmol Jaiswal Date: Tue, 10 Feb 2026 20:34:36 +0530 Subject: [PATCH] feat: add Databricks dialect awareness to DatabaseSessionService Adds dialect-specific handling for Databricks SQL in the session service schema types (DynamicJSON and PreciseTimestamp), and introduces a _DATABRICKS_DIALECT constant for future use. This ensures the TypeDecorator classes explicitly handle the Databricks dialect rather than falling through to default TEXT behavior. While the databricks-sqlalchemy driver currently lacks async support, this change prepares the codebase for when async compatibility is available. --- src/google/adk/sessions/database_session_service.py | 1 + src/google/adk/sessions/schemas/shared.py | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index c7c86e6e5f..6665669ed9 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -64,6 +64,7 @@ _MARIADB_DIALECT = "mariadb" _MYSQL_DIALECT = "mysql" _POSTGRESQL_DIALECT = "postgresql" +_DATABRICKS_DIALECT = "databricks" # Tuple key order for in-process per-session lock maps: # (app_name, user_id, session_id). _SessionLockKey: TypeAlias = tuple[str, str, str] diff --git a/src/google/adk/sessions/schemas/shared.py b/src/google/adk/sessions/schemas/shared.py index 25d4ea9e95..311c616ec6 100644 --- a/src/google/adk/sessions/schemas/shared.py +++ b/src/google/adk/sessions/schemas/shared.py @@ -37,6 +37,9 @@ def load_dialect_impl(self, dialect: Dialect): if dialect.name == "mysql": # Use LONGTEXT for MySQL to address the data too long issue return dialect.type_descriptor(mysql.LONGTEXT) + if dialect.name == "databricks": + # Databricks SQL stores JSON as STRING; use Text (the default) + return dialect.type_descriptor(Text) return dialect.type_descriptor(Text) # Default to Text for other dialects def process_bind_param(self, value, dialect: Dialect): @@ -64,4 +67,7 @@ class PreciseTimestamp(TypeDecorator): def load_dialect_impl(self, dialect): if dialect.name == "mysql": return dialect.type_descriptor(mysql.DATETIME(fsp=6)) + if dialect.name == "databricks": + # Databricks TIMESTAMP type natively supports microsecond precision + return dialect.type_descriptor(DateTime) return self.impl