Refactor to BaseCopy class, other refactoring to make multiprocessing work with HDFTableCopy objects

Author: bleonard33

pandas_to_postgres/_base_copy.py

@@ -0,0 +1,89 @@
+from .utilities import (
+    create_file_object,
+    df_generator,
+    logger,
+    classification_to_pandas,
+    cast_pandas,
+    add_level_metadata,
+    HDFMetadata,
+)
+
+import pandas as pd
+from sqlalchemy.schema import AddConstraint, DropConstraint
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.sql.schema import Table
+from sqlalchemy.engine.base import Connection
+
+
+class BaseCopy(object):
+    def __init__(
+        self,
+        defer_sql_objs: bool = False,
+        conn=None,
+        table_obj=None,
+        sql_table=None,
+        csv_chunksize: int = 10 ** 6,
+    ):
+
+        self.rows = 0
+        self.columns = None
+        self.csv_chunksize = csv_chunksize
+
+        if not defer_sql_objs:
+            self.instantiate_sql_objs(conn, table_obj)
+        else:
+            self.sql_table = sql_table
+
+    def instantiate_sql_objs(self, conn, table_obj):
+        """
+        When using multiprocessing, pickling of SQLAlchemy objects in __init__ causes
+        issues, so allow for deferring until after the pickling to fetch SQLAlchemy objs
+        """
+        self.conn = conn
+        self.table_obj = table_obj
+        self.sql_table = table_obj.name
+        self.primary_key = table_obj.primary_key
+        self.foreign_keys = table_obj.foreign_key_constraints
+
+    def drop_pk(self):
+        logger.info(f"Dropping {self.sql_table} primary key")
+        try:
+            with self.conn.begin_nested():
+                self.conn.execute(DropConstraint(self.primary_key, cascade=True))
+        except SQLAlchemyError:
+            logger.info(f"{self.sql_table} primary key not found. Skipping")
+
+    def create_pk(self):
+        logger.info(f"Creating {self.sql_table} primary key")
+        self.conn.execute(AddConstraint(self.primary_key))
+
+    def drop_fks(self):
+        for fk in self.foreign_keys:
+            logger.info(f"Dropping foreign key {fk.name}")
+            try:
+                with self.conn.begin_nested():
+                    self.conn.execute(DropConstraint(fk))
+            except SQLAlchemyError:
+                logger.warn(f"Foreign key {fk.name} not found")
+
+    def create_fks(self):
+        for fk in self.foreign_keys:
+            try:
+                logger.info(f"Creating foreign key {fk.name}")
+                self.conn.execute(AddConstraint(fk))
+            except SQLAlchemyError:
+                logger.warn(f"Error creating foreign key {fk.name}")
+
+    def truncate(self):
+        logger.info(f"Truncating {self.sql_table}")
+        self.conn.execute(f"TRUNCATE TABLE {self.sql_table};")
+
+    def analyze(self):
+        logger.info(f"Analyzing {self.sql_table}")
+        self.conn.execute(f"ANALYZE {self.sql_table};")
+
+    def copy_from_file(self, file_object):
+        cur = self.conn.connection.cursor()
+        cols = ", ".join([f"{col}" for col in self.columns])
+        sql = f"COPY {self.sql_table} ({cols}) FROM STDIN WITH CSV HEADER FREEZE"
+        cur.copy_expert(sql=sql, file=file_object)

pandas_to_postgres/copy_df.py

@@ -6,79 +6,29 @@
     add_level_metadata,
 )
 
+from ._base_copy import BaseCopy
+
 import pandas as pd
 from sqlalchemy.sql.schema import Table
 from sqlalchemy.engine.base import Connection
-from sqlalchemy.schema import AddConstraint, DropConstraint
-from sqlalchemy.exc import SQLAlchemyError
 
 
-class DataFrameCopy(object):
+class DataFrameCopy(BaseCopy):
     def __init__(
         self,
-        conn: Connection,
-        table_obj: Table,
         df: pd.DataFrame,
-        levels: dict = None,
+        defer_sql_objs: bool = False,
+        conn: Connection = None,
+        table_obj: Table = None,
         csv_chunksize: int = 10 ** 6,
+        levels: dict = None,
     ):
