Merge pull request #2 from miksyr/handle_nan_better

miksyr · web-flow · commit 04991e513348 · 2021-09-14T16:17:50.000+01:00
handing nan better and config update
diff --git a/.flake8 b/.flake8
@@ -0,0 +1,10 @@
+[flake8]
+ignore = E501, E203, W503
+exclude =
+  .git,
+  __pycache__,
+  build,
+  dist,
+  scripts,
+  tests
+max-complexity = 10
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -0,0 +1,15 @@
+
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v3.2.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-added-large-files
+      - id: detect-private-key
+      - id: requirements-txt-fixer
+  - repo: https://github.com/psf/black
+    rev: 21.5b2
+    hooks:
+      - id: black
diff --git a/easy_postgres_engine/postgres_engine.py b/easy_postgres_engine/postgres_engine.py
@@ -1,15 +1,16 @@
 import logging
+from typing import Any, Dict, Optional, Union
 
 import pandas as pd
 import psycopg2
 import psycopg2.extras
 
 from .retry_decorator import retry
+from .utils.dataframe_functions import replace_nan_with_none_in_dataframe
 
 
 class PostgresEngine:
-
-    def __init__(self, databaseName: str, user: str, password: str, host: str = 'localhost', port: int = 5432):
+    def __init__(self, databaseName: str, user: str, password: str, host: str = "localhost", port: int = 5432):
         """
         Class for accessing Postgres databases more easily.
 
@@ -27,72 +28,79 @@ def __init__(self, databaseName: str, user: str, password: str, host: str = 'loc
         self.connection = None
         self.cursor = None
 
-    def _get_connection(self):
+    def _get_connection(self) -> None:
         try:
-            self.connection = psycopg2.connect(user=self.user, password=self.password, host=self.host, port=self.port, database=self.databaseName)
+            self.connection = psycopg2.connect(
+                user=self.user, password=self.password, host=self.host, port=self.port, database=self.databaseName
+            )
         except Exception as ex:
-            logging.exception(f'Error connecting to PostgreSQL {ex}')
+            logging.exception(f"Error connecting to PostgreSQL {ex}")
             raise ex
 
-    def _get_cursor(self, isInsertionQuery: bool):
+    def _get_cursor(self, isInsertionQuery: bool) -> None:
         if isInsertionQuery:
             self.cursor = self.connection.cursor()
         else:
             self.cursor = self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
 
-    def _close_connection(self):
+    def _close_connection(self) -> None:
         self.connection.close()
 
-    def _close_cursor(self):
+    def _close_cursor(self) -> None:
         self.cursor.close()
 
-    def close(self):
+    def close(self) -> None:
         if self.connection is not None:
             self._close_connection()
         if self.cursor is not None:
             self._close_cursor()
 
-    def create_table(self, schema: str):
+    def create_table(self, schema: str) -> None:
         self._get_connection()
         self._get_cursor(isInsertionQuery=True)
         self.cursor.execute(schema)
         try:
             self.connection.commit()
         except Exception as ex:
-            logging.exception(f'error: {ex} \nschemaQuery: {schema}')
+            logging.exception(f"error: {ex} \nschemaQuery: {schema}")
             raise ex
         finally:
             self.close()
 
-    def create_index(self, tableName: str, column: str):
+    def create_index(self, tableName: str, column: str) -> None:
         self._get_connection()
         self._get_cursor(isInsertionQuery=True)
-        indexQuery = f'CREATE INDEX IF NOT EXISTS {tableName}_{column} ON {tableName}({column});'
+        indexQuery = f"CREATE INDEX IF NOT EXISTS {tableName}_{column} ON {tableName}({column});"
         self.cursor.execute(indexQuery)
         try:
             self.connection.commit()
         except Exception as ex:
-            logging.exception(f'error: {ex} \nindexQuery: {indexQuery}')
+            logging.exception(f"error: {ex} \nindexQuery: {indexQuery}")
             raise ex
         finally:
             self.close()
 
     @retry(numRetries=5, retryDelaySeconds=3, backoffScalingFactor=2)
-    def run_select_query(self, query: str, parameters: dict = None):
+    def run_select_query_with_retry(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> pd.DataFrame:
+        return self.run_select_query(query=query, parameters=parameters)
+
+    def run_select_query(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> pd.DataFrame:
         self._get_connection()
         self._get_cursor(isInsertionQuery=False)
         self.cursor.execute(query, parameters)
         outputs = self.cursor.fetchall()
         self.close()
         outputDataframe = pd.DataFrame(outputs)
-        return outputDataframe.where(outputDataframe.notnull(), None).dropna(axis=0, how='all')
+        return replace_nan_with_none_in_dataframe(dataframe=outputDataframe)
 
     @retry(numRetries=5, retryDelaySeconds=3, backoffScalingFactor=2)
-    def run_update_query(self, query: str, parameters: dict = None, returnId: bool = True):
+    def run_update_query(
+        self, query: str, parameters: Optional[Dict[str, Any]] = None, returnId: bool = True
+    ) -> Union[None, int]:
         self._get_connection()
         self._get_cursor(isInsertionQuery=True)
         if returnId:
-            query = f'{query}\nRETURNING id'
+            query = f"{query}\nRETURNING id"
         self.cursor.execute(query, parameters)
         if returnId:
             insertedId = self.cursor.fetchone()[0]
@@ -101,7 +109,7 @@ def run_update_query(self, query: str, parameters: dict = None, returnId: bool =
         try:
             self.connection.commit()
         except Exception as ex:
-            logging.exception(f'error: {ex} \nquery: {query} \nparameters: {parameters}')
+            logging.exception(f"error: {ex} \nquery: {query} \nparameters: {parameters}")
             raise ex
         finally:
             self.close()
diff --git a/easy_postgres_engine/retry_decorator.py b/easy_postgres_engine/retry_decorator.py
@@ -4,7 +4,6 @@
 
 
 def retry(numRetries: int = 5, retryDelaySeconds: int = 3, backoffScalingFactor: int = 2):
-
     def retry_decorator(func):
         @wraps(func)
         def retry_function(*args, **kwargs):
@@ -13,11 +12,13 @@ def retry_function(*args, **kwargs):
                 try:
                     return func(*args, **kwargs)
                 except Exception as ex:
-                    exceptionMessage = f'{ex}, Retrying in {currentDelay} seconds...'
+                    exceptionMessage = f"{ex}, Retrying in {currentDelay} seconds..."
                     logging.warning(exceptionMessage)
                     sleep(currentDelay)
                     numTries -= 1
                     currentDelay *= backoffScalingFactor
             return func(*args, **kwargs)
+
         return retry_function
+
     return retry_decorator
diff --git a/easy_postgres_engine/tests/test_postgres_engine.py b/easy_postgres_engine/tests/test_postgres_engine.py
@@ -1,18 +1,18 @@
 import testing.postgresql
 from unittest import TestCase
 
-from easy_postgres_engine.tests.test_table_schema import TEST_TABLE_SCHEMA
-from easy_postgres_engine.postgres_engine import PostgresEngine
+from .test_table_schema import TEST_TABLE_SCHEMA
+from ..postgres_engine import PostgresEngine
 
 
 def create_table(postgresqlConnection):
     config = postgresqlConnection.dsn()
     dbEngine = PostgresEngine(
-        databaseName=config['database'],
-        user=config['user'],
-        password=config.get('password'),
-        port=config['port'],
-        host=config['host']
+        databaseName=config["database"],
+        user=config["user"],
+        password=config.get("password"),
+        port=config["port"],
+        host=config["host"],
     )
     dbEngine.create_table(schema=TEST_TABLE_SCHEMA)
 
@@ -27,24 +27,23 @@ def tearDownModule():
 
 
 class TestPostgresEngine(TestCase):
-
-    def __init__(self, methodName='runTest'):
+    def __init__(self, methodName="runTest"):
         super(TestPostgresEngine, self).__init__(methodName=methodName)
         self.firstCustomerId = 10
-        self.firstCustomerName = 'Mary'
+        self.firstCustomerName = "Mary"
         self.secondCustomerId = 50
-        self.secondCustomerName = 'John'
+        self.secondCustomerName = "John"
 
     def setUp(self):
         super().setUp()
         self.postgresql = Postgresql()
         config = self.postgresql.dsn()
         self.dbEngine = PostgresEngine(
-            databaseName=config['database'],
-            user=config['user'],
-            password=config.get('password'),
-            port=config['port'],
-            host=config['host']
+            databaseName=config["database"],
+            user=config["user"],
+            password=config.get("password"),
+            port=config["port"],
+            host=config["host"],
         )
 
     def tearDown(self):
@@ -60,19 +59,18 @@ def test_engine(self):
                 (%(customerId)s, %(customerName)s)
         """
         insertedId1 = self.dbEngine.run_update_query(
-            query=TEST_INSERTION_QUERY,
-            parameters={'customerId': self.firstCustomerId, 'customerName': self.firstCustomerName}
+            query=TEST_INSERTION_QUERY, parameters={"customerId": self.firstCustomerId, "customerName": self.firstCustomerName}
         )
         self.assertEqual(insertedId1, 1)
         insertedId2 = self.dbEngine.run_update_query(
             query=TEST_INSERTION_QUERY,
-            parameters={'customerId': self.secondCustomerId, 'customerName': self.secondCustomerName}
+            parameters={"customerId": self.secondCustomerId, "customerName": self.secondCustomerName},
         )
         self.assertEqual(insertedId2, 2)
 
