From c140868bb5c7d1f0edd586fb2ca55bc613caf5d4 Mon Sep 17 00:00:00 2001
From: Denis DelGrosso <85250797+ddelgrosso1@users.noreply.github.com>
Date: Wed, 29 Nov 2023 13:56:19 -0500
Subject: [PATCH] feat: implement object retention lock for bucket / files
 (#2365)

---
 src/bucket.ts          |  3 +++
 src/file.ts            | 12 ++++++++++-
 src/storage.ts         |  8 ++++++++
 system-test/storage.ts | 45 ++++++++++++++++++++++++++++++++++++++++++
 test/file.ts           | 14 +++++++++++++
 test/index.ts          | 11 +++++++++++
 6 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/src/bucket.ts b/src/bucket.ts
index 8780f2e65..34df1a64a 100644
--- a/src/bucket.ts
+++ b/src/bucket.ts
@@ -327,6 +327,9 @@ export interface BucketMetadata extends BaseMetadata {
   };
   metageneration?: string;
   name?: string;
+  objectRetention?: {
+    mode?: string;
+  };
   owner?: {
     entity?: string;
     entityId?: string;
diff --git a/src/file.ts b/src/file.ts
index e17630576..d2274dd57 100644
--- a/src/file.ts
+++ b/src/file.ts
@@ -423,6 +423,7 @@ export interface FileMetadata extends BaseMetadata {
   };
   customTime?: string;
   eventBasedHold?: boolean | null;
+  readonly eventBasedHoldReleaseTime?: string;
   generation?: string | number;
   kmsKeyName?: string;
   md5Hash?: string;
@@ -436,6 +437,10 @@ export interface FileMetadata extends BaseMetadata {
     entity?: string;
     entityId?: string;
   };
+  retention?: {
+    retainUntilTime?: string;
+    mode?: string;
+  } | null;
   retentionExpirationTime?: string;
   size?: string | number;
   storageClass?: string;
@@ -3813,7 +3818,8 @@ class File extends ServiceObject<File, FileMetadata> {
     optionsOrCallback: SetMetadataOptions | MetadataCallback<FileMetadata>,
     cb?: MetadataCallback<FileMetadata>
   ): Promise<SetMetadataResponse<FileMetadata>> | void {
-    const options =
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const options: any =
       typeof optionsOrCallback === 'object' ? optionsOrCallback : {};
     cb =
       typeof optionsOrCallback === 'function'
@@ -3826,6 +3832,10 @@ class File extends ServiceObject<File, FileMetadata> {
       options
     );
 
+    if (metadata.retention !== undefined) {
+      options.overrideUnlockedRetention = true;
+    }
+
     super
       .setMetadata(metadata, options)
       .then(resp => cb!(null, ...resp))
diff --git a/src/storage.ts b/src/storage.ts
index a86f3a13b..e6f251acc 100644
--- a/src/storage.ts
+++ b/src/storage.ts
@@ -48,6 +48,7 @@ export interface GetServiceAccountCallback {
 export interface CreateBucketQuery {
   project: string;
   userProject: string;
+  enableObjectRetention: boolean;
 }
 
 export enum IdempotencyStrategy {
@@ -121,6 +122,7 @@ export interface CreateBucketRequest {
   cors?: Cors[];
   customPlacementConfig?: CustomPlacementConfig;
   dra?: boolean;
+  enableObjectRetention?: boolean;
   location?: string;
   multiRegional?: boolean;
   nearline?: boolean;
@@ -862,6 +864,7 @@ export class Storage extends Service {
    *     For more information, see {@link https://cloud.google.com/storage/docs/locations| Bucket Locations}.
    * @property {boolean} [dra=false] Specify the storage class as Durable Reduced
    *     Availability.
+   * @property {boolean} [enableObjectRetention=false] Specifiy whether or not object retention should be enabled on this bucket.
    * @property {string} [location] Specify the bucket's location. If specifying
    *     a dual-region, the `customPlacementConfig` property should be set in conjunction.
    *     For more information, see {@link https://cloud.google.com/storage/docs/locations| Bucket Locations}.
@@ -1023,6 +1026,11 @@ export class Storage extends Service {
       delete body.userProject;
     }
 
+    if (body.enableObjectRetention) {
+      query.enableObjectRetention = body.enableObjectRetention;
+      delete body.enableObjectRetention;
+    }
+
     this.request(
       {
         method: 'POST',
diff --git a/system-test/storage.ts b/system-test/storage.ts
index e4b384430..3a010da10 100644
--- a/system-test/storage.ts
+++ b/system-test/storage.ts
@@ -1602,6 +1602,51 @@ describe('storage', () => {
     });
   });
 
+  describe('object retention lock', () => {
+    const fileName = generateName();
+    let objectRetentionBucket: Bucket;
+
+    before(async () => {
+      objectRetentionBucket = storage.bucket(generateName());
+    });
+
+    after(async () => {
+      await objectRetentionBucket.deleteFiles({force: true});
+      await objectRetentionBucket.delete();
+    });
+
+    it('should create a bucket with object retention enabled', async () => {
+      const result = await objectRetentionBucket.create({
+        enableObjectRetention: true,
+      });
+
+      assert.deepStrictEqual(result[0].metadata.objectRetention, {
+        mode: 'Enabled',
+      });
+    });
+
+    it('should create a file with object retention enabled', async () => {
+      const time = new Date();
+      time.setMinutes(time.getSeconds() + 1);
+      const retention = {mode: 'Unlocked', retainUntilTime: time.toISOString()};
+      const file = new File(objectRetentionBucket, fileName);
+      await objectRetentionBucket.upload(FILES.big.path, {
+        metadata: {
+          retention,
+        },
+        destination: fileName,
+      });
+      const [metadata] = await file.getMetadata();
+      assert.deepStrictEqual(metadata.retention, retention);
+    });
+
+    it('should disable object retention on the file', async () => {
+      const file = new File(objectRetentionBucket, fileName);
+      const [metadata] = await file.setMetadata({retention: null});
+      assert.strictEqual(metadata.retention, undefined);
+    });
+  });
+
   describe('requester pays', () => {
     const HAS_2ND_PROJECT =
       process.env.GCN_STORAGE_2ND_PROJECT_ID !== undefined;
diff --git a/test/file.ts b/test/file.ts
index 5ca9f847b..76d488404 100644
--- a/test/file.ts
+++ b/test/file.ts
@@ -4567,6 +4567,20 @@ describe('File', () => {
     });
   });
 
+  describe('setMetadata', () => {
+    it('should set query parameter overrideUnlockedRetention', done => {
+      const newFile = new File(BUCKET, 'new-file');
+
+      newFile.parent.request = (reqOpts: DecorateRequestOptions) => {
+        console.log(reqOpts.qs);
+        assert.strictEqual(reqOpts.qs.overrideUnlockedRetention, true);
+        done();
+      };
+
+      newFile.setMetadata({retention: null}, assert.ifError);
+    });
+  });
+
   describe('setStorageClass', () => {
     const STORAGE_CLASS = 'new_storage_class';
 
diff --git a/test/index.ts b/test/index.ts
index 8e39ae28d..610cf0d54 100644
--- a/test/index.ts
+++ b/test/index.ts
@@ -905,6 +905,17 @@ describe('Storage', () => {
       }, /Both `coldline` and `storageClass` were provided./);
     });
 
+    it('should allow enabling object retention', done => {
+      storage.request = (
+        reqOpts: DecorateRequestOptions,
+        callback: Function
+      ) => {
+        assert.strictEqual(reqOpts.qs.enableObjectRetention, true);
+        callback();
+      };
+      storage.createBucket(BUCKET_NAME, {enableObjectRetention: true}, done);
+    });
+
     describe('storage classes', () => {
       it('should expand metadata.archive', done => {
         storage.request = (reqOpts: DecorateRequestOptions) => {