-        self.conn = conn
-        self.table_obj = table_obj
-        self.sql_table = self.table_obj.name
+        BaseCopy(defer_sql_objs, conn, table_obj, csv_chunksize)
+
         self.df = df
         self.levels = levels
         self.columns = self.df.columns
         self.rows = self.df.shape[0]
-        self.csv_chunksize = csv_chunksize
-        self.primary_key = self.table_obj.primary_key
-        self.foreign_keys = self.table_obj.foreign_key_constraints
-
-    def close_conn(self):
-        self.conn.close()
-        del self.conn
-
-    def drop_pk(self):
-        logger.info(f"Dropping {self.sql_table} primary key")
-        try:
-            with self.conn.begin_nested():
-                self.conn.execute(DropConstraint(self.primary_key, cascade=True))
-        except SQLAlchemyError:
-            logger.info(f"{self.sql_table} primary key not found. Skipping")
-
-    def create_pk(self):
-        logger.info(f"Creating {self.sql_table} primary key")
-        self.conn.execute(AddConstraint(self.primary_key))
-
-    def drop_fks(self):
-        for fk in self.foreign_keys:
-            logger.info(f"Dropping foreign key {fk.name}")
-            try:
-                with self.conn.begin_nested():
-                    self.conn.execute(DropConstraint(fk))
-            except SQLAlchemyError:
-                logger.warn(f"Foreign key {fk.name} not found")
-
-    def create_fks(self):
-        for fk in self.foreign_keys:
-            try:
-                logger.info(f"Creating foreign key {fk.name}")
-                self.conn.execute(AddConstraint(fk))
-            except SQLAlchemyError:
-                logger.warn(f"Error creating foreign key {fk.name}")
-
-    def truncate(self):
-        logger.info(f"Truncating {self.sql_table}")
-        self.conn.execute(f"TRUNCATE TABLE {self.sql_table};")
-
-    def analyze(self):
-        logger.info(f"Analyzing {self.sql_table}")
-        self.conn.execute(f"ANALYZE {self.sql_table};")
-
-    def copy_from_file(self, file_object):
-        cur = self.conn.connection.cursor()
-        cols = ", ".join([f"{col}" for col in self.columns])
-        sql = f"COPY {self.sql_table} ({cols}) FROM STDIN WITH CSV HEADER FREEZE"
-        cur.copy_expert(sql=sql, file=file_object)
 
     def format_df(self):
         # Handle NaN --> None type casting

pandas_to_postgres/copy_hdf.py

@@ -5,91 +5,49 @@
     classification_to_pandas,
     cast_pandas,
     add_level_metadata,
+    HDFMetadata,
 )
 
-import pandas as pd
-from sqlalchemy.schema import AddConstraint, DropConstraint
-from sqlalchemy.exc import SQLAlchemyError
-
+from ._base_copy import BaseCopy
 
-class HDFTableCopy(object):
-
-    rows = 0
-    columns = None
+import pandas as pd
+from sqlalchemy.sql.schema import Table
+from sqlalchemy.engine.base import Connection
+
+
+class HDFTableCopy(BaseCopy):
+    def __init__(
+        self,
+        hdf_tables: list,
+        hdf_meta: HDFMetadata,
+        defer_sql_objs: bool = False,
+        conn=None,
+        table_obj=None,
+        sql_table=None,
+        csv_chunksize: int = 10 ** 6,
+    ):
+        BaseCopy.__init__(
+            self, defer_sql_objs, conn, table_obj, sql_table, csv_chunksize
+        )
 
-    def __init__(self, sql_table, hdf_tables, hdf_meta, csv_chunksize=10 ** 6):
-        self.sql_table = sql_table
         self.hdf_tables = hdf_tables
-        self.csv_chunksize = csv_chunksize
 
         # Info from the HDFMetadata object
         self.levels = hdf_meta.levels
         self.file_name = hdf_meta.file_name
         self.hdf_chunksize = hdf_meta.chunksize
 
