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