feat: Add DatabaseSessionService with Hibernate and Flyway support #654

prasadskarmarkar · 2025-12-08T08:50:13Z

Summary

Implements a production-ready database-backed session service that provides
persistent storage for ADK sessions, events, and state using JPA/Hibernate.

Key Features

Hibernate 6.6 + HikariCP connection pooling for optimal performance
Flyway migrations for schema versioning and zero-downtime deployments
Multi-database support: PostgreSQL, MySQL, H2, and other RDBMS
Thread-safe operations with pessimistic locking for concurrent updates
Comprehensive test coverage with H2 in-memory database
Dialect-aware JSON storage (JSONB for PostgreSQL, CLOB for others)
Event filtering and pagination for efficient data retrieval

Architecture

Database dependencies are non-optional in the core module for ease of use,
consistent with Python ADK's approach where `DatabaseSessionService` is in the
main package.

Note: Database drivers (PostgreSQL, MySQL) are marked as
`true`, so users must add their specific driver
dependency to their project.

Files Changed

28 files, 3,874 insertions
New `DatabaseSessionService` implementation
JPA entities and converters for session/event storage
Flyway migration scripts for PostgreSQL, MySQL, and H2
Comprehensive unit tests

Test Plan

Run `DatabaseSessionServiceTest` - all 10 tests pass
Verify Flyway migrations apply correctly
Test multi-pod deployment scenarios with migration locks
Manual testing with PostgreSQL
Manual testing with MySQL"

google-cla · 2025-12-08T08:50:18Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

gemini-code-assist · 2025-12-08T08:51:04Z

Summary of Changes

Hello @prasadskarmarkar, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new, production-ready DatabaseSessionService designed to provide persistent storage for ADK sessions, events, and application/user state. It leverages JPA/Hibernate for ORM capabilities and integrates Flyway for robust schema management, supporting multiple database types. The service is built with performance and data integrity in mind, incorporating connection pooling, pessimistic locking, and dialect-aware JSON storage, making it a foundational improvement for stateful ADK applications.

Highlights

New Database Session Service: Introduces a DatabaseSessionService for persistent storage of ADK sessions, events, and state using JPA/Hibernate, replacing in-memory or other non-persistent solutions.
Database Persistence Layer: Implements a robust persistence layer utilizing Hibernate 6.6 with HikariCP for efficient connection pooling, ensuring optimal performance and scalability for session management.
Automated Schema Management: Integrates Flyway for database schema versioning and migrations, enabling zero-downtime deployments and consistent schema evolution across different environments.
Multi-Database Support: Provides out-of-the-box support for multiple relational databases including PostgreSQL, MySQL, and H2, with a flexible architecture to extend to other RDBMS.
Advanced Data Handling: Features dialect-aware JSON storage (JSONB for PostgreSQL, CLOB for others) and thread-safe operations with pessimistic locking for concurrent updates, ensuring data integrity.
Session and Event Management Features: Includes comprehensive functionalities for creating, retrieving, listing, appending events to, and deleting sessions, along with event filtering and pagination for efficient data retrieval.
Comprehensive Testing: Adds extensive unit tests for the DatabaseSessionService, Flyway migrations, and state management, primarily using an H2 in-memory database for reliability.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a significant and well-structured feature: a database-backed session service using Hibernate and Flyway. The implementation demonstrates strong practices, including multi-database support, Flyway for schema migrations, and pessimistic locking for concurrency control. The code is generally of high quality. However, I've identified a few critical and high-severity issues that should be addressed. The most critical concern is the use of Java serialization for persistence, which is fragile and insecure. There are also bugs causing data loss during the serialization/deserialization of certain event fields. Additionally, I've noted some misleading tests and opportunities for code cleanup and dependency improvements. My detailed comments are below.

core/src/main/java/com/google/adk/sessions/db/entity/StorageEvent.java

core/src/main/java/com/google/adk/sessions/DatabaseSessionService.java

core/src/main/java/com/google/adk/sessions/db/entity/StorageEvent.java

core/pom.xml

