From 651d39ead6d5ba01ba9dd6d08a7f537973ab2a19 Mon Sep 17 00:00:00 2001 From: Yashas Shetty Date: Fri, 29 May 2026 12:42:22 +0530 Subject: [PATCH 1/2] Enhance PostgresArtifactService with S3 and Kafka integration --- core/pom.xml | 5 + .../PostgresArtifactService.GUIDE.md | 75 ++++++ .../artifacts/PostgresArtifactService.java | 140 +++++++++- .../migration/PostgresToS3KafkaMigrator.java | 248 ++++++++++++++++++ .../adk/store/PostgresArtifactStore.java | 111 +++++++- 5 files changed, 561 insertions(+), 18 deletions(-) create mode 100644 core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.GUIDE.md create mode 100644 core/src/main/java/com/google/adk/artifacts/migration/PostgresToS3KafkaMigrator.java diff --git a/core/pom.xml b/core/pom.xml index c770927cb..a5dfb75c4 100644 --- a/core/pom.xml +++ b/core/pom.xml @@ -66,6 +66,11 @@ com.google.cloud google-cloud-storage + + software.amazon.awssdk + s3 + 2.42.31 + com.google.genai google-genai diff --git a/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.GUIDE.md b/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.GUIDE.md new file mode 100644 index 000000000..8cafc9aaa --- /dev/null +++ b/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.GUIDE.md @@ -0,0 +1,75 @@ +# PostgresArtifactService Guide + +## What It Does + +`PostgresArtifactService` stores artifacts in PostgreSQL and can optionally mirror them to S3 +and emit Kafka events. It is the runtime artifact service used by the application. + +Core behavior: +- Always writes artifact bytes and metadata to Postgres. +- Uploads bytes to S3 when `S3_BUCKET` is set. +- Publishes a Kafka event when `use_kafka=true` and `kafka_topic` is set. + +## Save Flow (Write Path) + +1. Extract bytes and MIME type from the `Part`. +2. Save to Postgres (data + metadata). +3. If S3 is enabled, upload bytes to `s3:///`. +4. If Kafka is enabled, publish a JSON event describing the artifact. + +## Configuration + +### Postgres (required) + +Environment variables: +- `DBURL` +- `DBUSER` +- `DBPASSWORD` + +### S3 (optional) + +Environment variables: +- `S3_BUCKET` (enables S3 upload) +- `S3_REGION` (recommended) +- `S3_ENDPOINT` (optional, for S3-compatible storage) +- `S3_PATH_STYLE` (`true` for path-style access) + +### Kafka (optional) + +Properties (via `PropertiesHelper`): +- `use_kafka=true` +- `kafka_topic=` + +## S3 Object Key Format + +``` +//// +``` + +If `filename` starts with `user:`, the key becomes: + +``` +//user// +``` + +## Kafka Event Payload + +Published fields: +- `type` = `artifact` +- `appName`, `userId`, `sessionId`, `filename`, `version` +- `mimeType` +- `fileUri` (S3 URI) +- `sizeBytes` +- `metadata` (only included when non-null) + +## Common Checks + +- If S3 upload fails, the Postgres write still succeeds (fail-open). +- If Kafka is disabled or topic is missing, no event is published. +- Missing `S3_REGION` or invalid AWS credentials will cause S3 failures. + +## Related Files + +- `PostgresArtifactService.java` +- `PostgresArtifactStore.java` +- `S3ArtifactService.java` diff --git a/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.java b/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.java index 74c60bd8e..e53a42255 100644 --- a/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.java +++ b/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.java @@ -16,17 +16,27 @@ package com.google.adk.artifacts; +import com.google.adk.kafka.consumer.KafkaWriter; import com.google.adk.store.PostgresArtifactStore; import com.google.adk.store.PostgresArtifactStore.ArtifactData; +import com.google.adk.utils.PropertiesHelper; import com.google.common.collect.ImmutableList; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.schedulers.Schedulers; +import java.net.URI; import java.sql.SQLException; import java.util.List; +import org.json.JSONObject; import org.jspecify.annotations.Nullable; +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; /** * A PostgreSQL-backed implementation of the {@link BaseArtifactService}. @@ -63,7 +73,12 @@ public final class PostgresArtifactService implements BaseArtifactService { private static final String DEFAULT_TABLE_NAME = "artifacts"; + private static final String DEFAULT_MIME_TYPE = "application/octet-stream"; private final PostgresArtifactStore dbHelper; + private final @Nullable S3Client s3Client; + private final @Nullable String s3Bucket; + private final boolean kafkaEnabled; + private final @Nullable String kafkaTopic; /** * Creates a new PostgresArtifactService using environment variables for database connection. Uses @@ -80,6 +95,10 @@ public final class PostgresArtifactService implements BaseArtifactService { */ public PostgresArtifactService() { this.dbHelper = PostgresArtifactStore.getInstance(DEFAULT_TABLE_NAME); + this.s3Bucket = getenvOrNull("S3_BUCKET"); + this.s3Client = s3Bucket != null ? buildS3Client() : null; + this.kafkaEnabled = Boolean.parseBoolean(PropertiesHelper.getInstance().getValue("use_kafka")); + this.kafkaTopic = PropertiesHelper.getInstance().getValue("kafka_topic"); } /** @@ -95,6 +114,10 @@ public PostgresArtifactService() { public PostgresArtifactService(String dbUrl, String dbUser, String dbPassword) { this.dbHelper = PostgresArtifactStore.createInstance(dbUrl, dbUser, dbPassword, DEFAULT_TABLE_NAME); + this.s3Bucket = getenvOrNull("S3_BUCKET"); + this.s3Client = s3Bucket != null ? buildS3Client() : null; + this.kafkaEnabled = Boolean.parseBoolean(PropertiesHelper.getInstance().getValue("use_kafka")); + this.kafkaTopic = PropertiesHelper.getInstance().getValue("kafka_topic"); } @Override @@ -107,10 +130,12 @@ public Single saveArtifact( byte[] data = extractBytesFromPart(artifact); String mimeType = extractMimeTypeFromPart(artifact); - // Save to database without metadata (metadata = null) - // Applications should use saveArtifact(..., metadata) if they need custom metadata - return dbHelper.saveArtifact( - appName, userId, sessionId, filename, data, mimeType, null); + int version = + dbHelper.saveArtifact( + appName, userId, sessionId, filename, data, mimeType, null); + persistToS3AndKafka( + appName, userId, sessionId, filename, version, data, mimeType, null); + return version; } catch (SQLException e) { throw new RuntimeException("Failed to save artifact: " + e.getMessage(), e); } @@ -154,9 +179,12 @@ public Single 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 +285,107 @@ 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); + String fileUri = "s3://" + 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) { + if (filename != null && filename.startsWith("user:")) { + return String.format("%s/%s/user/%s/%d", appName, userId, filename, version); + } + return String.format("%s/%s/%s/%s/%d", appName, userId, sessionId, filename, version); + } + + 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; } } } From 3b912530c7154be61f2ce7a450fe53b5bf43d040 Mon Sep 17 00:00:00 2001 From: "alfred.jimmy" Date: Sat, 20 Jun 2026 11:30:57 +0530 Subject: [PATCH 2/2] added flexibility to add any s3 path, store https object url --- .gitignore | 2 ++ .../artifacts/PostgresArtifactService.java | 30 +++++++++++++++++-- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/.gitignore b/.gitignore index 09f3849bf..4f6801095 100644 --- a/.gitignore +++ b/.gitignore @@ -33,3 +33,5 @@ Thumbs.db # Local documentation and plans docs/ plans/ + +contrib/samples/a2a_basic/bin/ \ No newline at end of file diff --git a/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.java b/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.java index e53a42255..b95378cea 100644 --- a/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.java +++ b/core/src/main/java/com/google/adk/artifacts/PostgresArtifactService.java @@ -80,6 +80,8 @@ public final class PostgresArtifactService implements BaseArtifactService { private final boolean kafkaEnabled; private final @Nullable String kafkaTopic; + private final @Nullable String s3BasePath; + /** * Creates a new PostgresArtifactService using environment variables for database connection. Uses * the default "artifacts" table. Per JVM, only one table is used for all artifact operations, @@ -96,6 +98,7 @@ public final class PostgresArtifactService implements BaseArtifactService { public PostgresArtifactService() { this.dbHelper = PostgresArtifactStore.getInstance(DEFAULT_TABLE_NAME); this.s3Bucket = getenvOrNull("S3_BUCKET"); + this.s3BasePath = getenvOrNull("S3_BASE_PATH"); this.s3Client = s3Bucket != null ? buildS3Client() : null; this.kafkaEnabled = Boolean.parseBoolean(PropertiesHelper.getInstance().getValue("use_kafka")); this.kafkaTopic = PropertiesHelper.getInstance().getValue("kafka_topic"); @@ -115,6 +118,7 @@ public PostgresArtifactService(String dbUrl, String dbUser, String dbPassword) { this.dbHelper = PostgresArtifactStore.createInstance(dbUrl, dbUser, dbPassword, DEFAULT_TABLE_NAME); this.s3Bucket = getenvOrNull("S3_BUCKET"); + this.s3BasePath = getenvOrNull("S3_BASE_PATH"); this.s3Client = s3Bucket != null ? buildS3Client() : null; this.kafkaEnabled = Boolean.parseBoolean(PropertiesHelper.getInstance().getValue("use_kafka")); this.kafkaTopic = PropertiesHelper.getInstance().getValue("kafka_topic"); @@ -303,7 +307,17 @@ private void persistToS3AndKafka( return; } String key = buildObjectKey(appName, userId, sessionId, filename, version); - String fileUri = "s3://" + s3Bucket + "/" + key; + + // 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); @@ -353,10 +367,20 @@ private void publishKafkaEvent( private String buildObjectKey( String appName, String userId, String sessionId, String filename, int version) { + String key; if (filename != null && filename.startsWith("user:")) { - return String.format("%s/%s/user/%s/%d", appName, userId, filename, version); + 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 String.format("%s/%s/%s/%s/%d", appName, userId, sessionId, filename, version); + return key; } private static @Nullable String getenvOrNull(String key) {