Load via INSERT on Redshift when scratch bucket not available

metrics/coverage_high_water_mark

@@ -1 +1 @@
-93.5700
+93.6900

metrics/mypy_high_water_mark

@@ -1 +1 @@
-92.3400
+92.3500

records_mover/db/loader.py

@@ -47,6 +47,12 @@ def temporary_loadable_directory_loc(self) -> Iterator[BaseDirectoryUrl]:
         with TemporaryDirectory(prefix='temporary_loadable_directory_loc') as dirname:
             yield FilesystemDirectoryUrl(dirname)
 
+    def has_temporary_loadable_directory_loc(self) -> bool:
+        # The default implementation uses the local filesystem where
+        # Records Mover runs, and we assume we can make temporary
+        # files.
+        return True
+
     @abstractmethod
     def known_supported_records_formats_for_load(self) -> List[BaseRecordsFormat]:
         """Candidates to look through when negotiating a common records format

records_mover/db/redshift/loader.py

@@ -10,8 +10,8 @@
 import logging
 from .records_copy import redshift_copy_options
 from ...records.load_plan import RecordsLoadPlan
-from ..errors import CredsDoNotSupportS3Import
-from typing import Optional, Union, Callable, ContextManager, List, Iterator
+from ..errors import CredsDoNotSupportS3Import, NoTemporaryBucketConfiguration
+from typing import Optional, Union, List, Iterator
 from ...url import BaseDirectoryUrl
 from botocore.credentials import Credentials
 from ...records.delimited import complain_on_unhandled_hints
@@ -23,11 +23,19 @@ class RedshiftLoader(LoaderFromRecordsDirectory):
     def __init__(self,
                  db: Union[sqlalchemy.engine.Engine, sqlalchemy.engine.Connection],
                  meta: sqlalchemy.MetaData,
-                 temporary_s3_directory_loc: Callable[[], ContextManager[BaseDirectoryUrl]])\
+                 s3_temp_base_loc: Optional[BaseDirectoryUrl])\
             -> None:
         self.db = db
         self.meta = meta
-        self.temporary_s3_directory_loc = temporary_s3_directory_loc
+        self.s3_temp_base_loc = s3_temp_base_loc
+
+    @contextmanager
+    def temporary_s3_directory_loc(self) -> Iterator[BaseDirectoryUrl]:
+        if self.s3_temp_base_loc is None:
+            raise NoTemporaryBucketConfiguration('Please provide a scratch S3 URL in your config')
+        else:
+            with self.s3_temp_base_loc.temporary_directory() as temp_loc:
+                yield temp_loc
 
     def load(self,
              schema: str,
@@ -158,3 +166,6 @@ def best_scheme_to_load_from(self) -> str:
     def temporary_loadable_directory_loc(self) -> Iterator[BaseDirectoryUrl]:
         with self.temporary_s3_directory_loc() as temp_loc:
             yield temp_loc
+
+    def has_temporary_loadable_directory_loc(self) -> bool:
+        return self.s3_temp_base_loc is not None

records_mover/db/redshift/redshift_db_driver.py

@@ -10,15 +10,14 @@
                              num_digits)
 from .sql import schema_sql_from_admin_views
 import timeout_decorator
-from contextlib import contextmanager
-from typing import Iterator, Optional, Union, Dict, List, Tuple
+from typing import Optional, Union, Dict, List, Tuple
 from ...url.base import BaseDirectoryUrl
 from records_mover.db.quoting import quote_group_name, quote_schema_and_table
 from .unloader import RedshiftUnloader
 from ..unloader import Unloader
 from .loader import RedshiftLoader
 from ..loader import LoaderFromRecordsDirectory
-from ..errors import NoTemporaryBucketConfiguration
+
 
 logger = logging.getLogger(__name__)
 
@@ -33,11 +32,11 @@ def __init__(self,
         self._redshift_loader =\
             RedshiftLoader(db=db,
                            meta=self.meta,
-                           temporary_s3_directory_loc=self.temporary_s3_directory_loc)
+                           s3_temp_base_loc=s3_temp_base_loc)
         self._redshift_unloader =\
             RedshiftUnloader(db=db,
                              table=self.table,
-                             temporary_s3_directory_loc=self.temporary_s3_directory_loc)
+                             s3_temp_base_loc=s3_temp_base_loc)
 
     def schema_sql(self, schema: str, table: str) -> str:
         out = schema_sql_from_admin_views(schema, table, self.db)
@@ -46,14 +45,6 @@ def schema_sql(self, schema: str, table: str) -> str:
         else:
             return out
 
-    @contextmanager
-    def temporary_s3_directory_loc(self) -> Iterator[BaseDirectoryUrl]:
-        if self.s3_temp_base_loc is None:
-            raise NoTemporaryBucketConfiguration('Please provide a scratch S3 URL in your config')
-        else:
-            with self.s3_temp_base_loc.temporary_directory() as temp_loc:
-                yield temp_loc
-
     # if this timeout goes off (at least for Redshift), it's probably
     # because memory is filling because sqlalchemy's cache of all
     # tables and columns filled up memory in the job.

records_mover/db/redshift/unloader.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from sqlalchemy_redshift.commands import UnloadFromSelect
 from ...records.records_directory import RecordsDirectory
 import sqlalchemy
@@ -11,10 +12,10 @@
 from ...records.records_format import (
     BaseRecordsFormat, DelimitedRecordsFormat, ParquetRecordsFormat
 )
-from typing import Union, Callable, Optional, ContextManager, List
+from typing import Union, Callable, Optional, List, Iterator
 from ...url.base import BaseDirectoryUrl
 from botocore.credentials import Credentials
-from ..errors import CredsDoNotSupportS3Export
+from ..errors import CredsDoNotSupportS3Export, NoTemporaryBucketConfiguration
 from ...records.delimited import complain_on_unhandled_hints
 from ..unloader import Unloader
 
@@ -26,11 +27,19 @@ class RedshiftUnloader(Unloader):
     def __init__(self,
                  db: Union[sqlalchemy.engine.Engine, sqlalchemy.engine.Connection],
                  table: Callable[[str, str], Table],
-                 temporary_s3_directory_loc: Callable[[], ContextManager[BaseDirectoryUrl]],
+                 s3_temp_base_loc: Optional[BaseDirectoryUrl],
                  **kwargs) -> None:
         super().__init__(db=db)
         self.table = table
-        self.temporary_s3_directory_loc = temporary_s3_directory_loc
+        self.s3_temp_base_loc = s3_temp_base_loc
+
+    @contextmanager
+    def temporary_s3_directory_loc(self) -> Iterator[BaseDirectoryUrl]:
+        if self.s3_temp_base_loc is None:
+            raise NoTemporaryBucketConfiguration('Please provide a scratch S3 URL in your config')
+        else:
+            with self.s3_temp_base_loc.temporary_directory() as temp_loc:
+                yield temp_loc
 
     def unload_to_s3_directory(self,
                                schema: str,

records_mover/records/mover.py

@@ -2,7 +2,7 @@
                       SupportsToFileobjsSource,
                       FileobjsSource, SupportsToDataframesSource)
 from .targets.base import (RecordsTarget, SupportsMoveFromRecordsDirectory,
-                           SupportsMoveFromTempLocAfterFillingIt,
+                           MightSupportMoveFromTempLocAfterFillingIt,
                            MightSupportMoveFromFileobjsSource,
                            SupportsMoveFromDataframes)
 from .sources import base as sources_base
@@ -123,8 +123,9 @@ def move(records_source: RecordsSource,
                 as fileobjs_source:
             return move(fileobjs_source, records_target, processing_instructions)
     elif (isinstance(records_source, SupportsMoveToRecordsDirectory) and
-          isinstance(records_target, SupportsMoveFromTempLocAfterFillingIt) and
-          records_source.has_compatible_format(records_target)):
+          isinstance(records_target, MightSupportMoveFromTempLocAfterFillingIt) and
+          records_source.has_compatible_format(records_target) and
+          records_target.can_move_from_temp_loc_after_filling_it()):
         logger.info(f"Mover: copying from {records_source} to {records_target} "
                     f"by filling in a temporary location...")
         return records_target.move_from_temp_loc_after_filling_it(records_source,

records_mover/records/targets/base.py

@@ -123,7 +123,14 @@ def can_move_from_fileobjs_source(self) -> bool:
         pass
 
 
-class SupportsMoveFromTempLocAfterFillingIt(NegotiatesRecordsFormat, metaclass=ABCMeta):
+class MightSupportMoveFromTempLocAfterFillingIt(NegotiatesRecordsFormat, metaclass=ABCMeta):
+    @abstractmethod
+    def can_move_from_temp_loc_after_filling_it(self) -> bool:
+        """Returns True if target as currently configured can be handed a
+        temporary location and fill it.
+        """
+        pass
+
     @abstractmethod
     def move_from_temp_loc_after_filling_it(self,
                                             records_source:

records_mover/records/targets/data_url.py

@@ -1,8 +1,6 @@
 from ..results import MoveResult
 from ..records_directory import RecordsDirectory
 from .base import (SupportsMoveFromDataframes,
-                   SupportsMoveFromTempLocAfterFillingIt,
-                   SupportsMoveToRecordsDirectory,
                    SupportsMoveFromRecordsDirectory)
 from ..processing_instructions import ProcessingInstructions
 from ...url.base import BaseFileUrl
@@ -16,7 +14,6 @@
 
 
 class DataUrlTarget(SupportsMoveFromDataframes,
-                    SupportsMoveFromTempLocAfterFillingIt,
                     SupportsMoveFromRecordsDirectory):
     def __init__(self,
                  output_loc: BaseFileUrl,
@@ -63,25 +60,6 @@ def move_from_records_directory(self,
                               records_format.generate_filename('data'): self.output_loc.url
                           })
 
-    def move_from_temp_loc_after_filling_it(self,
-                                            records_source:
-                                            SupportsMoveToRecordsDirectory,
-                                            processing_instructions:
-                                            ProcessingInstructions) -> MoveResult:
-        pis = processing_instructions
-        records_format = records_source.compatible_format(self)
-        if records_format is None:
-            raise NotImplementedError("No compatible records format between "
-                                      f"{records_source} and {self}")
-        with self.output_loc.temporary_directory() as temp_loc:
-            directory = RecordsDirectory(records_loc=temp_loc)
-            records_source.\
-                move_to_records_directory(directory,
-                                          records_format=records_format,
-                                          processing_instructions=pis)
-            return self.move_from_records_directory(directory,
-                                                    processing_instructions)
-
     def can_move_from_this_format(self,
                                   source_records_format: BaseRecordsFormat) -> bool:
         if self.records_format is None:

records_mover/records/targets/table/move_from_dataframes_source.py

@@ -30,25 +30,27 @@ def __init__(self,
         super().__init__(prep, target_table_details, processing_instructions)
 
     def move(self) -> MoveResult:
-        if len(self.table_target.known_supported_records_formats()) != 0:
-            if self.table_target.can_move_from_fileobjs_source():
-                return self.move_from_dataframes_source_via_fileobjs()
-            else:
-                # Some databases, like Redshift, can't load from a
-                # stream, but can load from files on an object store
-                # they're pointed to.
-                return self.move_from_dataframes_source_via_records_directory()
+        target_supports_formats = len(self.table_target.known_supported_records_formats()) != 0
+        if (target_supports_formats and self.table_target.can_move_from_fileobjs_source()):
+            return self.move_from_dataframes_source_via_fileobjs()
+        elif (target_supports_formats and
+              self.table_target.can_move_from_temp_loc_after_filling_it()):
+            # Some databases, like Redshift, can't load from a
+            # stream, but can load from files on an object store
+            # they're pointed to.
+            return self.move_from_dataframes_source_via_temporary_records_directory()
         else:
             logger.info("Known formats for target database: "
                         f"{self.table_target.known_supported_records_formats()}")
             logger.info("Table target can move from fileobjs source? "
                         f"{self.table_target.can_move_from_fileobjs_source()}")
             logger.warning("Loading via INSERT statement as this DB "
-                           "driver does not yet support more direct LOAD methods.  "
+                           "driver does not yet support or is not configured for "
+                           "more direct load methods.  "
                            "This may be very slow.")
             return self.move_from_dataframes_source_via_insert()
 
-    def move_from_dataframes_source_via_records_directory(self) -> MoveResult:
+    def move_from_dataframes_source_via_temporary_records_directory(self) -> MoveResult:
         records_format = next(iter(self.table_target.known_supported_records_formats()), None)
         with self.dfs_source.to_fileobjs_source(self.processing_instructions,
                                                 records_format) as fileobjs_source:

records_mover/records/targets/table/target.py

@@ -1,6 +1,6 @@
 from records_mover.records.targets.base import (
     SupportsMoveFromRecordsDirectory,
-    SupportsMoveFromTempLocAfterFillingIt,
+    MightSupportMoveFromTempLocAfterFillingIt,
     MightSupportMoveFromFileobjsSource,
     SupportsMoveFromDataframes,
 )
@@ -31,7 +31,7 @@
 
 
 class TableRecordsTarget(SupportsMoveFromRecordsDirectory,
-                         SupportsMoveFromTempLocAfterFillingIt,
+                         MightSupportMoveFromTempLocAfterFillingIt,
                          MightSupportMoveFromFileobjsSource,
                          SupportsMoveFromDataframes,
                          TargetTableDetails):
@@ -109,6 +109,13 @@ def can_move_from_this_format(self,
             return False
         return loader.can_load_this_format(source_records_format)
 
+    def can_move_from_temp_loc_after_filling_it(self) -> bool:
+        driver = self.db_driver(self.db_engine)
+        loader = driver.loader()
+        if loader is None:
+            return False
+        return loader.has_temporary_loadable_directory_loc()
+
     def move_from_temp_loc_after_filling_it(self,
                                             records_source:
                                             SupportsMoveToRecordsDirectory,

tests/unit/airflow/test_google_cloud_credentials_hook.py

@@ -0,0 +1,16 @@
+from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
+from records_mover.airflow.hooks.google_cloud_credentials_hook import GoogleCloudCredentialsHook
+from mock import Mock
+import unittest
+
+
+class TestGoogleCloudCredentialsHook(unittest.TestCase):
+    def test_get_conn(self):
+        mock_init = Mock('__init__')
+        GoogleCloudBaseHook.__init__ = mock_init
+        mock_init.return_value = None
+        hook = GoogleCloudCredentialsHook()
+        mock_get_credentials = Mock('get_credentials')
+        hook._get_credentials = mock_get_credentials
+        conn = hook.get_conn()
+        self.assertEqual(conn, mock_get_credentials.return_value)

tests/unit/db/redshift/test_loader.py

@@ -9,12 +9,12 @@ class TestRedshiftLoader(unittest.TestCase):
     def setUp(self):
         self.mock_db = Mock(name='db')
         self.mock_meta = Mock(name='meta')
-        self.mock_temporary_s3_directory_loc = MagicMock(name='temporary_s3_directory_loc')
+        self.s3_temp_base_loc = MagicMock(name='s3_temp_base_loc')
 
         self.redshift_loader =\
             RedshiftLoader(db=self.mock_db,
                            meta=self.mock_meta,
-                           temporary_s3_directory_loc=self.mock_temporary_s3_directory_loc)
+                           s3_temp_base_loc=self.s3_temp_base_loc)
 
     @patch('records_mover.db.redshift.loader.redshift_copy_options')
     @patch('records_mover.db.redshift.loader.ProcessingInstructions')
@@ -80,7 +80,7 @@ def test_load_non_s3(self,
         mock_directory = Mock(name='directory')
         mock_directory.scheme = 'mumble'
 
-        mock_temp_s3_loc = self.mock_temporary_s3_directory_loc.return_value.__enter__.return_value
+        mock_temp_s3_loc = self.s3_temp_base_loc.temporary_directory().__enter__()
         mock_s3_directory = mock_directory.copy_to.return_value
         mock_s3_directory.scheme = 's3'
 

tests/unit/db/redshift/test_redshift_db_driver.py

@@ -1,5 +1,4 @@
 from .base_test_redshift_db_driver import BaseTestRedshiftDBDriver
-from records_mover.db.redshift.redshift_db_driver import NoTemporaryBucketConfiguration
 from unittest.mock import patch
 import sqlalchemy
 
@@ -102,9 +101,3 @@ def test_type_for_floating_point(self):
                                                                 input_fp_significand_bits)
             self.assertEqual(type(actual_col_type), sqlalchemy.sql.sqltypes.Float)
             self.assertEqual(actual_col_type.precision, expected_fp_significand_bits)
-
-    def test_temporary_s3_directory_loc_unset(self):
-        self.redshift_db_driver.s3_temp_base_loc = None
-        with self.assertRaises(NoTemporaryBucketConfiguration):
-            with self.redshift_db_driver.temporary_s3_directory_loc():
-                pass

tests/unit/db/redshift/test_unloader.py

@@ -12,20 +12,19 @@ def test_can_unload_this_format_true(self,
                                          mock_redshift_unload_options):
         mock_db = Mock(name='db')
         mock_table = Mock(name='table')
-        mock_temporary_s3_directory_loc = Mock(name='temporary_s3_directory_loc')
 
         mock_target_records_format = Mock(name='target_records_format', spec=DelimitedRecordsFormat)
         mock_unload_plan = mock_RecordsUnloadPlan.return_value
         mock_unload_plan.records_format = mock_target_records_format
 
         mock_processing_instructions = mock_unload_plan.processing_instructions
-        mock_temporary_s3_directory_loc = Mock(name='temporary_s3_directory_loc')
+        mock_s3_temp_base_loc = Mock(name='s3_temp_base_loc')
         mock_target_records_format.hints = {}
 
         redshift_unloader =\
             RedshiftUnloader(db=mock_db,
                              table=mock_table,
-                             temporary_s3_directory_loc=mock_temporary_s3_directory_loc)
+                             s3_temp_base_loc=mock_s3_temp_base_loc)
         out = redshift_unloader.can_unload_this_format(mock_target_records_format)
         mock_RecordsUnloadPlan.\
             assert_called_with(records_format=mock_target_records_format)