core/src/main/java/com/google/adk/sessions/DatabaseSessionService.java

core/src/main/java/com/google/adk/sessions/db/converter/DialectAwareJsonConverter.java

core/src/test/java/com/google/adk/sessions/StateManagementTest.java

prasadskarmarkar · 2025-12-10T06:50:19Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces a robust, production-ready DatabaseSessionService using Hibernate and Flyway, which is a significant and well-executed feature. The implementation correctly uses pessimistic locking for concurrency control and Flyway for schema migrations, supporting multiple database backends. The test coverage is comprehensive and impressive.

My review focuses on improving maintainability, robustness, and consistency. Key suggestions include:

Refactoring duplicated transaction management logic in DatabaseSessionService.
Removing potential write side-effects from the getSession method to adhere to command-query separation.
Enhancing error logging by replacing System.err with a proper logger and logging exceptions that are currently swallowed.
Improving schema consistency in the MySQL migration script.
Cleaning up potentially unused code and over-broad .gitignore rules.

core/src/main/java/com/google/adk/sessions/DatabaseSessionService.java

.gitignore

core/src/main/java/com/google/adk/runner/Runner.java

core/src/main/java/com/google/adk/sessions/DatabaseSessionService.java

core/src/main/java/com/google/adk/sessions/db/converter/DialectAwareJsonConverter.java

core/src/main/java/com/google/adk/sessions/db/entity/StorageEvent.java

core/src/main/resources/db/migration/mysql/V1__Initial_schema.sql

prasadskarmarkar · 2025-12-10T17:16:01Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces a robust, production-ready DatabaseSessionService using Hibernate for persistence and Flyway for database migrations. The implementation is impressive, with strong support for multiple database dialects, thread-safe operations using pessimistic locking, and comprehensive test coverage. My review focuses on a few areas for improvement, including strengthening the database schema for timestamp handling, simplifying state management logic, and removing potentially obsolete code. Overall, this is a high-quality contribution that significantly enhances the ADK's session management capabilities.

core/src/main/resources/db/migration/postgresql/V1__Initial_schema.sql

core/src/main/java/com/google/adk/sessions/DatabaseSessionService.java

core/src/main/java/com/google/adk/sessions/db/util/SchemaInitializer.java

prasadskarmarkar · 2025-12-10T22:25:49Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces a robust, production-ready DatabaseSessionService using Hibernate for persistence and Flyway for database migrations. The implementation is comprehensive, covering multi-database support, thread-safe operations with pessimistic locking, and extensive test coverage. The architecture is well-thought-out, particularly the handling of migrations in multi-pod environments.

My review focuses on improving concurrency handling and ensuring data type consistency. I've identified a race condition during the creation of new application/user states and an issue with the deep copying of state maps, both of which could affect correctness under load. I've also suggested an improvement to the MySQL schema for better performance and consistency. Overall, this is an excellent and thorough implementation.

gemini-code-assist · 2025-12-10T22:28:07Z

core/src/main/java/com/google/adk/sessions/DatabaseSessionService.java

+  private StorageAppState getOrCreateAppState(EntityManager em, String appName) {
+    StorageAppState appState = em.find(StorageAppState.class, appName);
+    if (appState == null) {
+      appState = new StorageAppState();
+      appState.setAppName(appName);
+      appState.setState(new HashMap<>());
+      appState.setUpdateTime(Instant.now());
+      em.persist(appState);
+    }
+    return appState;
+  }


There is a race condition in the getOrCreateAppState method (and similarly in getOrCreateUserState). The current "check-then-act" logic (em.find() followed by em.persist()) is not atomic across different transactions. If two threads concurrently call createSession for the same new application, both could find that the StorageAppState does not exist and then both attempt to persist a new one. This will cause the second transaction to fail with a PersistenceException (due to a primary key constraint violation) when it tries to commit. This can lead to failed session creations under concurrent load.

To fix this, you should handle the potential concurrency more robustly. One common pattern is to attempt the persist operation and catch the PersistenceException. If the exception is a constraint violation, it means another transaction created the entity in the meantime, so you can then safely find the existing entity.

