saveArtifact(
byte[] data = extractBytesFromPart(artifact);
String mimeType = extractMimeTypeFromPart(artifact);
- // Save to database with caller-provided metadata
- return dbHelper.saveArtifact(
- appName, userId, sessionId, filename, data, mimeType, metadata);
+ int version =
+ dbHelper.saveArtifact(
+ appName, userId, sessionId, filename, data, mimeType, metadata);
+ persistToS3AndKafka(
+ appName, userId, sessionId, filename, version, data, mimeType, metadata);
+ return version;
} catch (SQLException e) {
throw new RuntimeException("Failed to save artifact: " + e.getMessage(), e);
}
@@ -257,13 +289,127 @@ private byte[] extractBytesFromPart(Part part) {
*/
private String extractMimeTypeFromPart(Part part) {
if (part.inlineData() != null && part.inlineData().isPresent()) {
- return part.inlineData().get().mimeType().orElse("application/octet-stream");
+ return part.inlineData().get().mimeType().orElse(DEFAULT_MIME_TYPE);
}
- return "application/octet-stream"; // Default fallback
+ return DEFAULT_MIME_TYPE; // Default fallback
+ }
+
+ private void persistToS3AndKafka(
+ String appName,
+ String userId,
+ String sessionId,
+ String filename,
+ int version,
+ byte[] data,
+ String mimeType,
+ @Nullable String metadata) {
+ if (s3Client == null || s3Bucket == null) {
+ return;
+ }
+ String key = buildObjectKey(appName, userId, sessionId, filename, version);
+
+ // Construct an HTTPS URL instead of an s3:// path
+ String region = getenvOrNull("S3_REGION");
+ String fileUri;
+ if (region != null && !region.isBlank()) {
+ fileUri = String.format("https://%s.s3.%s.amazonaws.com/%s", s3Bucket, region, key);
+ } else {
+ // Fallback to global endpoint if region is not set
+ fileUri = String.format("https://%s.s3.amazonaws.com/%s", s3Bucket, key);
+ }
+
+ try {
+ PutObjectRequest.Builder requestBuilder =
+ PutObjectRequest.builder().bucket(s3Bucket).key(key).contentType(mimeType);
+ s3Client.putObject(requestBuilder.build(), RequestBody.fromBytes(data));
+ dbHelper.updateFileUri(appName, userId, sessionId, filename, version, fileUri);
+ publishKafkaEvent(
+ appName, userId, sessionId, filename, version, mimeType, metadata, fileUri, data.length);
+ } catch (S3Exception | SQLException e) {
+ // Fail-open: Postgres save already succeeded; log and continue.
+ System.err.println("⚠️ Failed to persist artifact to S3/Kafka: " + e.getMessage());
+ }
+ }
+
+ private void publishKafkaEvent(
+ String appName,
+ String userId,
+ String sessionId,
+ String filename,
+ int version,
+ String mimeType,
+ @Nullable String metadata,
+ String fileUri,
+ int sizeBytes) {
+ if (!kafkaEnabled || kafkaTopic == null || kafkaTopic.isBlank()) {
+ return;
+ }
+ JSONObject payload = new JSONObject();
+ payload.put("type", "artifact");
+ payload.put("appName", appName);
+ payload.put("userId", userId);
+ payload.put("sessionId", sessionId);
+ payload.put("filename", filename);
+ payload.put("version", version);
+ payload.put("mimeType", mimeType);
+ payload.put("fileUri", fileUri);
+ payload.put("sizeBytes", sizeBytes);
+ if (metadata != null) {
+ payload.put("metadata", metadata);
+ }
+ String key = appName + ":" + userId + ":" + sessionId + ":" + filename + ":" + version;
+ try {
+ KafkaWriter.PublishWithStatus(kafkaTopic, key, payload.toString());
+ } catch (Exception e) {
+ System.err.println("⚠️ Failed to publish artifact event to Kafka: " + e.getMessage());
+ }
+ }
+
+ private String buildObjectKey(
+ String appName, String userId, String sessionId, String filename, int version) {
+ String key;
+ if (filename != null && filename.startsWith("user:")) {
+ key = String.format("%s/%s/user/%s/%d", appName, userId, filename, version);
+ } else {
+ key = String.format("%s/%s/%s/%s/%d", appName, userId, sessionId, filename, version);
+ }
+
+ if (s3BasePath != null && !s3BasePath.isBlank()) {
+ // Ensure the base path doesn't end with a slash to avoid double slashes
+ String cleanBasePath =
+ s3BasePath.endsWith("/") ? s3BasePath.substring(0, s3BasePath.length() - 1) : s3BasePath;
+ return cleanBasePath + "/" + key;
+ }
+ return key;
+ }
+
+ private static @Nullable String getenvOrNull(String key) {
+ String value = System.getenv(key);
+ return value == null || value.isBlank() ? null : value;
+ }
+
+ private static S3Client buildS3Client() {
+ var builder = S3Client.builder();
+ String region = getenvOrNull("S3_REGION");
+ if (region != null) {
+ builder.region(Region.of(region));
+ }
+ String endpoint = getenvOrNull("S3_ENDPOINT");
+ if (endpoint != null) {
+ builder.endpointOverride(URI.create(endpoint));
+ }
+ String pathStyle = getenvOrNull("S3_PATH_STYLE");
+ if (pathStyle != null && Boolean.parseBoolean(pathStyle)) {
+ builder.serviceConfiguration(S3Configuration.builder().pathStyleAccessEnabled(true).build());
+ }
+ return builder.build();
}
/** Closes the database connection pool. Call this when shutting down the application. */
public void close() {
dbHelper.close();
+ if (s3Client != null) {
+ s3Client.close();
+ }
}
}
diff --git a/core/src/main/java/com/google/adk/artifacts/migration/PostgresToS3KafkaMigrator.java b/core/src/main/java/com/google/adk/artifacts/migration/PostgresToS3KafkaMigrator.java
new file mode 100644
index 000000000..4ac978893
--- /dev/null
+++ b/core/src/main/java/com/google/adk/artifacts/migration/PostgresToS3KafkaMigrator.java
@@ -0,0 +1,248 @@
+/*
+ * 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.artifacts.migration;
+
+import com.google.adk.kafka.consumer.KafkaWriter;
+import com.google.adk.utils.PropertiesHelper;
+import java.net.URI;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Timestamp;
+import java.time.Instant;
+import java.util.Optional;
+import org.json.JSONObject;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.core.sync.RequestBody;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.S3Configuration;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.S3Exception;
+
+/**
+ * Migrates artifacts from Postgres to S3 and publishes metadata to Kafka.
+ *
+ * Required environment variables:
+ *
+ *
+ * - DBURL
+ *
- DBUSER
+ *
- DBPASSWORD
+ *
- S3_BUCKET
+ *
+ *
+ * Optional environment variables:
+ *
+ *
+ * - ARTIFACTS_TABLE (defaults to "artifacts")
+ *
- S3_REGION (uses SDK default provider if unset)
+ *
- S3_ENDPOINT (for S3-compatible storage)
+ *
- S3_PATH_STYLE (true/false, for S3-compatible storage)
+ *
- KAFKA_TOPIC (defaults to kafka_topic in PropertiesHelper)
+ *
+ *
+ * Kafka message fields:
+ *
+ *
+ * - appName
+ *
- userId
+ *
- sessionId
+ *
- filename
+ *
- version
+ *
- metadata
+ *
- mimeType
+ *
- createdAt
+ *
- path (S3 URI)
+ *
+ */
+public final class PostgresToS3KafkaMigrator {
+ private static final Logger logger = LoggerFactory.getLogger(PostgresToS3KafkaMigrator.class);
+ private static final String DEFAULT_TABLE_NAME = "artifacts";
+
+ public static void main(String[] args) throws SQLException {
+ String dbUrl = requiredEnv("DBURL");
+ String dbUser = requiredEnv("DBUSER");
+ String dbPassword = requiredEnv("DBPASSWORD");
+ String bucketName = requiredEnv("S3_BUCKET");
+ String s3Prefix = Optional.ofNullable(System.getenv("S3_PREFIX")).orElse("");
+ String tableName =
+ Optional.ofNullable(System.getenv("ARTIFACTS_TABLE")).orElse(DEFAULT_TABLE_NAME);
+ int rowLimit = parseRowLimit(System.getenv("ROW_LIMIT"));
+ String kafkaTopic =
+ Optional.ofNullable(System.getenv("KAFKA_TOPIC"))
+ .orElse(PropertiesHelper.getInstance().getValue("kafka_topic"));
+ if (kafkaTopic == null || kafkaTopic.isBlank()) {
+ throw new IllegalStateException(
+ "Kafka topic not configured. Set KAFKA_TOPIC or kafka_topic.");
+ }
+
+ try (S3Client s3Client = buildS3Client();
+ Connection connection = DriverManager.getConnection(dbUrl, dbUser, dbPassword)) {
+ connection.setAutoCommit(false);
+ String sql =
+ String.format(
+ "SELECT app_name, user_id, session_id, filename, version, mime_type, data, metadata, created_at "
+ + "FROM %s ORDER BY app_name, user_id, session_id, filename, version",
+ tableName);
+ if (rowLimit > 0) {
+ sql = sql + " LIMIT " + rowLimit;
+ logger.info("Limiting migration to {} rows.", rowLimit);
+ }
+ try (PreparedStatement statement = connection.prepareStatement(sql)) {
+ statement.setFetchSize(500);
+ int migrated = 0;
+ try (ResultSet rs = statement.executeQuery()) {
+ while (rs.next()) {
+ String appName = rs.getString("app_name");
+ String userId = rs.getString("user_id");
+ String sessionId = rs.getString("session_id");
+ String filename = rs.getString("filename");
+ int version = rs.getInt("version");
+ String mimeType = rs.getString("mime_type");
+ byte[] data = rs.getBytes("data");
+ String metadataJson = rs.getString("metadata");
+ Timestamp createdAt = rs.getTimestamp("created_at");
+
+ String key = buildObjectKey(s3Prefix, appName, userId, sessionId, filename, version);
+ PutObjectRequest.Builder requestBuilder =
+ PutObjectRequest.builder().bucket(bucketName).key(key);
+ if (mimeType != null && !mimeType.isBlank()) {
+ requestBuilder.contentType(mimeType);
+ }
+
+ try {
+ s3Client.putObject(requestBuilder.build(), RequestBody.fromBytes(data));
+ } catch (S3Exception e) {
+ throw new RuntimeException("Failed to upload artifact to S3 for key: " + key, e);
+ }
+
+ String s3Uri = "s3://" + bucketName + "/" + key;
+ JSONObject payload = new JSONObject();
+ payload.put("appName", appName);
+ payload.put("userId", userId);
+ payload.put("sessionId", sessionId);
+ payload.put("filename", filename);
+ payload.put("version", version);
+ if (metadataJson != null) {
+ payload.put("metadata", metadataJson);
+ }
+ if (mimeType != null) {
+ payload.put("mimeType", mimeType);
+ }
+ payload.put("createdAt", formatTimestamp(createdAt));
+ payload.put("path", s3Uri);
+
+ String keyForKafka =
+ appName + ":" + userId + ":" + sessionId + ":" + filename + ":" + version;
+ try {
+ KafkaWriter.PublishWithStatus(kafkaTopic, keyForKafka, payload.toString());
+ } catch (Exception e) {
+ throw new RuntimeException(
+ "Failed to publish Kafka event for key: " + keyForKafka, e);
+ }
+
+ migrated++;
+ if (migrated % 500 == 0) {
+ logger.info("Migrated {} artifacts so far...", migrated);
+ }
+ }
+ }
+ logger.info("Migration complete. Total artifacts migrated: {}", migrated);
+ }
+ }
+ }
+
+ private static S3Client buildS3Client() {
+ var builder = S3Client.builder();
+ String region = System.getenv("S3_REGION");
+ if (region != null && !region.isBlank()) {
+ builder.region(Region.of(region));
+ }
+ String endpoint = System.getenv("S3_ENDPOINT");
+ if (endpoint != null && !endpoint.isBlank()) {
+ builder.endpointOverride(URI.create(endpoint));
+ }
+ String pathStyle = System.getenv("S3_PATH_STYLE");
+ if (pathStyle != null && Boolean.parseBoolean(pathStyle)) {
+
+ builder.serviceConfiguration(S3Configuration.builder().pathStyleAccessEnabled(true).build());
+ }
+ return builder.build();
+ }
+
+ private static String requiredEnv(String key) {
+ String value = System.getenv(key);
+ if (value == null || value.isBlank()) {
+ throw new IllegalStateException("Missing required environment variable: " + key);
+ }
+ return value;
+ }
+
+ private static boolean fileHasUserNamespace(String filename) {
+ return filename != null && filename.startsWith("user:");
+ }
+
+ private static String buildObjectKey(
+ String prefix,
+ String appName,
+ String userId,
+ String sessionId,
+ String filename,
+ int version) {
+ if (fileHasUserNamespace(filename)) {
+ return String.format(
+ "%s%s/%s/user/%s/%d", normalizePrefix(prefix), appName, userId, filename, version);
+ }
+ return String.format(
+ "%s%s/%s/%s/%s/%d", normalizePrefix(prefix), appName, userId, sessionId, filename, version);
+ }
+
+ private static String formatTimestamp(Timestamp createdAt) {
+ if (createdAt == null) {
+ return Instant.EPOCH.toString();
+ }
+ return createdAt.toInstant().toString();
+ }
+
+ private static int parseRowLimit(String rawLimit) {
+ if (rawLimit == null || rawLimit.isBlank()) {
+ return 0;
+ }
+ try {
+ return Integer.parseInt(rawLimit.trim());
+ } catch (NumberFormatException ex) {
+ throw new IllegalStateException("ROW_LIMIT must be an integer.", ex);
+ }
+ }
+
+ private static String normalizePrefix(String prefix) {
+ if (prefix == null || prefix.isBlank()) {
+ return "";
+ }
+ String trimmed = prefix.trim();
+ if (trimmed.endsWith("/")) {
+ return trimmed;
+ }
+ return trimmed + "/";
+ }
+
+ private PostgresToS3KafkaMigrator() {}
+}
diff --git a/core/src/main/java/com/google/adk/store/PostgresArtifactStore.java b/core/src/main/java/com/google/adk/store/PostgresArtifactStore.java
index d31be8676..3c990c520 100644
--- a/core/src/main/java/com/google/adk/store/PostgresArtifactStore.java
+++ b/core/src/main/java/com/google/adk/store/PostgresArtifactStore.java
@@ -247,13 +247,40 @@ public int saveArtifact(
String mimeType,
String metadata)
throws SQLException {
+ return saveArtifact(appName, userId, sessionId, filename, data, mimeType, metadata, null);
+ }
+
+ /**
+ * Save artifact to database with metadata and file URI. Returns the assigned version number.
+ *
+ * @param appName the application name
+ * @param userId the user ID
+ * @param sessionId the session ID
+ * @param filename the artifact filename
+ * @param data the artifact binary data
+ * @param mimeType the MIME type
+ * @param metadata the metadata JSON string (can be null)
+ * @param fileUri the URI pointing to external storage like S3 (can be null)
+ * @return the version number assigned to this artifact
+ * @throws SQLException if save operation fails
+ */
+ public int saveArtifact(
+ String appName,
+ String userId,
+ String sessionId,
+ String filename,
+ byte[] data,
+ String mimeType,
+ String metadata,
+ String fileUri)
+ throws SQLException {
logger.debug(
"Saving artifact: app={}, user={}, session={}, file={}, size={}KB, mime={}",
appName,
userId,
sessionId,
filename,
- data.length / 1024,
+ data != null ? data.length / 1024 : 0,
mimeType);
Connection conn = null;
@@ -267,8 +294,8 @@ public int saveArtifact(
String sql =
String.format(
- "INSERT INTO %s (app_name, user_id, session_id, filename, version, mime_type, data, metadata) "
- + "VALUES (?, ?, ?, ?, ?, ?, ?, ?::jsonb)",
+ "INSERT INTO %s (app_name, user_id, session_id, filename, version, mime_type, data, metadata, file_uri) "
+ + "VALUES (?, ?, ?, ?, ?, ?, ?, ?::jsonb, ?)",
tableName);
try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
@@ -278,8 +305,21 @@ public int saveArtifact(
pstmt.setString(4, filename);
pstmt.setInt(5, nextVersion);
pstmt.setString(6, mimeType);
- pstmt.setBytes(7, data);
- pstmt.setString(8, metadata);
+ if (data != null) {
+ pstmt.setBytes(7, data);
+ } else {
+ pstmt.setNull(7, java.sql.Types.BINARY);
+ }
+ if (metadata != null) {
+ pstmt.setString(8, metadata);
+ } else {
+ pstmt.setNull(8, java.sql.Types.OTHER);
+ }
+ if (fileUri != null) {
+ pstmt.setString(9, fileUri);
+ } else {
+ pstmt.setNull(9, java.sql.Types.VARCHAR);
+ }
int rowsAffected = pstmt.executeUpdate();
@@ -335,6 +375,46 @@ public int saveArtifact(
}
}
+ /**
+ * Update the file URI for an existing artifact version.
+ *
+ * @param appName the application name
+ * @param userId the user ID
+ * @param sessionId the session ID
+ * @param filename the artifact filename
+ * @param version the artifact version
+ * @param fileUri the URI pointing to external storage (S3)
+ * @throws SQLException if update fails
+ */
+ public void updateFileUri(
+ String appName, String userId, String sessionId, String filename, int version, String fileUri)
+ throws SQLException {
+ String sql =
+ String.format(
+ "UPDATE %s SET file_uri = ? WHERE app_name = ? AND user_id = ? AND session_id = ? AND filename = ? AND version = ?",
+ tableName);
+ try (Connection conn = getConnection();
+ PreparedStatement pstmt = conn.prepareStatement(sql)) {
+ pstmt.setString(1, fileUri);
+ pstmt.setString(2, appName);
+ pstmt.setString(3, userId);
+ pstmt.setString(4, sessionId);
+ pstmt.setString(5, filename);
+ pstmt.setInt(6, version);
+ pstmt.executeUpdate();
+ } catch (SQLException e) {
+ logger.error(
+ "❌ Error updating file_uri: app={}, user={}, session={}, file={}, version={}, error={}",
+ appName,
+ userId,
+ sessionId,
+ filename,
+ version,
+ e.getMessage());
+ throw e;
+ }
+ }
+
/**
* Get next version number for an artifact with row-level locking to prevent race conditions.
*
@@ -421,14 +501,14 @@ public ArtifactData loadArtifact(
// Load specific version
sql =
String.format(
- "SELECT data, mime_type, version, created_at, metadata FROM %s "
+ "SELECT data, mime_type, version, created_at, metadata, file_uri FROM %s "
+ "WHERE app_name = ? AND user_id = ? AND session_id = ? AND filename = ? AND version = ?",
tableName);
} else {
// Load latest version
sql =
String.format(
- "SELECT data, mime_type, version, created_at, metadata FROM %s "
+ "SELECT data, mime_type, version, created_at, metadata, file_uri FROM %s "
+ "WHERE app_name = ? AND user_id = ? AND session_id = ? AND filename = ? "
+ "ORDER BY version DESC LIMIT 1",
tableName);
@@ -451,6 +531,7 @@ public ArtifactData loadArtifact(
int loadedVersion = rs.getInt("version");
Timestamp createdAt = rs.getTimestamp("created_at");
String metadata = rs.getString("metadata");
+ String fileUri = rs.getString("file_uri");
logger.info(
"✅ Artifact loaded: app={}, user={}, session={}, file={}, version={}, size={}KB",
@@ -459,9 +540,9 @@ public ArtifactData loadArtifact(
sessionId,
filename,
loadedVersion,
- data.length / 1024);
+ data != null ? data.length / 1024 : 0);
- return new ArtifactData(data, mimeType, loadedVersion, createdAt, metadata);
+ return new ArtifactData(data, mimeType, loadedVersion, createdAt, metadata, fileUri);
} else {
logger.warn(
"⚠️ Artifact not found: app={}, user={}, session={}, file={}, version={}",
@@ -668,14 +749,26 @@ public static class ArtifactData {
public final int version;
public final Timestamp createdAt;
public final String metadata;
+ public final String fileUri;
public ArtifactData(
byte[] data, String mimeType, int version, Timestamp createdAt, String metadata) {
+ this(data, mimeType, version, createdAt, metadata, null);
+ }
+
+ public ArtifactData(
+ byte[] data,
+ String mimeType,
+ int version,
+ Timestamp createdAt,
+ String metadata,
+ String fileUri) {
this.data = data;
this.mimeType = mimeType;
this.version = version;
this.createdAt = createdAt;
this.metadata = metadata;
+ this.fileUri = fileUri;
}
}
}