Merge branch 'development-state-handling-command-line-option' into development

RuslanBergenov · RuslanBergenov · commit ef9d6b2e63ea · 2021-10-21T18:12:30.000-06:00
# Conflicts:
#	.github/workflows/python-package.yml
diff --git a/README.md b/README.md
@@ -456,7 +456,16 @@ Add the following files to *sandbox* directory under project root directory:
     "project_id": "{your-project-id}",
     "dataset_id": "{your_dataset_id}",
     "validate_records":  false
-  }   
+  }     
+    ```
+  
+- **target_config_merge_state_false_flag.json**:  
+  ```
+  {
+    "project_id": "{your-project-id}",
+    "dataset_id": "{your_dataset_id}",
+    "merge_state_messages": 0
+  }     
     ```
 ## Config files in this project
 
diff --git a/target_bigquery/__init__.py b/target_bigquery/__init__.py
@@ -11,6 +11,7 @@
 from target_bigquery.encoders import DecimalEncoder
 from target_bigquery.process import process
 from target_bigquery.utils import emit_state, ensure_dataset
+from target_bigquery.state import State, LiteralState
 
 logger = singer.get_logger()
 
@@ -21,12 +22,24 @@ def main():
     parser.add_argument("-c", "--config", help="Config file", required=True)
     parser.add_argument("-t", "--tables", help="Table configs file", required=False)
     parser.add_argument("-s", "--state", help="Initial state file", required=False)
+
+    # https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
+    parser.add_argument('--merge_state_messages', help="Merge many state messages to construct a state file",
+                        dest='merge_state_messages', action='store_true')
+    parser.add_argument('--no-merge_state_messages',
+                        help="Don't merge many state messages into one message. The latest state message becomes the state file.",
+                        dest='merge_state_messages', action='store_false')
+    parser.set_defaults(merge_state_messages=None)
+    # default needs to be None. If it's None, it means it's not supplied and we need to check the config file
+    # if default is True here, then setting it in config file will not work
+    # in the config file, default will be True
+
     parser.add_argument("-ph", "--processhandler",
                         help="Defines the loading process. Partial loads by default.",
                         required=False,
                         choices=["load-job", "partial-load-job", "bookmarks-partial-load-job"],
                         default="partial-load-job"
-    )
+                        )
 
     flags = parser.parse_args()
 
@@ -58,10 +71,22 @@ def main():
     location = config.get("location", "US")
     validate_records = config.get("validate_records", True)
     add_metadata_columns = config.get("add_metadata_columns", True)
+
+    # we can pass merge state option via CLI param
+    merge_state_messages_cli = flags.merge_state_messages
+
+    # we can pass merge state option via config file per Meltano request
+    merge_state_messages_config = config.get("merge_state_messages", True)
+
+    # merge state option via CLI trumps one passed via config file
+    # we need to check if CLI option was passed at all. if not, we check the config file
+    merge_state_messages = merge_state_messages_cli if type(
+        merge_state_messages_cli) == bool else merge_state_messages_config
+
     project_id, dataset_id = config["project_id"], config["dataset_id"]
 
     table_configs = tables.get("streams", {})
-    max_cache = 1024 * 1024 * config.get("max_cache", 50) # this is needed for partial loads
+    max_cache = 1024 * 1024 * config.get("max_cache", 50)  # this is needed for partial loads
 
     tap_stream = io.TextIOWrapper(sys.stdin.buffer, encoding="utf-8")
 
@@ -87,6 +112,7 @@ def main():
             ph,
             tap_stream,
             initial_state=state,
+            state_handler=State if merge_state_messages else LiteralState,
             project_id=project_id,
             dataset=dataset,
             location=location,
diff --git a/target_bigquery/processhandler.py b/target_bigquery/processhandler.py
@@ -13,7 +13,7 @@
 
 from target_bigquery.encoders import DecimalEncoder
 from target_bigquery.schema import build_schema, cleanup_record, format_record_to_schema
-from target_bigquery.state import State
+
 from target_bigquery.simplify_json_schema import simplify
 from target_bigquery.validate_json_schema import validate_json_schema_completeness, \
     check_schema_for_dupes_in_field_names
@@ -131,7 +131,9 @@ def __init__(self, logger, **kwargs):
         # self.table_configs = kwargs.get("table_configs", {}) or {}
         #
         # self.INIT_STATE = kwargs.get("initial_state") or {}
-        self.STATE = State(**self.INIT_STATE)
+        # self.STATE = State(**self.INIT_STATE)
+        self.STATE_HANDLER = kwargs.get("state_handler")
+        self.STATE = self.STATE_HANDLER(**self.INIT_STATE)
 
         self.bq_schema_dicts = {}
         self.rows = {}
@@ -298,7 +300,7 @@ def _load_to_bq(self,
                 field=partition_field
             )
 
-        # clusteing
+        # clustering
         if cluster_fields:
             load_config.clustering_fields = cluster_fields
 
@@ -372,7 +374,8 @@ class BookmarksStatePartialLoadJobProcessHandler(PartialLoadJobProcessHandler):
     def __init__(self, logger, **kwargs):
         super(BookmarksStatePartialLoadJobProcessHandler, self).__init__(logger, **kwargs)
 
-        self.EMITTED_STATE = State(**self.INIT_STATE)
+        self.STATE_HANDLER = kwargs.get("state_handler")
+        self.EMITTED_STATE = self.STATE_HANDLER(**self.INIT_STATE)
 
     def handle_state_message(self, msg):
         assert isinstance(msg, singer.StateMessage)
diff --git a/target_bigquery/state.py b/target_bigquery/state.py
@@ -31,3 +31,12 @@ def merge(self, state):
 
             elif isinstance(m, Remove):
                 pass
+
+
+class LiteralState(State):
+
+    def merge(self, state):
+        if state:
+            self.clear()
+            self.update(state)
+
diff --git a/tests/test_state.py b/tests/test_state.py
@@ -15,7 +15,7 @@
 """
 
 from tests import unittestcore
-from target_bigquery.state import State
+from target_bigquery.state import State, LiteralState
 import os
 
 
@@ -45,6 +45,9 @@ def test_flat_schema(self):
         )
         print(state)
 
+        assert state == {'bookmarks': {'stream_one': {'timestamp': '2020-01-11T00:00:00.000000Z'},
+                                       'stream_two': {'timestamp': '2020-01-11T00:00:00.000000Z'}}}
+
     def test_state_but_no_data(self):
         from target_bigquery import main
 
@@ -82,3 +85,162 @@ def test_state_but_no_data_bookmarks_load(self):
 
         self.assertEqual(ret, 0, msg="Exit code is not 0!")
         self.assertDictEqual(state[-1], {"bookmarks": {"simple_stream": {"timestamp": "2020-01-11T00:00:00.000000Z"}}})
+
+    def test_state_facebook_stream(self):
+        from target_bigquery import main
+
+        self.set_cli_args(
+            stdin=os.path.join(os.path.join(
+                os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'tests'), 'rsc'),
+                'data'), 'facebook_stream.json'),
+            config=os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'sandbox'),
+                                'target-config.json'),
+            processhandler="load-job"
+        )
+
+        ret = main()
+        state = self.get_state()[-1]
+
+        self.assertEqual(ret, 0, msg="Exit code is not 0!")
+
+        self.assertEqual(state, {"bookmarks": {"ads": {"updated_time": "2020-07-24T13:03:56-05:00"},
+                                               "adsets": {"updated_time": "2020-07-23T16:16:54-05:00"},
+                                               "campaigns": {"updated_time": "2020-07-23T16:16:52-05:00"},
+                                               "ads_insights": {"date_start": "2020-07-24T00:00:00+00:00"}}}
+                         )
+
+    def test_state_facebook_stream_merge_state_cli_trumps_config(self):
+        from target_bigquery import main
+
+        self.set_cli_args(
+            stdin=os.path.join(os.path.join(
+                os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'tests'), 'rsc'),
+                'data'), 'facebook_stream.json'),
+            config=os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'sandbox'),
+                                'target_config_merge_state_false_flag.json'),
+            processhandler="load-job",
+            flag='--merge_state_messages'
+
+        )
+
+        ret = main()
+
+        state = self.get_state()[-1]
+
+        self.assertEqual(ret, 0, msg="Exit code is not 0!")
+
+        self.assertEqual(state, {"bookmarks": {"ads": {"updated_time": "2020-07-24T13:03:56-05:00"},
+                                               "adsets": {"updated_time": "2020-07-23T16:16:54-05:00"},
+                                               "campaigns": {"updated_time": "2020-07-23T16:16:52-05:00"},
+                                               "ads_insights": {"date_start": "2020-07-24T00:00:00+00:00"}}}
+                         )
+
+
+class TestSimpleStreamLiteralStateNoMerging(TestSimpleStream):
+
+    def test_init(self):
+        s = LiteralState(**{"a": 1})
+        print(s)
+
+    def test_flat_schema(self):
+        state = LiteralState()
+        state.merge(
+            {"bookmarks": {"stream_one": {"timestamp": "2020-01-10T00:00:00.000000Z"}}}
+        )
+        print(state)
+
+        state.merge(
+            {"bookmarks": {"stream_one": {"timestamp": "2020-01-11T00:00:00.000000Z"}}}
+        )
+        print(state)
+
+        state.merge(
+            {"bookmarks": {"stream_two": {"timestamp": "2020-01-11T00:00:00.000000Z"}}}
+        )
+        print(state)
+
+        assert state == {"bookmarks": {"stream_two": {"timestamp": "2020-01-11T00:00:00.000000Z"}}}
+
+    def test_state_but_no_data(self):
+        from target_bigquery import main
+
+        self.set_cli_args(
+            stdin=os.path.join(os.path.join(
+                os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'tests'), 'rsc'),
+                'partial_load_streams'), 'no_data_stream.json'),
+            config=os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'sandbox'),
+                                'target_config_cache.json'),
+            processhandler="partial-load-job",
+            flag='--no-merge_state_messages'
+        )
+
+        ret = main()
+        state = self.get_state()
+        self.assertEqual(1, len(state))  # only 1 state is expected: end state emit
+
+        self.assertEqual(ret, 0, msg="Exit code is not 0!")
+        self.assertDictEqual(state[-1], {"bookmarks": {"simple_stream": {"timestamp": "2020-01-11T00:00:00.000000Z"}}})
+
+    def test_state_but_no_data_bookmarks_load(self):
+        from target_bigquery import main
+
+        self.set_cli_args(
+            stdin=os.path.join(os.path.join(
+                os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'tests'), 'rsc'),
+                'partial_load_streams'), 'no_data_stream.json'),
+            config=os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'sandbox'),
+                                'target_config_cache.json'),
+            processhandler="bookmarks-partial-load-job",
+            flag='--no-merge_state_messages'
+        )
+
+        ret = main()
+        state = self.get_state()
+        self.assertEqual(1, len(state))  # only 1 state is expected: end state emit
+
+        self.assertEqual(ret, 0, msg="Exit code is not 0!")
+        self.assertDictEqual(state[-1], {"bookmarks": {"simple_stream": {"timestamp": "2020-01-11T00:00:00.000000Z"}}})
+
+    def test_state_facebook_stream(self):
+        from target_bigquery import main
+
+        self.set_cli_args(
+            stdin=os.path.join(os.path.join(
+                os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'tests'), 'rsc'),
+                'data'), 'facebook_stream.json'),
+            config=os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'sandbox'),
+                                'target-config.json'),
+            processhandler="load-job",
+            flag='--no-merge_state_messages'
+        )
+
+        ret = main()
+
+        state = self.get_state()[-1]
+
+        self.assertEqual(ret, 0, msg="Exit code is not 0!")
+
+        self.assertEqual(state, {"bookmarks": {"ads_insights": {"date_start": "2020-07-24T00:00:00+00:00"}}}
+                         )
+
+    def test_state_facebook_stream_merge_state_config(self):
+        from target_bigquery import main
+
+        self.set_cli_args(
+            stdin=os.path.join(os.path.join(
+                os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'tests'), 'rsc'),
+                'data'), 'facebook_stream.json'),
+            config=os.path.join(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'sandbox'),
+                                'target_config_merge_state_false_flag.json'),
+            processhandler="load-job"
+
+        )
+
+        ret = main()
+
+        state = self.get_state()[-1]
+
+        self.assertEqual(ret, 0, msg="Exit code is not 0!")
+
+        self.assertEqual(state, {"bookmarks": {"ads_insights": {"date_start": "2020-07-24T00:00:00+00:00"}}}
+                         )
diff --git a/tests/unittestcore.py b/tests/unittestcore.py
@@ -2,6 +2,7 @@
 import os
 import sys
 import unittest
+import re
 
 from google.cloud import bigquery
 
@@ -34,7 +35,12 @@ def setUp(self):
         os.environ["MALFORMED_TARGET_CONFIG"] = os.path.join(
             os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'sandbox'),
             'malformed_target_config.json')
-        #TODO: make naming convention of target config files consistent "_" vs "-". Use "_" as it's easier to copy with a click
+
+        os.environ["TARGET_CONFIG_MERGE_STATE_FALSE_FLAG"] = os.path.join(
+            os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'sandbox'),
+            'target_config_merge_state_false_flag.json')
+
+        # TODO: make naming convention of target config files consistent "_" vs "-". Use "_" as it's easier to copy with a click
         # I think we would just need to rename target-config.json to target_config.json (also update it in README)
         self.client = None
         self.project_id = None
@@ -46,11 +52,18 @@ def tearDown(self):
 
     def set_cli_args(self, ds_delete=True, *args, **kwargs):
         arg = [arg for arg in args]
+
         for k, v in kwargs.items():
             if k == "stdin":
                 sys.stdin = open(v, "r")
                 continue
 
+            # if some flag is being passed, such as --merge-state or --no-merge-state:
+            # we want to add this flag to CLI arguments
+            if k == "flag":
+                arg.append(v)
+                continue
+
             arg.append("--{}".format(k))
             arg.append("{}".format(v))
 
