--- /dev/null
+#!/usr/bin/env/env python3
+
+# Copyright (c) 2021 Cisco and/or its affiliates.
+# 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.
+
+"""Storage Backend Class."""
+
+from json import loads
+from struct import unpack
+from gzip import GzipFile
+
+from boto3 import Session
+from botocore import exceptions
+
+S3_API_LIMIT = 1048576
+
+
+class Storage:
+ """Class implementing storage object retrieval.
+ S3 Select API allows us to retrieve a subset of data by using simple SQL
+ expressions. By using Select API to retrieve only the data needed by the
+ application, drastic performance improvements can be achieved.
+ """
+ def __init__(self, endpoint_url, bucket, profile_name):
+ """Class init function to create S3 client object.
+
+ :param endpoint_url: S3 storage endpoint url.
+ :param bucket: S3 parent bucket.
+ :param profile_name: S3 storage configuration.
+ :type endpoint_url: str
+ :type bucket: str
+ :type profile_name: str
+ """
+ self.endpoint_url = endpoint_url
+ self.bucket = bucket
+ self.profile_name = profile_name
+
+ self.session = Session(profile_name=self.profile_name)
+ self.client = self.session.client(
+ service_name=u"s3", endpoint_url=self.endpoint_url
+ )
+ self.resource = self.session.resource(
+ service_name=u"s3", endpoint_url=self.endpoint_url
+ )
+
+ def __repr__(self):
+ """Return a string executable as Python constructor call.
+
+ :returns: Executable constructor call.
+ :rtype: str
+ """
+ return (
+ f"Storage(endpoint_url={self.endpoint_url!r}, "
+ f"bucket={self.bucket!r}, "
+ f"profile_name={self.profile_name!r})"
+ )
+
+ def _get_matching_s3_keys(
+ self, bucket, prefix=u"", suffix=u""):
+ """This function generates the keys in an S3 bucket. Function act as
+ a Python generator object.
+
+ :param bucket: Name of the S3 bucket.
+ :param prefix: Only fetch keys that start with this prefix (optional).
+ :param suffix: Only fetch keys that end with this suffix (optional).
+ :type bucket: str
+ :type prefix: str
+ :type suffix: str
+ :raises RuntimeError: If connection to storage fails.
+ """
+ kwargs = {
+ u"Bucket": bucket
+ }
+
+ prefixes = (prefix, ) if isinstance(prefix, str) else prefix
+
+ for key_prefix in prefixes:
+ kwargs[u"Prefix"] = key_prefix
+ try:
+ paginator = self.client.get_paginator(u"list_objects_v2")
+ for page in paginator.paginate(**kwargs):
+ try:
+ contents = page[u"Contents"]
+ except KeyError:
+ break
+
+ for obj in contents:
+ key = obj[u"Key"]
+ if key.endswith(suffix):
+ yield obj
+ except exceptions.EndpointConnectionError:
+ raise RuntimeError(
+ u"Connection Error!"
+ )
+
+ def _get_matching_s3_content(
+ self, key, expression):
+ """This function filters the contents of an S3 object based on a simple
+ structured query language (SQL) statement. In the request, along with
+ the SQL expression, we are specifying JSON serialization of the object.
+ S3 uses this format to parse object data into records, and returns only
+ records that match the specified SQL expression. Data serialization
+ format for the response is set to JSON.
+
+ :param key: S3 Key (file path).
+ :param expression: S3 compatible SQL query.
+ :type key: str
+ :type expression: str
+ :returns: JSON content of interest.
+ :rtype: str
+ :raises RuntimeError: If connection to storage fails.
+ :raises ValueError: If JSON reading fails.
+ """
+ try:
+ content = self.client.select_object_content(
+ Bucket=self.bucket,
+ Key=key,
+ ExpressionType=u"SQL",
+ Expression=expression,
+ InputSerialization={
+ u"JSON": {
+ u"Type": u"Document"
+ },
+ u"CompressionType": u"GZIP"
+ },
+ OutputSerialization={
+ u"JSON": {
+ u"RecordDelimiter": u""
+ }
+ }
+ )
+ records = u""
+ for event in content[u"Payload"]:
+ if u"Records" in event:
+ records = event[u"Records"][u"Payload"].decode(u"utf-8")
+ return records
+ except exceptions.EndpointConnectionError:
+ raise RuntimeError(
+ u"Connection Error!"
+ )
+ except exceptions.EventStreamError:
+ raise ValueError(
+ u"Malformed JSON content!"
+ )
+
+ def _get_matching_s3_object(
+ self, key):
+ """Gets full S3 object. If the file is gzip'd it will be unpacked.
+
+ :param key: Name of the S3 key (file).
+ :type key: str
+ :returns: JSON file of interest.
+ :rtype: str
+ :raises RuntimeError: If connection to storage fails.
+ """
+ try:
+ streaming_object = self.client.get_object(
+ Bucket=self.bucket,
+ Key=key
+ )[u"Body"]
+ with GzipFile(fileobj=streaming_object) as gzipfile:
+ content = gzipfile.read()
+ return content
+ except exceptions.EndpointConnectionError:
+ raise RuntimeError(
+ u"Connection Error!"
+ )
+
+ def _get_matching_s3_length(
+ self, key):
+ """Gets the file size of S3 object. If the file is gzip'd the packed
+ size is reported.
+
+ :param key: Name of the S3 key (file).
+ :type key: str
+ :returns: File size in bytes. Defaults to 0 if any error.
+ :rtype: int
+ :raises RuntimeError: If connection to storage fails.
+ """
+ try:
+ compressed_size = self.client.get_object(
+ Bucket=self.bucket,
+ Key=key
+ )[u"ContentLength"]
+ last_four_bytes = self.client.get_object(
+ Bucket=self.bucket,
+ Key=key,
+ Range=f"bytes={compressed_size-4}-{compressed_size}"
+ )[u"Body"]
+ return unpack(u"I", last_four_bytes.read(4))[0]
+ except exceptions.EndpointConnectionError:
+ return 0
+
+ def is_large_file(
+ self, key):
+ """Returns True if file is larger then 1MB that S3 select allows.
+
+ :param key: Name of the S3 key (file).
+ :type key: str
+ :returns: Returns True if file is large then 1MB that S3 select allows.
+ :rtype: bool
+ """
+ return bool(
+ self._get_matching_s3_length(key=key[u"Key"]) > S3_API_LIMIT
+ )
+
+ def s3_file_processing(
+ self, prefix=u"", suffix=u"json.gz",
+ expression=u"select * from s3object s"):
+ """Batch S3 key processing. Function retrieves list of files and use
+ S3 Select API to query content.
+
+ :param prefix: Only fetch keys that start with this prefix (optional).
+ :param suffix: Only fetch keys that end with this suffix (optional).
+ :param expression: S3 compatible SQL query (optional).
+ :type prefix: str
+ :type suffix: str
+ :type expression: str
+ """
+ key_iterator = self._get_matching_s3_keys(
+ bucket=self.bucket,
+ prefix=prefix,
+ suffix=suffix
+ )
+
+ for key in key_iterator:
+ try:
+ yield key[u"Key"], loads(
+ self._get_matching_s3_content(
+ key=key[u"Key"], expression=expression
+ )
+ )
+ except ValueError:
+ return
+
+ def s3_dump_file_processing(
+ self, prefix=u"", suffix=u"json.gz"):
+ """Batch S3 key processing. Function retrieves list of files and use
+ S3 Get Object API to query content.
+
+ :param prefix: Only fetch keys that start with this prefix (optional).
+ :param suffix: Only fetch keys that end with this suffix (optional).
+ :type prefix: str
+ :type suffix: str
+ """
+ key_iterator = self._get_matching_s3_keys(
+ bucket=self.bucket,
+ prefix=prefix,
+ suffix=suffix
+ )
+
+ for key in key_iterator:
+ try:
+ yield loads(
+ self._get_matching_s3_object(
+ key=key[u"Key"]
+ )
+ )
+ except ValueError:
+ return
Code Review / csit.git / blobdiff