-    def table_metadata(self):
-        self.table_obj = db.metadata.tables[self.sql_table]
-        self.primary_key = self.table_obj.primary_key
-        self.foreign_keys = self.table_obj.foreign_key_constraints
-
-    def set_conn(self, conn):
-        self.conn = conn
-
-    def delete_conn(self):
-        del self.conn
-
-    def drop_pk(self):
-        logger.info(f"Dropping {self.sql_table} primary key")
-        try:
-            with self.conn.begin_nested():
-                self.conn.execute(DropConstraint(self.primary_key, cascade=True))
-        except SQLAlchemyError:
-            logger.info(f"{self.sql_table} primary key not found. Skipping")
-
-    def create_pk(self):
-        logger.info(f"Creating {self.sql_table} primary key")
-        self.conn.execute(AddConstraint(self.primary_key))
-
-    def drop_fks(self):
-        for fk in self.foreign_keys:
-            logger.info(f"Dropping foreign key {fk.name}")
-            try:
-                with self.conn.begin_nested():
-                    self.conn.execute(DropConstraint(fk))
-            except SQLAlchemyError:
-                logger.warn(f"Foreign key {fk.name} not found")
-
-    def create_fks(self):
-        for fk in self.foreign_keys:
-            try:
-                logger.info(f"Creating foreign key {fk.name}")
-                self.conn.execute(AddConstraint(fk))
-            except SQLAlchemyError:
-                logger.warn(f"Error creating foreign key {fk.name}")
-
-    def truncate(self):
-        logger.info(f"Truncating {self.sql_table}")
-        self.conn.execute(f"TRUNCATE TABLE {self.sql_table};")
-
-    def analyze(self):
-        logger.info(f"Analyzing {self.sql_table}")
-        self.conn.execute(f"ANALYZE {self.sql_table};")
-
-    def copy_from_file(self, file_object):
-        cur = self.conn.connection.cursor()
-        cols = ", ".join([f"{col}" for col in self.columns])
-        sql = f"COPY {self.sql_table} ({cols}) FROM STDIN WITH CSV HEADER FREEZE"
-        cur.copy_expert(sql=sql, file=file_object)
-
     def copy_table(self):
-        self.table_metadata()
         self.drop_fks()
         self.drop_pk()
+
+        # These need to be one transaction to use COPY FREEZE
         with self.conn.begin():
             self.truncate()
             self.hdf_to_pg()
-            self.create_pk()
-            self.create_fks()
+
+        self.create_pk()
+        self.create_fks()
         self.analyze()
 
     def hdf_to_pg(self):
@@ -126,8 +84,26 @@ def hdf_to_pg(self):
 
 
 class ClassificationHDFTableCopy(HDFTableCopy):
-    def __init__(self, sql_table, hdf_tables, hdf_meta, csv_chunksize=10 ** 6):
-        HDFTableCopy.__init__(self, sql_table, hdf_tables, hdf_meta, csv_chunksize)
+    def __init__(
+        self,
+        hdf_tables: list,
+        hdf_meta: HDFMetadata,
+        defer_sql_objs: bool = False,
+        conn=None,
+        table_obj=None,
+        sql_table: str = None,
+        csv_chunksize: int = 10 ** 6,
+    ):
+        HDFTableCopy.__init__(
+            self,
+            hdf_tables,
+            hdf_meta,
+            defer_sql_objs,
+            conn,
+            table_obj,
+            sql_table,
+            csv_chunksize,
+        )
 
     def hdf_to_pg(self):
         if self.hdf_tables is None:
@@ -158,8 +134,26 @@ def hdf_to_pg(self):
 
 
 class BigHDFTableCopy(HDFTableCopy):
-    def __init__(self, sql_table, hdf_tables, hdf_meta, csv_chunksize=10 ** 6):
-        HDFTableCopy.__init__(self, sql_table, hdf_tables, hdf_meta, csv_chunksize)
+    def __init__(
+        self,
+        hdf_tables: list,
+        hdf_meta: HDFMetadata,
+        defer_sql_objs: bool = False,
+        conn=None,
+        table_obj=None,
+        sql_table=None,
+        csv_chunksize: int = 10 ** 6,
+    ):
+        HDFTableCopy.__init__(
+            self,
+            hdf_tables,
+            hdf_meta,
+            defer_sql_objs,
+            conn,
+            table_obj,
+            sql_table,
+            csv_chunksize,
+        )
 
     def hdf_to_pg(self):
         if self.hdf_tables is None: