Inital stab at refactoring for sources that can change.

Author: dames

relational_data_loader_project/BatchDataLoader.py

@@ -1,46 +1,28 @@
 import logging
-import pandas
 from io import StringIO
 import importlib
 
 from column_transformers.StringTransformers import ToUpper
 
 
-
 class BatchDataLoader(object):
-    def __init__(self, source_table_configuration, columns, batch_configuration, logger=None):
+    def __init__(self, data_source, source_table_configuration, target_table_configuration, columns, data_load_tracker, batch_configuration, target_engine, logger=None):
         self.logger = logger or logging.getLogger(__name__)
         self.source_table_configuration = source_table_configuration
         self.columns = columns
+        self.data_source = data_source
+        self.target_table_configuration = target_table_configuration
+        self.data_load_tracker = data_load_tracker
         self.batch_configuration = batch_configuration
-
-    def build_select_statement(self, previous_key=0):
-
-        column_array = list(map(lambda cfg: cfg['source_name'], self.columns))
-        column_names = ", ".join(column_array)
-
-
-        return "SELECT TOP ({0}) {1} FROM {2}.{3} WHERE {4} > {5} ORDER BY {4}".format(self.batch_configuration['size'],
-                                                            column_names,
-                                                            self.source_table_configuration['schema'],
-                                                            self.source_table_configuration['name'],
-                                                            self.batch_configuration['source_unique_column'],
-                                                            previous_key
-                                                            )
+        self.target_engine = target_engine
 
     # Imports rows, returns True if >0 rows were found
-    def import_batch(self, source_engine, target_engine, target_table_configuration, batch_tracker, previous_key):
-        self.logger.debug("ImportBatch Starting for source {0} target {1} previous_key {2}".format(self.source_table_configuration['name'],
-                                                                                                   target_table_configuration['name'],
-                                                                                                   previous_key))
-
-        sql = self.build_select_statement(previous_key)
+    def import_batch(self, previous_batch_key):
+        batch_tracker = self.data_load_tracker.start_batch()
 
-        self.logger.debug("Starting read of SQL Statement: {0}".format(sql))
-        data_frame = pandas.read_sql_query(sql, source_engine)
-        self.logger.debug("Completed read")
+        self.logger.debug("ImportBatch Starting from previous_batch_key: {0}".format(previous_batch_key))
 
-        batch_tracker.extract_completed_successfully(len(data_frame))
+        data_frame = self.data_source.get_next_data_frame(self.source_table_configuration, self.columns, self.batch_configuration, batch_tracker, previous_batch_key)
 
         if len(data_frame) == 0:
             self.logger.debug("There are no rows to import, returning -1")
@@ -49,7 +31,7 @@ def import_batch(self, source_engine, target_engine, target_table_configuration,
 
         data_frame = self.attach_column_transformers(data_frame)
 
-        self.write_data_frame_to_table(data_frame, target_table_configuration, target_engine)
+        self.write_data_frame_to_table(data_frame, self.target_table_configuration, self.target_engine)
         batch_tracker.load_completed_successfully()
 
         last_key_returned = data_frame.iloc[-1][self.batch_configuration['source_unique_column']]
@@ -66,6 +48,9 @@ def write_data_frame_to_table(self, data_frame, table_configuration, target_engi
         raw = target_engine.raw_connection()
         curs = raw.cursor()
 
+        #TODO: This is assuming that our destination schema column order matches the columns in the dataframe. This
+        #isn't always correct (especially in csv sources) - therefore, we should derive the column_array from the
+        #data frames' columns.
         column_array = list(map(lambda cfg: cfg['destination']['name'], self.columns))
 
         curs.copy_from(data, destination_table, sep=',', columns=column_array, null='')
@@ -78,13 +63,12 @@ def attach_column_transformers(self, data_frame):
         self.logger.debug("Attaching column transformers")
         for column in self.columns:
             if 'column_transformer' in column:
-                #transformer = self.create_column_transformer_type(column['column_transformer'])
+                # transformer = self.create_column_transformer_type(column['column_transformer'])
                 transformer = ToUpper.execute;
                 data_frame[column['source_name']] = data_frame[column['source_name']].map(transformer)
-                #print (data_frame)
+                # print (data_frame)
         return data_frame
 
-
     def create_column_transformer_type(self, type_name):
         module = importlib.import_module(type_name)
         class_ = getattr(module, type_name)

relational_data_loader_project/CsvDataSource.py

@@ -0,0 +1,21 @@
+import logging
+import pandas
+
+
+class CsvDataSource(object):
+    def __init__(self, source_path, source_table_configuration, columns, logger=None):
+        self.logger = logger or logging.getLogger(__name__)
+        self.source_path = source_path
+        self.columns = columns
+
+    def get_data_frame(self, batch_tracker, previous_key=0):
+        path_to_csv_file = "{0}{1}.csv".format(self.source_path, self.source_table_configuration['source_table']['name'])
+
+        self.logger.debug("Starting read of file: {0}".format(path_to_csv_file))
+        data_frame = pandas.read_csv(path_to_csv_file)
+        self.logger.debug("Completed read")
+
+        batch_tracker.extract_completed_successfully(len(data_frame))
+        return data_frame
+
+

relational_data_loader_project/DataLoadManager.py

@@ -2,46 +2,50 @@
 from relational_data_loader_project.BatchDataLoader import BatchDataLoader
 from relational_data_loader_project.DestinationTableManager import DestinationTableManager
 from relational_data_loader_project.DataLoadTracker import DataLoadTracker
-from relational_data_loader_project.SourceTableManager import SourceTableManager
 import os
 import json
 
 
 class DataLoadManager(object):
-    def __init__(self, configuration_path, logger=None):
+    def __init__(self, configuration_path, data_source, logger=None):
         self.logger = logger or logging.getLogger(__name__)
         self.configuration_path = configuration_path
+        self.data_source = data_source
 
-    def start_import(self, source_engine, target_engine, full_load):
+    def start_imports(self, target_engine, full_load):
         for file in os.listdir(self.configuration_path):
-            self.start_single_import(source_engine, target_engine, file, full_load)
+            self.start_single_import(target_engine, file, full_load)
 
-    def start_single_import(self, source_engine, target_engine, configuration_name, full_load):
+    def start_single_import(self, target_engine, configuration_name, full_load):
 
         with open("{0}{1}".format(self.configuration_path, configuration_name)) as json_data:
             pipeline_configuration = json.load(json_data)
 
         data_load_tracker = DataLoadTracker(configuration_name, json_data, full_load)
 
         self.logger.debug("Execute Starting")
+
         destination_table_manager = DestinationTableManager()
 
-        columns = self.remove_invalid_columns(pipeline_configuration['source_table'], pipeline_configuration['columns'],
-                                              source_engine)
+        columns = self.data_source.get_valid_columns(pipeline_configuration['source_table'], pipeline_configuration['columns'])
 
         if full_load:
             self.logger.info("Full-load is set. Recreating the staging table.")
             destination_table_manager.create_table(pipeline_configuration['stage_table'],
                                                    columns, target_engine, drop_first=True)
 
         # Import the data.
-        batch_importer = BatchDataLoader(pipeline_configuration['source_table'], columns,
-                                         pipeline_configuration['batch'])
+        batch_importer = BatchDataLoader(self.data_source,
+                                         pipeline_configuration['source_table'],
+                                         pipeline_configuration['stage_table'],
+                                         columns,
+                                         data_load_tracker,
+                                         pipeline_configuration['batch'],
+                                         target_engine)
 
         previous_unique_column_value = 0
         while previous_unique_column_value > -1:
-            previous_unique_column_value = batch_importer.import_batch(source_engine, target_engine, pipeline_configuration['stage_table'], data_load_tracker.start_batch(), previous_unique_column_value)
-
+            previous_unique_column_value = batch_importer.import_batch(previous_unique_column_value)
 
         self.logger.info("ImportBatch Completed")
 
@@ -57,13 +61,3 @@ def start_single_import(self, source_engine, target_engine, configuration_name,
         data_load_tracker.completed_successfully()
         self.logger.info(data_load_tracker.get_statistics())
 
-    def remove_invalid_columns(self, source_table_configuration, column_configration, source_engine):
-        source_table_manager = SourceTableManager()
-        existing_columns = source_table_manager.get_columns(source_table_configuration, source_engine)
-        return list(filter(lambda column: self.column_exists(column['source_name'], existing_columns), column_configration))
-
-    def column_exists(self, column_name, existing_columns):
-        if column_name in existing_columns:
-            return True
-        self.logger.warning("Column {0} does not exist in source. It will be ignored for now, however may cause downstream issues.".format(column_name))
-        return False

relational_data_loader_project/MsSqlDataSource.py

@@ -0,0 +1,61 @@
+import logging
+import pandas
+from sqlalchemy import create_engine
+from sqlalchemy import MetaData
+from sqlalchemy.schema import Table
+
+
+class MsSqlDataSource(object):
+
+    def __init__(self, connection_string, logger=None):
+        self.logger = logger or logging.getLogger(__name__)
+        self.connection_string = connection_string
+        self.database_engine = create_engine(connection_string)
+
+    def build_select_statement(self, table_configuration, columns, batch_configuration, previous_batch_key):
+        column_array = list(map(lambda cfg: cfg['source_name'], columns))
+        column_names = ", ".join(column_array)
+
+        return "SELECT TOP ({0}) {1} FROM {2}.{3} WHERE {4} > {5} ORDER BY {4}".format(batch_configuration['size'],
+                                                                                       column_names,
+                                                                                       table_configuration[
+                                                                                           'schema'],
+                                                                                       table_configuration[
+                                                                                           'name'],
+                                                                                       batch_configuration[
+                                                                                           'source_unique_column'],
+                                                                                       previous_batch_key)
+
+    # Returns an array of configured_columns containing only columns that this data source supports. Logs invalid ones.
+    def get_valid_columns(self, table_configuration, configured_columns):
+        columns_in_database = self.get_table_columns(table_configuration)
+
+        return list(
+            filter(lambda column: self.column_exists(column['source_name'], columns_in_database), configured_columns))
+
+    def column_exists(self, column_name, columns_in_database):
+        if column_name in columns_in_database:
+            return True
+        self.logger.warning(
+            "Column {0} does not exist in source. It will be ignored for now, however may cause downstream issues.".format(
+                column_name))
+        return False
+
+    def get_table_columns(self, table_configuration):
+        metadata = MetaData()
+        self.logger.debug("Reading definition for source table {0}.{1}".format(table_configuration['schema'],
+                                                                               table_configuration['name']))
+        table = Table(table_configuration['name'], metadata, schema=table_configuration['schema'], autoload=True,
+                      autoload_with=self.database_engine)
+        return list(map(lambda column: column.name, table.columns))
+
+    def get_next_data_frame(self, table_configuration, columns, batch_configuration, batch_tracker, previous_batch_key):
+        sql = self.build_select_statement(table_configuration, columns, batch_configuration, previous_batch_key)
+
+        self.logger.debug("Starting read of SQL Statement: {0}".format(sql))
+        data_frame = pandas.read_sql_query(sql, self.database_engine)
+        self.logger.debug("Completed read")
+
+        batch_tracker.extract_completed_successfully(len(data_frame))
+
+        return data_frame

relational_data_loader_project/SourceTableManager.py

@@ -1,5 +1,6 @@
 import logging
 from relational_data_loader_project.DataLoadManager import DataLoadManager
+from relational_data_loader_project.MsSqlDataSource import MsSqlDataSource
 from sqlalchemy import create_engine
 import argparse
 
@@ -9,11 +10,12 @@
 def main(args):
 
     configure_logging(args['log_level'])
-    source_engine = create_engine(args['source-engine'])
+    data_source = MsSqlDataSource(args['source-connection-string'])
+
     destination_engine = create_engine(args['destination-engine'])
 
-    data_load_manager = DataLoadManager(args['configuration-folder'])
-    data_load_manager.start_import(source_engine, destination_engine, True)
+    data_load_manager = DataLoadManager(args['configuration-folder'], data_source)
+    data_load_manager.start_imports(destination_engine, True)
 
 
 def configure_logging(log_level):
@@ -41,8 +43,8 @@ def _log_level_string_to_int(log_level_string):
 def get_arguments():
     parser = argparse.ArgumentParser(description='Relational Data Loader')
 
-    parser.add_argument('source-engine', metavar='source-engine',
-                        help='The source engine. Eg: mssql+pyodbc://dwsource')
+    parser.add_argument('source-connection-string', metavar='source-connection-string',
+                        help='The source connections string. Eg: mssql+pyodbc://dwsource or csv://c://some//Path//To//Csv//Files//')
 
     parser.add_argument('destination-engine', metavar='destination-engine',
                         help='The destination engine. Eg: postgresql+psycopg2://postgres:xxxx@localhost/dest_dw')

relational_data_loader_project/__main__.py

@@ -0,0 +1,4 @@
+id,StringColumn1,IntColumn1,StringColumn2,DateColumn1,DateColumn2,DateTimeColumn1,LongString,UnicodeString
+1,"String Column 1",555, "String Column 2, 01-Dec-1976,,01-dec-1976 1:00 am,"This is a really long string",""
+
+