progress commit

Author: ChintanRaval

rdl/data_sources/AWSLambdaDataSource.py

@@ -0,0 +1,101 @@
+import logging
+import pandas
+import json
+import boto3
+
+from rdl.data_sources.ChangeTrackingInfo import ChangeTrackingInfo
+from rdl.data_sources.SourceTableInfo import SourceTableInfo
+from rdl.shared import Providers
+from rdl.shared.Utils import prevent_senstive_data_logging
+
+
+class AWSLambdaDataSource(object):
+    CONNECTION_STRING_PREFIX = 'aws-lambda-arn://'  # 'aws-lambda://tenant=543_dc2;arn=123456789012:function:my-function;'
+
+    def __init__(self, connection_string, logger=None):
+        self.logger = logger or logging.getLogger(__name__)
+        self.connection_string = connection_string
+        self.aws_lambda_client = boto3.client('lambda')
+
+    @staticmethod
+    def can_handle_connection_string(connection_string):
+        return connection_string.startswith(AWSLambdaDataSource.CONNECTION_STRING_PREFIX)
+
+    @staticmethod
+    def connection_string_prefix():
+        return AWSLambdaDataSource.CONNECTION_STRING_PREFIX
+
+    def get_table_info(self, table_config, last_known_sync_version):
+        column_names, last_sync_version, sync_version, full_refresh_required, data_changed_since_last_sync = \
+            self.__get_table_info(table_config, last_known_sync_version)
+        columns_in_database = column_names
+        change_tracking_info = ChangeTrackingInfo(
+            last_sync_version=last_sync_version,
+            sync_version=sync_version + 1,
+            force_full_load=full_refresh_required,
+            data_changed_since_last_sync=data_changed_since_last_sync)
+        source_table_info = SourceTableInfo(columns_in_database, change_tracking_info)
+        return source_table_info
+
+    @prevent_senstive_data_logging
+    def get_table_data_frame(self, table_config, columns, batch_config, batch_tracker, batch_key_tracker,
+                             full_refresh, change_tracking_info):
+        self.logger.debug(f"Starting read data from lambda.. : \n{None}")
+        column_names, data = self.__get_table_data(table_config, batch_config, change_tracking_info, full_refresh, columns, batch_key_tracker)
+        self.logger.debug(f"Finished read data from lambda.. : \n{None}")
+        # should we log size of data extracted?
+        data_frame = pandas.DataFrame(data=data, columns=column_names)
+        batch_tracker.extract_completed_successfully(len(data_frame))
+        return data_frame
+
+    def __get_table_info(self, table_config, last_known_sync_version):
+        pay_load = {
+            "command": "GetTableInfo",
+            "tenantId": 543,  # self.connection_string.tenant.split('_')[0] as int
+            "table": {
+                "schema": table_config['schema'],
+                "name": table_config['name']
+            },
+            "commandPayload": {
+                "lastSyncVersion": last_known_sync_version,
+            }
+        }
+
+        result = self.__invoke_lambda(pay_load)
+
+        return result['ColumnNames'], result['Data']
+
+    def __get_table_data(self, table_config, batch_config, change_tracking_info, full_refresh, columns, batch_key_tracker):
+        pay_load = {
+            "command": "GetTableData",
+            "tenantId": 543,  # self.connection_string.tenant.split('_')[0] as int
+            "table": {
+                "schema": table_config['schema'],
+                "name": table_config['name']
+            },
+            "commandPayload": {
+                "auditColumnNameForChangeVersion": Providers.AuditColumnsNames.CHANGE_VERSION,
+                "auditColumnNameForDeletionFlag": Providers.AuditColumnsNames.IS_DELETED,
+                "batchSize": batch_config['size'],
+                "lastSyncVersion": change_tracking_info.last_sync_version,
+                "fullRefresh": full_refresh,
+                "columnNames": columns,
+                "primaryKeyColumnNames": table_config['primary_keys'],
+                "lastBatchPrimaryKeys": [{k: v} for k, v in batch_key_tracker.bookmarks.items()]
+            }
+        }
+
+        result = self.__invoke_lambda(pay_load)
+
+        return result['ColumnNames'], result['Data']
+
+    def __invoke_lambda(self, pay_load):
+        lambda_response = self.aws_lambda_client.invoke(
+            FunctionName='string',  # self.connection_string.arn
+            InvocationType='RequestResponse',
+            LogType='None',  # |'Tail', Set to Tail to include the execution log in the response
+            Payload=json.dump(pay_load).encode()
+        )
+        result = json.loads(lambda_response.Payload.read())  # .decode()
+        return result
+

rdl/data_sources/DataSourceFactory.py

@@ -1,12 +1,12 @@
 import logging
 from rdl.data_sources.MsSqlDataSource import MsSqlDataSource
-
+from rdl.data_sources.AWSLambdaDataSource import AWSLambdaDataSource
 
 class DataSourceFactory(object):
 
     def __init__(self, logger=None):
         self.logger = logger or logging.getLogger(__name__)
-        self.sources = [MsSqlDataSource]
+        self.sources = [MsSqlDataSource, AWSLambdaDataSource]
 
     def create_source(self, connection_string):
         for source in self.sources:

setup.py

@@ -14,6 +14,7 @@
           'SQLAlchemy==1.3.3',
           'sqlalchemy-citext==1.3.post0',
           'alembic==1.0.9',
+          'boto3==1.9.187',
       ],
       package_data={
           '': ['alembic.ini', 'alembic/*.py', 'alembic/**/*.py'],