-        queryResults = self.dbEngine.run_select_query(query='SELECT * FROM tbl_example')
-        self.assertSequenceEqual(list(queryResults['customer_id'].values), [self.firstCustomerId, self.secondCustomerId])
-        self.assertSequenceEqual(list(queryResults['customer_name'].values), [self.firstCustomerName, self.secondCustomerName])
+        queryResults = self.dbEngine.run_select_query(query="SELECT * FROM tbl_example")
+        self.assertSequenceEqual(list(queryResults["customer_id"].values), [self.firstCustomerId, self.secondCustomerId])
+        self.assertSequenceEqual(list(queryResults["customer_name"].values), [self.firstCustomerName, self.secondCustomerName])
 
         specificQueryResults = self.dbEngine.run_select_query(
             query="""
@@ -83,7 +81,7 @@ def test_engine(self):
                 WHERE
                     customer_id = %(customerId)s
             """,
-            parameters={'customerId': self.firstCustomerId}
+            parameters={"customerId": self.firstCustomerId},
         )
         self.assertEqual(len(specificQueryResults), 1)
-        self.assertEqual(specificQueryResults['customer_name'].iloc[0], self.firstCustomerName)
+        self.assertEqual(specificQueryResults["customer_name"].iloc[0], self.firstCustomerName)
diff --git a/easy_postgres_engine/utils/__init__.py b/easy_postgres_engine/utils/__init__.py
diff --git a/easy_postgres_engine/utils/dataframe_functions.py b/easy_postgres_engine/utils/dataframe_functions.py
@@ -0,0 +1,7 @@
+import numpy as np
+import pandas as pd
+
+
+def replace_nan_with_none_in_dataframe(dataframe: pd.DataFrame) -> pd.DataFrame:
+    dataframe = dataframe.where(dataframe.notnull(), None).dropna(axis=0, how="all")
+    return dataframe.replace({np.nan: None})
diff --git a/pyproject.toml b/pyproject.toml
@@ -0,0 +1,4 @@
+[tool.black]
+line-length = 127
+target-version = ['py36', 'py37', 'py38']
+include = '\.pyi?$'
diff --git a/requirements.txt b/requirements.txt
@@ -1,3 +1,7 @@
+black==21.5b2
+flake8==3.9.2
+nose2[coverage_plugin]==0.10.0
+numpy==1.19.2
 pandas==1.3.2
 psycopg2-binary==2.9.1
 testing.postgresql==1.3.0
