fix: Log management (#237)

* some basic log management

* code automatically formatted

Signed-off-by: sentential[bot] <bot@sentential>

---------

Signed-off-by: sentential[bot] <bot@sentential>
Co-authored-by: sentential[bot] <bot@sentential>

sentential/cli/root.py

@@ -4,6 +4,7 @@
 from sentential.lib.drivers.local_images import LocalImagesDriver
 from sentential.lib.drivers.local_lambda import LocalLambdaDriver
 from sentential.lib.drivers.aws_ecr import AwsEcrDriver
+from sentential.lib.drivers.aws_lambda import AwsLambdaDriver
 from sentential.lib.template import Init
 from sentential.lib.shapes import Architecture, Runtimes
 from sentential.lib.ontology import Ontology
@@ -59,9 +60,13 @@ def ls(verbose: bool = typer.Option(False)):
 
 
 @root.command()
-def clean(remote: bool = typer.Option(False)):
-    """clean images"""
+def clean(remote: bool = typer.Option(False), remote_logs: bool = typer.Option(False)):
+    """clean images (and logs)"""
     ontology = Ontology()
     LocalImagesDriver(ontology).clean()
     if remote:
         AwsEcrDriver(ontology).clean()
+    if (
+        remote_logs
+    ):  # MAYBE: is it time to graduate to `clean local` and `clean remote`?
+        AwsLambdaDriver(ontology).clean()

sentential/lib/clients.py

@@ -17,7 +17,7 @@ def __init__(self, env: str = getenv("SENTENTIAL_ENV", "none")) -> None:
         self.kms = boto3.client("kms", **boto3_config[env])
         self.ebr = boto3.client("events", **boto3_config[env])
         self.api_gw = boto3.client("apigatewayv2", **boto3_config[env])
-        self.cloudwatch = boto3.client("logs", **boto3_config[env])
+        self.logs = boto3.client("logs", **boto3_config[env])
         self.docker = docker
 
 

sentential/lib/drivers/aws_lambda.py

@@ -22,6 +22,7 @@ def __init__(self, ontology: Ontology) -> None:
         self.ontology = ontology
         self.function_name = ontology.context.resource_name
         self.policy_arn = f"arn:aws:iam::{self.ontology.context.account_id}:policy/{self.ontology.context.resource_name}"
+        self.log_group = f"/aws/lambda/{self.function_name}"
 
     @property
     def provision(self) -> Provision:
@@ -41,6 +42,7 @@ def deploy(self, image: AwsImageDetail, arch: Union[Architecture, None]) -> str:
             PolicyArn=self._put_policy(tags)["Policy"]["Arn"],
         )
         self._put_lambda(chosen_dist, tags)
+        self._put_log_policy()
 
         return f"deployed {self.ontology.context.resource_name} to aws"
 
@@ -122,6 +124,9 @@ def invoke(self, payload: str) -> str:
         response["Payload"] = response["Payload"].decode("utf-8")
         return LambdaInvokeResponse(**response).json()
 
+    def clean(self) -> None:
+        self._clean_logs()
+
     def _put_role(self, tags: Optional[Dict[str, str]] = None) -> Dict:
         role_name = self.function_name
         try:
@@ -249,3 +254,20 @@ def _put_lambda(
             clients.lmb.tag_resource(Resource=function["FunctionArn"], Tags=tags)
 
         return function
+
+    def _put_log_policy(self) -> None:
+        try:
+            clients.logs.put_retention_policy(
+                logGroupName=self.log_group,
+                retentionInDays=self.provision.log_retention,
+            )
+        except clients.logs.exceptions.ResourceNotFoundException:
+            pass
+
+    def _clean_logs(self) -> None:
+        try:
+            clients.logs.delete_log_group(
+                logGroupName=self.log_group,
+            )
+        except clients.logs.exceptions.ResourceNotFoundException:
+            pass

sentential/lib/shapes.py

@@ -30,6 +30,7 @@ class Provision(Shaper):
     allow_methods: List[str] = Field(default=["*"], description="CORS AllowMethods")
     allow_origins: List[str] = Field(default=["*"], description="CORS AllowOrigins")
     expose_headers: List[str] = Field(default=["*"], description="CORS ExposeHeaders")
+    log_retention: int = Field(default=7, description="log retention (days)")
 
     @validator("auth_type")
     def is_valid_auth_type(cls, v):

tests/lib/drivers/test_aws_lambda.py

@@ -4,19 +4,29 @@
 from sentential.lib.clients import clients
 from sentential.lib.drivers.aws_lambda import AwsLambdaDriver
 from sentential.lib.drivers.aws_ecr import AwsEcrDriver
+from sentential.lib.exceptions import AwsDriverError
 from sentential.lib.exceptions import AWS_EXCEPTIONS
 
 
 @pytest.mark.usefixtures(
     "moto", "init", "ontology", "mock_repo", "aws_lambda_driver", "aws_ecr_driver"
 )
 class TestAwsLambdaDriver:
+    def create_log_group(self, log_group: str):
+        clients.logs.create_log_group(logGroupName=log_group)
+        clients.logs.put_retention_policy(logGroupName=log_group, retentionInDays=7)
+
     def get_lambda(self, function_name: str):
         return clients.lmb.get_function(FunctionName=function_name)
 
     def get_lambda_config(self, function_name: str):
         return clients.lmb.get_function_configuration(FunctionName=function_name)
 
+    def get_log_policy(self, log_group: str):
+        return clients.logs.describe_log_groups(logGroupNamePattern=log_group)[
+            "logGroups"
+        ][0]["retentionInDays"]
+
     def test_deploy(
         self, aws_lambda_driver: AwsLambdaDriver, aws_ecr_driver: AwsEcrDriver
     ):
@@ -51,8 +61,12 @@ def test_deploy_w_provisions(
 
         image = aws_ecr_driver.get_image("0.0.1")
         aws_lambda_driver.deploy(image, Architecture.system())
+        self.create_log_group(
+            aws_lambda_driver.log_group
+        )  # Is this something that moto should be doing?
 
         lambda_config = self.get_lambda_config(function_name)
+        log_policy = self.get_log_policy(aws_lambda_driver.log_group)
         ssm_config = cast(Provision, configs.parameters)
 
         assert lambda_config["MemorySize"] == ssm_config.memory
@@ -62,3 +76,9 @@ def test_deploy_w_provisions(
             lambda_config["VpcConfig"]["SecurityGroupIds"]
             == ssm_config.security_group_ids
         )
+        assert log_policy == ssm_config.log_retention
+
+    def test_clean(self, aws_lambda_driver: AwsLambdaDriver):
+        aws_lambda_driver.clean()
+        with pytest.raises(IndexError):
+            self.get_log_policy(aws_lambda_driver.log_group)