A simplified version of this pattern would look like this:

// Inside a try-catch block for em.persist() and em.flush() } catch (PersistenceException e) { // A more robust check would inspect the exception cause for a constraint violation. em.clear(); // Detach the failed entity from the persistence context. StorageAppState existingState = em.find(StorageAppState.class, appName); if (existingState != null) { return existingState; } // If it's still null, something else went wrong. throw e; }

This change will make the session creation process more resilient in a multi-threaded or multi-pod environment.

core/src/main/java/com/google/adk/sessions/db/converter/JsonUserType.java

core/src/main/resources/db/migration/mysql/V1__Initial_schema.sql

prasadskarmarkar · 2025-12-11T00:43:36Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces a comprehensive and well-designed DatabaseSessionService, adding persistent storage for sessions using JPA/Hibernate and Flyway for migrations. The implementation is robust, with good handling of concurrency through pessimistic locking, transaction management, and support for multiple database dialects. The test coverage is excellent, covering unit, integration, and concurrency scenarios.

My review focuses on a few key areas to further improve correctness and maintainability:

Data Serialization: There are opportunities to make the serialization of Content objects more robust by fully leveraging Jackson's capabilities, preventing potential data loss.
Schema Definition: A minor adjustment in the H2 database schema can prevent potential data truncation issues.
Code Consistency: A small refactoring in one of the methods can improve consistency with the rest of the class.

Overall, this is a high-quality contribution that significantly enhances the ADK's capabilities.

core/src/main/java/com/google/adk/sessions/db/entity/StorageEvent.java

core/src/main/resources/db/migration/h2/V1__Initial_schema.sql

core/src/main/java/com/google/adk/sessions/DatabaseSessionService.java

prasadskarmarkar · 2025-12-13T01:00:32Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces a comprehensive and production-ready DatabaseSessionService using Hibernate for persistence and Flyway for database migrations. The implementation is robust, featuring support for multiple database backends (PostgreSQL, MySQL, H2), pessimistic locking for concurrent operations, and a well-structured approach to schema management. The addition of extensive unit and integration tests covering concurrency, state management, and migrations is commendable. My review focuses on a few minor areas for improvement in error handling and code clarity within the new service.

core/src/main/java/com/google/adk/sessions/DatabaseSessionService.java

prasadskarmarkar · 2025-12-14T07:57:45Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces a comprehensive and production-ready DatabaseSessionService. The implementation is robust, featuring support for multiple databases (PostgreSQL, MySQL, H2) via Hibernate, schema management with Flyway, and excellent concurrency control with pessimistic locking. The separation of state into app, user, and session scopes is well-designed. The test coverage is extensive and covers critical aspects like concurrency, migrations, and state management. My review includes a few suggestions for improving dependency management for MySQL, minor code style enhancements, and ensuring consistency in the database schema scripts. Overall, this is an excellent contribution.

gemini-code-assist · 2025-12-14T08:00:27Z

core/src/main/resources/db/migration/mysql/V1__Initial_schema.sql

+    session_id VARCHAR(128) NOT NULL,
+    invocation_id VARCHAR(256),
+    author VARCHAR(256),
+    actions JSON,


There's an inconsistency in the data types used for JSON storage in the MySQL schema. The events.actions column is defined as JSON, while other columns intended for JSON data (like sessions.state, events.content, etc.) are defined as LONGTEXT. The JsonUserType and EventActionsUserType implementations write a string to the database for non-PostgreSQL dialects. To maintain consistency with the UserType implementations and other columns, events.actions should also be LONGTEXT.

Suggested change

actions JSON,

actions LONGTEXT,

gemini-code-assist · 2025-12-14T08:00:28Z

core/pom.xml

+    <!-- Database Migration -->
+
+    <!-- Flyway Core - Database schema migration and versioning tool -->
+    <dependency>
+      <groupId>org.flywaydb</groupId>
+      <artifactId>flyway-core</artifactId>
+    </dependency>
+
+    <!-- Flyway PostgreSQL Support - PostgreSQL-specific migration support -->
+    <dependency>
+      <groupId>org.flywaydb</groupId>
+      <artifactId>flyway-database-postgresql</artifactId>
+      <scope>runtime</scope>
+    </dependency>