diff --git a/setup.py b/setup.py
@@ -1,24 +1,20 @@
 from distutils.core import setup
 
 setup(
-    name='easy_postgres_engine',
-    packages=['easy_postgres_engine'],
-    version='0.2',
-    description='Engine class for easier connections to postgres databases',
-    author='Michael Doran',
-    author_email='mikrdoran@gmail.com',
-    url='https://github.com/miksyr/easy_postgres_engine',
-    download_url='https://github.com/miksyr/easy_postgres_engine/archive/v_02.tar.gz',
-    keywords=['postgreSQL', 'postgres'],
-    install_requires=[
-            'pandas==1.3.2',
-            'psycopg2-binary==2.9.1',
-            'testing.postgresql==1.3.0'
-      ],
+    name="easy_postgres_engine",
+    packages=["easy_postgres_engine"],
+    version="0.2",
+    description="Engine class for easier connections to postgres databases",
+    author="Michael Doran",
+    author_email="mikrdoran@gmail.com",
+    url="https://github.com/miksyr/easy_postgres_engine",
+    download_url="https://github.com/miksyr/easy_postgres_engine/archive/v_02.tar.gz",
+    keywords=["postgreSQL", "postgres"],
+    install_requires=["pandas==1.3.2", "psycopg2-binary==2.9.1", "testing.postgresql==1.3.0"],
     classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
         "Operating System :: OS Independent",
     ],
-    python_requires='>=3.8'
+    python_requires=">=3.8",
 )
