S3: Implement restore object with select job type

moto/s3/exceptions.py

@@ -604,3 +604,23 @@ def __init__(self) -> None:
             "DaysMustProvidedExceptForSelectRequest",
             "`Days` must be provided except for select requests",
         )
+
+
+class InvalidRestoreRequestType(S3ClientError):
+    code = 400
+
+    def __init__(self) -> None:
+        super().__init__(
+            "InvalidRestoreRequestType",
+            "Invalid type for restore request. It must be `SELECT`",
+        )
+
+
+class MissingRequiredParametersForSelectRequest(S3ClientError):
+    code = 400
+
+    def __init__(self, parameter_name: str) -> None:
+        super().__init__(
+            "MissingRequiredParametersForSelectRequest",
+            f'Missing required parameter in RestoreRequest: "{parameter_name}" for select request',
+        )

moto/s3/models.py

@@ -48,12 +48,14 @@
     InvalidPart,
     InvalidPublicAccessBlockConfiguration,
     InvalidRequest,
+    InvalidRestoreRequestType,
     InvalidStorageClass,
     InvalidTagError,
     InvalidTargetBucketForLogging,
     MalformedXML,
     MissingBucket,
     MissingKey,
+    MissingRequiredParametersForSelectRequest,
     NoSuchPublicAccessBlockConfiguration,
     NoSuchUpload,
     ObjectLockConfigurationNotFoundError,
@@ -2885,7 +2887,13 @@
         ]
 
     def restore_object(
-        self, bucket_name: str, key_name: str, days: Optional[str], type_: Optional[str]
+        self,
+        bucket_name: str,
+        key_name: str,
+        days: Optional[str],
+        type_: Optional[str],
+        select_parameters: Optional[Dict[str, Any]],
+        output_location: Optional[Dict[str, Any]],
     ) -> bool:
         key = self.get_object(bucket_name, key_name)
         if not key:
@@ -2897,8 +2905,72 @@
         if days and type_:
             raise DaysMustNotProvidedForSelectRequest()
 
+        if type_ and type_ != "SELECT":
+            raise InvalidRestoreRequestType()
+
+        if type_ and select_parameters is None:
+            raise MissingRequiredParametersForSelectRequest("SelectParameters")
+
+        if type_ and output_location is None:
+            raise MissingRequiredParametersForSelectRequest("OutputLocation")
+
         if key.storage_class not in ARCHIVE_STORAGE_CLASSES:
             raise InvalidObjectState(storage_class=key.storage_class)
+
+        # select request
+        if type_:
+            restored_key = copy.deepcopy(key)
+            output_location_s3 = output_location["S3"]
+
+            # TODO: 'Encryption', 'CannedACL', 'AccessControlList', 'Tagging' and 'UserMetadata' are need to be supported for OutputLocation.S3 configuration.
+            # set bucket name
+            _bucket_name = output_location_s3.get("BucketName")
+            if _bucket_name is not None:
+                # check if bucket is exists here.
+                bucket = self.get_bucket(_bucket_name)
+            else:
+                bucket = self.bucket(key.bucket_name)
+            restored_key.bucket_name = bucket.name
+
+            # set key name
+            _key_name = output_location_s3.get("Prefix")
+            if _key_name is not None:
+                restored_key.name = _key_name
+                key_name = _key_name
+            else:
+                key_name = key.name
+
+            # set storage class
+            _storage = output_location_s3.get("StorageClass")
+            if _storage is not None:
+                restored_key.set_storage_class(_storage)
+
+            # set value
+            contents = self.select_object_content(
+                key.bucket_name,
+                key.name,
+                select_parameters["Expression"],
+                select_parameters["InputSerialization"],
+            )
+
+            # serialize selected object
+            # TODO: Delimiter extraction is also defined in `serialize_select` method in `moto/s3/select_object_content.py`. This part might be extracted into a common method.
+            delimiter = (
+                (select_parameters["OutputSerialization"].get("JSON") or {}).get(
+                    "RecordDelimiter"
+                )
+                or "\n"
+            ).encode("utf-8")
+            payload_bytes = b""
+            for data in contents:
+                payload_bytes += data + delimiter
+            restored_key.value = payload_bytes
+
+            # save object to bucket
+            bucket.keys.setlist(key_name, restored_key)
+            return True
+
+        # restore object request
         had_expiry_date = key.expiry_date is not None
         if days:
             key.restore(int(days))

moto/s3/responses.py

@@ -2278,6 +2278,8 @@ def _key_response_post(
                 key_name,
                 params.get("Days", None),
                 params.get("Type", None),
+                params.get("SelectParameters", None),
+                params.get("OutputLocation", None),
             )
             status_code = 200 if previously_restored else 202
             return status_code, {}, ""

tests/test_s3/test_s3.py

@@ -596,6 +596,55 @@ def test_restore_key_transition():
     state_manager.unset_transition(model_name="s3::keyrestore")
 
 