For full support of MySQL, especially in multi-pod environments where migration locking is crucial, it's recommended to include the flyway-mysql dependency. This is analogous to how flyway-database-postgresql is included for PostgreSQL-specific features. Please consider adding it here and to the parent pom.xml's dependencyManagement section.

gemini-code-assist · 2025-12-14T08:00:28Z

core/src/main/java/com/google/adk/sessions/DatabaseSessionService.java

+    return Completable.fromCallable(
+            () -> {
+              executeInTransaction(
+                  em -> {
+                    // Find session
+                    SessionId id = new SessionId(appName, userId, sessionId);
+                    StorageSession session = em.find(StorageSession.class, id);
+                    if (session == null) {
+                      throw new SessionNotFoundException(
+                          String.format(
+                              "Session not found: appName=%s, userId=%s, sessionId=%s",
+                              appName, userId, sessionId));
+                    }
+
+                    // Delete session (cascade will delete events)
+                    em.remove(session);
+
+                    return null;
+                  },
+                  "Error deleting session");
+              return null;
+            })


The use of Completable.fromCallable with a Callable that just returns null is a bit indirect for an operation that has no return value. The Completable.fromAction operator is a better semantic fit here.

Consider refactoring to use fromAction to make the code slightly cleaner and more idiomatic, like so:

return Completable.fromAction(() -> executeInTransaction(em -> { ... }, "Error deleting session")) .subscribeOn(Schedulers.io());

gemini-code-assist · 2025-12-14T08:00:28Z

core/src/main/java/com/google/adk/sessions/db/converter/ByteArrayConverter.java

+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.adk.sessions.db.converter;
+
+import jakarta.persistence.AttributeConverter;
+import jakarta.persistence.Converter;
+
+/** JPA attribute converter that passes through byte[] for databases like PostgreSQL using bytea. */
+@Converter(autoApply = false)
+public class ByteArrayConverter implements AttributeConverter<byte[], byte[]> {
+
+  @Override
+  public byte[] convertToDatabaseColumn(byte[] attribute) {
+    return attribute; // store as-is for bytea
+  }
+
+  @Override
+  public byte[] convertToEntityAttribute(byte[] dbData) {
+    return dbData; // read as-is
+  }
+}


This converter is added but doesn't appear to be used by any of the new entities or services in this pull request. If it's not intended for immediate use, it should be removed to avoid adding unused code to the codebase.

gemini-code-assist · 2025-12-14T08:00:28Z

core/src/main/resources/META-INF/persistence.xml

+      <property name="hibernate.hikari.idleTimeout" value="30000"/>
+    </properties>
+  </persistence-unit>
+</persistence>


This file is missing a final newline character. It's a common convention and good practice to end files with a newline.

Suggested change

</persistence>

</persistence>

gemini-code-assist · 2025-12-14T08:00:28Z

core/src/main/resources/db/migration/h2/V1__Initial_schema.sql

+CREATE INDEX IF NOT EXISTS idx_events_session ON events(app_name, user_id, session_id);
+
+-- Index for sorting events by timestamp
+CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp);


This file is missing a final newline character. It's a common convention and good practice to end files with a newline.

Suggested change

CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp);

CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp);

gemini-code-assist · 2025-12-14T08:00:28Z

core/src/main/resources/db/migration/postgresql/V1__Initial_schema.sql

+CREATE INDEX IF NOT EXISTS idx_events_session ON events(app_name, user_id, session_id);
+
+-- Index for sorting events by timestamp
+CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp);


This file is missing a final newline character. It's a common convention and good practice to end files with a newline.

Suggested change

CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp);

CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp);

