Modify features related with log

cryeo · Sep 7, 2018 · 105ab02 · 105ab02
1 parent 09abf3f
commit 105ab02
Show file tree

Hide file tree

Showing 7 changed files with 12 additions and 17 deletions.
diff --git a/README.md b/README.md
@@ -28,7 +28,8 @@ ru = RedshiftUnloader(host='<redshift host>',
                       s3_bucket='<s3 bucket name>',
                       access_key_id='<aws access key id>',
                       secret_access_key='<aws secret access key>',
-                      region='<aws region>')
+                      region='<aws region>',
+                      verbose=False)
 
 # If you don't need header, set with_header as False
 ru.unload(query="SELECT * FROM my_table WHERE log_time >= ''", 

diff --git a/redshift_unloader/__init__.py b/redshift_unloader/__init__.py
@@ -1,4 +1,4 @@
 from redshift_unloader.redshift_unloader import RedshiftUnloader
 
 
-__version__ = '0.1.0'
+__version__ = '0.1.1'
diff --git a/redshift_unloader/logger.py b/redshift_unloader/logger.py
@@ -1,9 +1,8 @@
 import logging
 
 logging.basicConfig(
-    level=logging.INFO,
+    level=logging.WARNING,
     format="[%(asctime)s / %(name)s] [%(levelname)s] %(message)s"
 )
 
 logger = logging.getLogger("redshift-unloader")
-logger.setLevel(logging.DEBUG)
diff --git a/redshift_unloader/redshift_unloader.py b/redshift_unloader/redshift_unloader.py
@@ -3,7 +3,7 @@
 import tempfile
 import uuid
 import gzip
-import io
+import logging
 
 from redshift_unloader.credential import Credential
 from redshift_unloader.redshift import Redshift
@@ -21,7 +21,7 @@ class RedshiftUnloader:
 
     def __init__(self, host: str, port: int, user: str, password: str,
                  database: str, s3_bucket: str, access_key_id: str,
-                 secret_access_key: str, region: str) -> None:
+                 secret_access_key: str, region: str, verbose: bool = False) -> None:
         credential = Credential(
             access_key_id=access_key_id, secret_access_key=secret_access_key)
         self.__redshift = Redshift(
@@ -32,6 +32,8 @@ def __init__(self, host: str, port: int, user: str, password: str,
             database=database,
             credential=credential)
         self.__s3 = S3(credential=credential, bucket=s3_bucket, region=region)
+        if verbose:
+            logger.setLevel(logging.DEBUG)
 
     def unload(self, query: str, filename: str, with_header: bool = True) -> None:
         session_id = self.__generate_session_id()
@@ -52,6 +54,7 @@ def unload(self, query: str, filename: str, with_header: bool = True) -> None:
             delimiter=',',
             null_string='',
             add_quotes=True,
+            escape=True,
             allow_overwrite=True)
 
         logger.debug("Fetch the list of objects")

diff --git a/tests/test_redshift.py b/tests/test_redshift.py
@@ -4,7 +4,6 @@
 from mock import PropertyMock
 from collections import namedtuple
 
-from redshift_unloader.logger import logger
 from redshift_unloader.credential import Credential
 from redshift_unloader.redshift import Redshift
 
@@ -25,10 +24,9 @@ def setUp(self, mock_connect):
         self.credential = Credential(access_key_id=self.ACCESS_KEY_ID, secret_access_key=self.SECRET_ACCESS_KEY)
         self.redshift = Redshift(host=self.HOST, port=self.PORT, user=self.USER, password=self.PASSWORD,
                                  database=self.DATABASE, credential=self.credential)
-        logger.disabled = True
 
     def tearDown(self):
-        logger.disabled = False
+        pass
 
     def test_get_columns(self):
         query = "SELECT * FROM some_table WHERE date_column >= '2018-01-01'"

diff --git a/tests/test_redshift_unloader.py b/tests/test_redshift_unloader.py
@@ -4,7 +4,6 @@
 from mock import call
 
 from redshift_unloader import RedshiftUnloader
-from redshift_unloader.logger import logger
 
 
 class TestRedshiftUnloader(unittest.TestCase):
@@ -28,10 +27,9 @@ def setUp(self, mock_redshift, mock_s3):
                                          database=self.DATABASE, s3_bucket=self.S3_BUCKET,
                                          access_key_id=self.ACCESS_KEY_ID, secret_access_key=self.SECRET_ACCESS_KEY,
                                          region=self.REGION)
-        logger.disabled = True
 
     def tearDown(self):
-        logger.disabled = False
+        pass
 
     @mock.patch('builtins.open')
     @mock.patch('shutil.copyfileobj')
@@ -73,6 +71,7 @@ def test_unload(self, mock_gettempdir, mock_mkdir, mock_rmtree, mock_copyfileobj
                                                                                      delimiter=',',
                                                                                      null_string='',
                                                                                      add_quotes=True,
+                                                                                     escape=True,
                                                                                      allow_overwrite=True)
 
             self.unloader._RedshiftUnloader__s3.list.assert_called_once_with(s3_path)

diff --git a/tests/test_s3.py b/tests/test_s3.py
@@ -5,7 +5,6 @@
 
 from moto import mock_s3
 
-from redshift_unloader.logger import logger
 from redshift_unloader.s3 import S3
 from redshift_unloader.credential import Credential
 
@@ -35,13 +34,9 @@ def setUp(self):
         self.credential = Credential(access_key_id=self.ACCESS_KEY_ID, secret_access_key=self.SECRET_ACCESS_KEY)
         self.s3 = S3(self.credential, bucket=self.BUCKET, region=self.REGION)
 
-        logger.disabled = True
-
     def tearDown(self):
         self.mock_s3.stop()
 
-        logger.disabled = False
-
     def test_uri(self):
         self.assertEqual(self.s3.uri('/path/to'), f's3://{self.BUCKET}/path/to')
         self.assertEqual(self.s3.uri('/path/to/'), f's3://{self.BUCKET}/path/to/')