[PROD-1103] Access report (#31)

Author: rmoneys

sync/__init__.py

@@ -1,4 +1,4 @@
 """Library for leveraging the power of Sync"""
-__version__ = "0.1.0"
+__version__ = "0.1.1"
 
 TIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ"

sync/api/__init__.py

@@ -0,0 +1,33 @@
+from sync.clients.sync import get_default_client
+from sync.models import AccessReport, AccessReportLine, AccessStatusCode
+
+
+def get_access_report() -> AccessReport:
+    """Reports access to the Sync API
+
+    :return: access report
+    :rtype: AccessReport
+    """
+    response = get_default_client().get_products()
+
+    error = response.get("error")
+    if error:
+        return AccessReport(
+            [
+                AccessReportLine(
+                    name="Sync Authentication",
+                    status=AccessStatusCode.RED,
+                    message=f"{error['code']}: {error['message']}",
+                )
+            ]
+        )
+
+    return AccessReport(
+        [
+            AccessReportLine(
+                name="Sync Authentication",
+                status=AccessStatusCode.GREEN,
+                message="Sync credentials are valid",
+            )
+        ]
+    )

sync/awsdatabricks.py

@@ -12,14 +12,18 @@
 from urllib.parse import urlparse
 
 import boto3 as boto
+import orjson
 from botocore.exceptions import ClientError
-from orjson import orjson
 
+from sync.api import get_access_report as get_api_access_report
 from sync.api.predictions import create_prediction_with_eventlog_bytes, get_prediction
 from sync.api.projects import get_project
 from sync.clients.databricks import get_default_client
 from sync.config import CONFIG, DB_CONFIG
 from sync.models import (
+    AccessReport,
+    AccessReportLine,
+    AccessStatusCode,
     DatabricksAPIError,
     DatabricksClusterReport,
     DatabricksError,
@@ -31,6 +35,100 @@
 logger = logging.getLogger(__name__)
 
 
+def get_access_report(log_url: str = None) -> AccessReport:
+    """Reports access to Databricks, AWS and Sync required for integrating jobs with Sync.
+    Access is partially determined by the configuration of this library and boto3.
+
+    :param log_url: location of event logs, defaults to None
+    :type log_url: str, optional
+    :return: access report
+    :rtype: AccessReport
+    """
+    report = get_api_access_report()
+    dbx_client = get_default_client()
+
+    response = dbx_client.get_current_user()
+    user_name = response.get("userName")
+    if user_name:
+        report.append(
+            AccessReportLine(
+                name="Databricks Authentication",
+                status=AccessStatusCode.GREEN,
+                message=f"Authenticated as '{user_name}'",
+            )
+        )
+    else:
+        report.append(
+            AccessReportLine(
+                name="Databricks Authentication",
+                status=AccessStatusCode.RED,
+                message=f"{response.get('error_code')}: {response.get('message')}",
+            )
+        )
+
+    response = boto.client("sts").get_caller_identity()
+    arn = response.get("Arn")
+    if arn:
+        report.append(
+            AccessReportLine(
+                name="AWS Authentication",
+                status=AccessStatusCode.GREEN,
+                message=f"Authenticated as '{arn}'",
+            )
+        )
+
+        ec2 = boto.client("ec2", region_name=DB_CONFIG.aws_region_name)
+        report.add_boto_method_call(ec2.describe_instances, AccessStatusCode.YELLOW, DryRun=True)
+    else:
+        report.append(
+            AccessReportLine(
+                name="AWS Authentication",
+                status=AccessStatusCode.RED,
+                message="Failed to authenticate AWS credentials",
+            )
+        )
+
+    if log_url:
+        parsed_log_url = urlparse(log_url)
+
+        if parsed_log_url.scheme == "s3" and arn:
+            s3 = boto.client("s3")
+            report.add_boto_method_call(
+                s3.list_objects_v2,
+                Bucket=parsed_log_url.netloc,
+                Prefix=parsed_log_url.params.rstrip("/"),
+                MaxKeys=1,
+            )
+        elif parsed_log_url.scheme == "dbfs":
+            response = dbx_client.list_dbfs_directory(parsed_log_url.geturl())
+            if "error_code" not in response:
+                report.append(
+                    AccessReportLine(
+                        name="Log Access",
+                        status=AccessStatusCode.GREEN,
+                        message=f"Can list objects at {parsed_log_url.geturl()}",
+                    )
+                )
+            else:
+                report.append(
+                    AccessReportLine(
+                        name="Log Access",
+                        status=AccessStatusCode.RED,
+                        message=f"Can list objects at {parsed_log_url.geturl()}",
+                    )
+                )
+        else:
+            report.append(
+                AccessReportLine(
+                    name="Log Access",
+                    status=AccessStatusCode.RED,
+                    message=f"scheme in {parsed_log_url.geturl()} is not supported",
+                )
+            )
+
+    return report
+
+
 def create_prediction(
     plan_type: str,
     compute_type: str,
@@ -314,15 +412,18 @@ def _get_cluster_instances(cluster: dict) -> Response[dict]:
     #  are associated with this cluster
     if cluster_instances is None:
         ec2 = boto.client("ec2", region_name=aws_region_name)
-        cluster_instances = ec2.describe_instances(
-            Filters=[
-                {"Name": "tag:Vendor", "Values": ["Databricks"]},
-                {"Name": "tag:ClusterId", "Values": [cluster_id]},
-                # {'Name': 'tag:JobId', 'Values': []}
-            ]
-        )
+        try:
+            cluster_instances = ec2.describe_instances(
+                Filters=[
+                    {"Name": "tag:Vendor", "Values": ["Databricks"]},
+                    {"Name": "tag:ClusterId", "Values": [cluster_id]},
+                    # {'Name': 'tag:JobId', 'Values': []}
+                ]
+            )
+        except Exception as exc:
+            logger.warning(exc)
 
-    if not cluster_instances["Reservations"]:
+    if not cluster_instances or not cluster_instances["Reservations"]:
         no_instances_message = (
             f"Unable to find any active or recently terminated instances for cluster `{cluster_id}` in `{aws_region_name}`. "
             + "Please refer to the following documentation for options on how to address this - "
@@ -1179,10 +1280,10 @@ def _get_eventlog_from_dbfs(
     run_end_time_millis: int,
     poll_duration_seconds: int,
 ):
-    prefix = format_dbfs_filepath(f"{base_filepath}/eventlog/")
-
-    root_dir = get_default_client().list_dbfs_directory(prefix)
+    dbx_client = get_default_client()
 
+    prefix = format_dbfs_filepath(f"{base_filepath}/eventlog/")
+    root_dir = dbx_client.list_dbfs_directory(prefix)
     eventlog_files = [f for f in root_dir["files"] if f["is_dir"]]
     matching_subdirectory = None
 
@@ -1195,7 +1296,7 @@ def _get_eventlog_from_dbfs(
         eventlog_file_metadata = eventlog_files.pop()
         path = eventlog_file_metadata["path"]
 
-        subdir = get_default_client().list_dbfs_directory(path)
+        subdir = dbx_client.list_dbfs_directory(path)
 
         subdir_files = subdir["files"]
         matching_subdirectory = next(
@@ -1204,7 +1305,7 @@ def _get_eventlog_from_dbfs(
         )
 
     if matching_subdirectory:
-        eventlog_dir = get_default_client().list_dbfs_directory(matching_subdirectory["path"])
+        eventlog_dir = dbx_client.list_dbfs_directory(matching_subdirectory["path"])
 
         poll_num_attempts = 0
         poll_max_attempts = 20  # 5 minutes / 15 seconds = 20 attempts
@@ -1219,14 +1320,13 @@ def _get_eventlog_from_dbfs(
                 )
                 sleep(poll_duration_seconds)
 
-            eventlog_dir = get_default_client().list_dbfs_directory(matching_subdirectory["path"])
+            eventlog_dir = dbx_client.list_dbfs_directory(matching_subdirectory["path"])
             poll_num_attempts += 1
 
         eventlog_zip = io.BytesIO()
         eventlog_zip_file = zipfile.ZipFile(eventlog_zip, "a", zipfile.ZIP_DEFLATED)
 
         eventlog_files = eventlog_dir["files"]
-        dbx_client = get_default_client()
         for eventlog_file_metadata in eventlog_files:
             filename: str = eventlog_file_metadata["path"].split("/")[-1]
             filesize: int = eventlog_file_metadata["file_size"]

sync/awsemr.py

@@ -7,25 +7,119 @@
 import logging
 import re
 from copy import deepcopy
+from typing import Tuple
 from urllib.parse import urlparse
 from uuid import uuid4
 
 import boto3 as boto
 import orjson
 from dateutil.parser import parse as dateparse
-from typing import Tuple
 
 from sync import TIME_FORMAT
+from sync.api import get_access_report as get_api_access_report
 from sync.api.predictions import create_prediction, wait_for_prediction
 from sync.api.projects import get_project
-from sync.models import EMRError, Platform, ProjectError, Response
+from sync.models import (
+    AccessReport,
+    AccessReportLine,
+    AccessStatusCode,
+    EMRError,
+    Platform,
+    ProjectError,
+    Response,
+)
 
 logger = logging.getLogger(__name__)
 
 
 RUN_DIR_PATTERN_TEMPLATE = r"{project_prefix}/{project_id}/(?P<timestamp>\d{{4}}-[^/]+)/{run_id}"
 
 
+def get_access_report(
+    log_url: str = None, cluster_id: str = None, region_name: str = None
+) -> AccessReport:
+    """Reports access to systems required for integration of EMR jobs with Sync.
+    Access is partially determined by the configuration of this library and boto3.
+
+    :param log_url: location of event logs, defaults to None
+    :type log_url: str, optional
+    :param cluster_id: cluster ID with which to test EMR access, defaults to None
+    :type cluster_id: str, optional
+    :param region_name: region override, defaults to None
+    :type region_name: str, optional
+    :return: access report
+    :rtype: AccessReport
+    """
+    report = get_api_access_report()
+    sts = boto.client("sts")
+    response = sts.get_caller_identity()
+
+    arn = response.get("Arn")
+    if not arn:
+        report.append(
+            AccessReportLine(
+                name="AWS Authentication",
+                status=AccessStatusCode.RED,
+                message="Failed to authenticate AWS credentials",
+            )
+        )
+    else:
+        report.append(
+            AccessReportLine(
+                name="AWS Authentication",
+                status=AccessStatusCode.GREEN,
+                message=f"Authenticated as '{arn}'",
+            )
+        )
+
+    if log_url:
+        parsed_log_url = urlparse(log_url)
+
+        if parsed_log_url.scheme == "s3" and arn:
+            s3 = boto.client("s3")
+            report.add_boto_method_call(
+                s3.list_objects_v2,
+                Bucket=parsed_log_url.netloc,
+                Prefix=parsed_log_url.params.rstrip("/"),
+                MaxKeys=1,
+            )
+        else:
+            report.append(
+                AccessReportLine(
+                    name="Logging",
+                    status=AccessStatusCode.RED,
+                    message=f"scheme in {parsed_log_url.geturl()} is not supported",
+                )
+            )
+
+    if arn and cluster_id:
+        emr = boto.client("emr", region_name=region_name)
+
+        try:
+            response = emr.describe_cluster(ClusterId=cluster_id)
+            report.append(
+                AccessReportLine(
+                    "EMR DescribeCluster",
+                    AccessStatusCode.GREEN,
+                    "describe_cluster call succeeded",
+                )
+            )
+
+            if response["Cluster"]["InstanceCollectionType"] == "INSTANCE_FLEET":
+                report.add_boto_method_call(emr.list_instance_fleets, ClusterId=cluster_id)
+            elif response["Cluster"]["InstanceCollectionType"] == "INSTANCE_GROUP":
+                report.add_boto_method_call(emr.list_instance_groups, ClusterId=cluster_id)
+
+            report.add_boto_method_call(emr.list_bootstrap_actions, ClusterId=cluster_id)
+            report.add_boto_method_call(emr.list_instances, ClusterId=cluster_id)
+            report.add_boto_method_call(emr.list_steps, ClusterId=cluster_id)
+
+        except Exception as exc:
+            report.append(AccessReportLine("EMR DescribeCluster", AccessStatusCode.RED, str(exc)))
+
+    return report
+
+
 def get_project_job_flow(job_flow: dict, project_id: str) -> Response[dict]:
     """Returns a copy of the incoming job flow with project configuration.
 

sync/cli/awsdatabricks.py

@@ -1,5 +1,5 @@
 import click
-from orjson import orjson
+import orjson
 
 from sync import awsdatabricks
 from sync.cli.util import validate_project
@@ -10,7 +10,13 @@
 @click.group
 def aws_databricks():
     """Databricks on AWS commands"""
-    pass
+
+
+@aws_databricks.command
+@click.option("--log-url")
+def access_report(log_url: str = None):
+    """Get access report"""
+    click.echo(awsdatabricks.get_access_report(log_url))
 
 
 @aws_databricks.command

sync/cli/awsemr.py

@@ -17,6 +17,15 @@ def aws_emr():
     pass
 
 
+@aws_emr.command
+@click.option("--cluster-id")
+@click.option("--log-url")
+@click.option("-r", "--region")
+def access_report(cluster_id: str = None, log_url: str = None, region: str = None):
+    """Get access report"""
+    click.echo(awsemr.get_access_report(log_url=log_url, cluster_id=cluster_id, region_name=region))
+
+
 @aws_emr.command
 @click.argument("job-flow", type=click.File("r"))
 @click.option("-p", "--project", callback=validate_project)