Summary Implements a production-ready database-backed session service that provides persistent storage for ADK sessions, events, and state using JPA/Hibernate. Key Features Hibernate 6.6 + HikariCP connection pooling for optimal performance Flyway migrations for schema versioning and zero-downtime deployments Multi-database support: PostgreSQL, MySQL, H2, and other RDBMS Thread-safe operations with pessimistic locking for concurrent updates Comprehensive test coverage with H2 in-memory database Dialect-aware JSON storage (JSONB for PostgreSQL, CLOB for others) Event filtering and pagination for efficient data retrieval Architecture Database dependencies are non-optional in the core module for ease of use, consistent with Python ADK's approach where `DatabaseSessionService` is in the main package. Note: Database drivers (PostgreSQL, MySQL) are marked as `true`, so users must add their specific driver dependency to their project.

gemini-code-assist bot reviewed Dec 8, 2025

View reviewed changes

prasadskarmarkar force-pushed the feature/database-session-service branch 4 times, most recently from fc948ac to 16e0ad4 Compare December 10, 2025 06:48

gemini-code-assist bot reviewed Dec 10, 2025

View reviewed changes

prasadskarmarkar force-pushed the feature/database-session-service branch 2 times, most recently from 89c8fe5 to fd2fdc7 Compare December 10, 2025 17:15

gemini-code-assist bot reviewed Dec 10, 2025

View reviewed changes

prasadskarmarkar force-pushed the feature/database-session-service branch 2 times, most recently from c1aa2b4 to 7a16f1a Compare December 10, 2025 22:22

gemini-code-assist bot reviewed Dec 10, 2025

View reviewed changes

prasadskarmarkar force-pushed the feature/database-session-service branch from 7a16f1a to 1fd831a Compare December 11, 2025 00:42

gemini-code-assist bot reviewed Dec 11, 2025

View reviewed changes

prasadskarmarkar force-pushed the feature/database-session-service branch 3 times, most recently from 7c46d73 to f0e2716 Compare December 13, 2025 00:54

prasadskarmarkar marked this pull request as draft December 13, 2025 00:59

gemini-code-assist bot reviewed Dec 13, 2025

View reviewed changes

prasadskarmarkar force-pushed the feature/database-session-service branch from 8030712 to 912d271 Compare December 14, 2025 07:53

gemini-code-assist bot reviewed Dec 14, 2025

View reviewed changes

prasadskarmarkar force-pushed the feature/database-session-service branch from ae36c9e to ac716fd Compare December 14, 2025 17:40

prasadskarmarkar force-pushed the feature/database-session-service branch 2 times, most recently from 3fe9316 to 0126f61 Compare December 14, 2025 18:55

prasadskarmarkar marked this pull request as ready for review December 14, 2025 18:56

prasadskarmarkar mentioned this pull request Dec 15, 2025

Implement Database-Backed Session Storage (DatabaseSessionService) for ADK Java #665

Open

prasadskarmarkar force-pushed the feature/database-session-service branch from 7265a33 to 78506b1 Compare December 15, 2025 18:44

prasadskarmarkar mentioned this pull request Dec 18, 2025

feat: Implementation of MySQL-backed session service for ADK #647

Open

prasadskarmarkar force-pushed the feature/database-session-service branch from 78506b1 to c8e0960 Compare December 19, 2025 20:34

prasadskarmarkar mentioned this pull request Dec 20, 2025

deleteSession() of the InMemorySessionService does not delete the user to session mapping #687

Open

	CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp);
	CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp);

feat: Add DatabaseSessionService with Hibernate and Flyway support #654

Are you sure you want to change the base?

feat: Add DatabaseSessionService with Hibernate and Flyway support #654

Uh oh!

Conversation

prasadskarmarkar commented Dec 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Key Features

Architecture

Files Changed

Test Plan

Uh oh!

google-cla bot commented Dec 8, 2025

Uh oh!

gemini-code-assist bot commented Dec 8, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

prasadskarmarkar commented Dec 10, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

prasadskarmarkar commented Dec 10, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

prasadskarmarkar commented Dec 10, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Dec 10, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

prasadskarmarkar commented Dec 11, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

prasadskarmarkar commented Dec 13, 2025

Uh oh!

gemini-code-assist bot left a comment

prasadskarmarkar commented Dec 8, 2025 •

edited

Loading