+@freeze_time("2012-01-01 12:00:00")
+@mock_aws
+def test_restore_key_with_select():
+    if not settings.TEST_DECORATOR_MODE:
+        raise SkipTest("Can't set transition directly in ServerMode")
+
+    state_manager.set_transition(
+        model_name="s3::keyrestore", transition={"progression": "manual", "times": 1}
+    )
+
+    s3_resource = boto3.resource("s3", region_name=DEFAULT_REGION_NAME)
+    bucket_name = "foobar"
+    bucket = s3_resource.Bucket(bucket_name)
+    bucket.create()
+
+    key = bucket.put_object(
+        Key="the-key", Body=json.dumps({"a1": "b1", "a2": "b2"}), StorageClass="GLACIER"
+    )
+    assert key.restore is None
+
+    # restore successfully completed
+    restored_key_name = "restored.json"
+    key.restore_object(
+        RestoreRequest={
+            "Type": "SELECT",
+            "SelectParameters": {
+                "Expression": "SELECT COUNT(*) as cnt FROM S3Object",
+                "ExpressionType": "SQL",
+                "InputSerialization": {"JSON": {"Type": "DOCUMENT"}},
+                "OutputSerialization": {"JSON": {"RecordDelimiter": ","}},
+            },
+            "OutputLocation": {
+                "S3": {
+                    "BucketName": bucket_name,
+                    "Prefix": restored_key_name,
+                    "StorageClass": "STANDARD",
+                }
+            },
+        }
+    )
+    resp = boto3.client("s3", region_name=DEFAULT_REGION_NAME).get_object(
+        Bucket=bucket_name, Key=restored_key_name
+    )
+    assert resp["ContentLength"] == 9
+    assert resp["Body"].read() == b'{"_1":1},'
+
+    state_manager.unset_transition(model_name="s3::keyrestore")
+
+
 @mock_aws
 def test_restore_unknown_key():
     client = boto3.client("s3", region_name=DEFAULT_REGION_NAME)
@@ -633,7 +682,8 @@ def test_restore_object_invalid_request_params():
         raise SkipTest("Can't set transition directly in ServerMode")
 
     s3_resource = boto3.resource("s3", region_name=DEFAULT_REGION_NAME)
-    bucket = s3_resource.Bucket("foobar")
+    bucket_name = "foobar"
+    bucket = s3_resource.Bucket(bucket_name)
     bucket.create()
 
     key = bucket.put_object(Key="the-key", Body=b"somedata", StorageClass="GLACIER")
@@ -652,6 +702,85 @@ def test_restore_object_invalid_request_params():
     assert err["Code"] == "DaysMustNotProvidedForSelectRequest"
     assert err["Message"] == "`Days` must not be provided for select requests"
 
+    # Invalid type for restore request. It must be `SELECT`.
+    with pytest.raises(ClientError) as exc:
+        key.restore_object(RestoreRequest={"Type": "INVALID_TYPE"})
+    err = exc.value.response["Error"]
+    assert err["Code"] == "InvalidRestoreRequestType"
+    assert err["Message"] == "Invalid type for restore request. It must be `SELECT`"
+
+    # `SelectParameters` must be provided for select requests.
+    with pytest.raises(ClientError) as exc:
+        key.restore_object(RestoreRequest={"Type": "SELECT"})
+    err = exc.value.response["Error"]
+    assert err["Code"] == "MissingRequiredParametersForSelectRequest"
+    assert (
+        err["Message"]
+        == 'Missing required parameter in RestoreRequest: "SelectParameters" for select request'
+    )
+
+    # `OutputLocation` must be provided for select requests.
+    with pytest.raises(ClientError) as exc:
+        key.restore_object(
+            RestoreRequest={
+                "Type": "SELECT",
+                "SelectParameters": {
+                    "Expression": "SELECT * FROM S3Object",
+                    "ExpressionType": "SQL",
+                    "InputSerialization": {"JSON": {"Type": "DOCUMENT"}},
+                    "OutputSerialization": {"JSON": {"RecordDelimiter": ","}},
+                },
+            },
+        )
+    err = exc.value.response["Error"]
+    assert err["Code"] == "MissingRequiredParametersForSelectRequest"
+    assert (
+        err["Message"]
+        == 'Missing required parameter in RestoreRequest: "OutputLocation" for select request'
+    )
+
+    # raise error when supecified bucket does not exists.
+    with pytest.raises(ClientError) as exc:
+        key.restore_object(
+            RestoreRequest={
+                "Type": "SELECT",
+                "SelectParameters": {
+                    "Expression": "SELECT COUNT(*) as cnt FROM S3Object",
+                    "ExpressionType": "SQL",
+                    "InputSerialization": {"JSON": {"Type": "DOCUMENT"}},
+                    "OutputSerialization": {"JSON": {"RecordDelimiter": ","}},
+                },
+                "OutputLocation": {
+                    "S3": {"BucketName": "other-bucket", "Prefix": "test"}
+                },
+            }
+        )
+    err = exc.value.response["Error"]
+    err["Code"] == "NoSuchBucket"
+
+    # raise error when specified storage class is invalid
+    with pytest.raises(ClientError) as exc:
+        key.restore_object(
+            RestoreRequest={
+                "Type": "SELECT",
+                "SelectParameters": {
+                    "Expression": "SELECT COUNT(*) as cnt FROM S3Object",
+                    "ExpressionType": "SQL",
+                    "InputSerialization": {"JSON": {"Type": "DOCUMENT"}},
+                    "OutputSerialization": {"JSON": {"RecordDelimiter": ","}},
+                },
+                "OutputLocation": {
+                    "S3": {
+                        "BucketName": bucket_name,
+                        "Prefix": "test.json",
+                        "StorageClass": "INVALID_CLASS",
+                    }
+                },
+            }
+        )
+    err = exc.value.response["Error"]
+    err["Code"] == "InvalidStorageClass"
+
 
 @mock_aws
 def test_get_versioning_status():