Use new bulk upload POST route to send introspection results up to gate. (#74)

Author: James Robinson

noteable_magics/sql/meta_commands.py

@@ -519,9 +519,6 @@ class IntrospectAndStoreDatabaseCommand(MetaCommand):
     # new datasource type supported.
     AVOID_SCHEMAS = set(('information_schema', 'pg_catalog', 'crdb_internal'))
 
-    # As expected in real kernels in Notable. Overridden in test suite.
-    JWT_PATHNAME = "/vault/secrets/jwt"
-
     def run(self, invoked_as: str, args: List[str]) -> None:
         """Drive introspecting whole database, POSTing results back to Gate for storage.
 
@@ -542,10 +539,6 @@ def run(self, invoked_as: str, args: List[str]) -> None:
 
         inspector = self.get_inspector()
 
-        # On successful completion, will tell Gate to delete any now-orphaned
-        # relations from prior introspections older than this timestamp.
-        introspection_started_at = datetime.utcnow()
-
         # This and delta() just for development timing figures. Could become yet another
         # timer context manager implementation.
         start = time.monotonic()
@@ -565,14 +558,11 @@ def delta() -> float:
         relations_and_kinds = self.all_table_and_views(inspector)
         print(f'Discovered {len(relations_and_kinds)} relations in {delta()}')
 
-        auth_header = self.get_auth_header()
-
-        # JSON-able schema descriptions to tell Gate about.
-        message_queue: List[RelationStructureDescription] = []
-
         # Introspect each relation concurrently.
         # TODO: Take minimum concurrency as a param?
-        with ThreadPoolExecutor(max_workers=self.MAX_INTROSPECTION_THREADS) as executor:
+        with ThreadPoolExecutor(
+            max_workers=self.MAX_INTROSPECTION_THREADS
+        ) as executor, RelationStructureMessager(ds_id) as messenger:
 
             future_to_relation = {
                 executor.submit(
@@ -584,40 +574,15 @@ def delta() -> float:
             for future in as_completed(future_to_relation):
                 schema_name, relation_name, kind = future_to_relation[future]
                 try:
-                    message_queue.append(future.result())
+                    messenger.queue_for_delivery(future.result())
                 except Exception as exc:
                     print(f'Exception for {schema_name}.{relation_name}: {exc}')
 
         table_introspection_delta = delta()
         print(
-            f'Done introspecting in {table_introspection_delta}, amortized {table_introspection_delta / len(relations_and_kinds)}s per relation'
+            f'Done introspecting and messaging gate in {table_introspection_delta}, amortized {table_introspection_delta / len(relations_and_kinds)}s per relation'
         )
 
-        # Now report each result back to gate, currently one at a time. Done completely
-        # after just for relative timing measurements.
-
-        # TODO: Could do the messaging back to gate either one-at-a-time in the worker threads,
-        # or in bulkier chunks either in the main thread when Gate offers a bulk API,
-        # or could use httpx and an event loop, ...
-
-        # But right now doing the messaging back to gate back here in the main thread just one
-        # request at a time.
-
-        session = requests.Session()
-        session.headers.update(auth_header)
-
-        if message_queue:
-            for message in message_queue:
-                self.inform_gate_relation(session, ds_id, message)
-
-            # Clear out any prior known relations which may not exist anymore in this datasource.
-            #
-            # We do this at the tail end of things, and not the beginning, so as to not eagerly delete
-            # prior known data if we happen to croak due to some unforseen exception while introspecting.
-            self.inform_gate_completed(session, ds_id, introspection_started_at)
-
-            print(f'Done storing discovered table and view structures in {delta()}')
-
         # run() contract: return what to bind to the SQL cell variable name, and if display() needs
         # to be called on it. Nothing and nope!
         return (None, False)
@@ -873,11 +838,93 @@ def get_datasource_id(self) -> UUID:
         handle = self.conn.name
         return UUID(handle[1:])
 
-    def get_auth_header(self):
-        """Set up the auth header dict for making requests directly to Gate"""
+
+class RelationStructureMessager:
+    """Context manager that collects the single-relation descriptions discovered
+    within IntrospectAndStoreDatabaseCommand, buffers up to CAPACITY at a time, then
+    POSTs the current collected group up to Gate in a single bulk POST (which accepts
+    at most 10).
+
+    Upon context exit, be sure to POST any partial remainder, then tell gate
+    we're all done introspecting, allowing it to delete any old now no longer
+    existing structures from prior introspections.
+
+    Helper class simplifying IntrospectAndStoreDatabaseCommand.run().
+    """
+
+    # Overridden in test suite.
+    CAPACITY = 10
+    # As expected in real kernels in Notable. Overridden in test suite.
+    JWT_PATHNAME = "/vault/secrets/jwt"
+
+    _relations: List[RelationStructureDescription]
+    _session: requests.Session
+    _datasource_id: UUID
+    _partition_counter: int
+    _started_at: datetime
+
+    def __init__(self, datasource_id: UUID):
+        self._datasource_id = datasource_id
+
+        # Set up HTTP session to message Gate about what is discovered.
+        self._session = requests.Session()
         jwt = pathlib.Path(self.JWT_PATHNAME).read_text()
+        self._session.headers.update({"Authorization": f"Bearer {jwt}"})
+
+    def __enter__(self):
+        self._relations = []
+        self._started_at = datetime.utcnow()
+        self._partition_counter = 1
+
+        return self
+
+    def queue_for_delivery(self, relation: RelationStructureDescription):
+        """Stash this discovered relation. If we have enough already
+        to send up to Gate, then do so.
+        """
+        self._relations.append(relation)
+
+        # Time for an intermediate flush?
+        if len(self._relations) == self.CAPACITY:
+            self._message_gate()
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self._message_gate(completed_introspection=True)
+
+    def _message_gate(self, completed_introspection: bool = False):
+
+        base_url = f"http://gate.default/api/v1/datasources/{self._datasource_id}/schema/relations"
+
+        if self._relations:
+            # Assemble buffered relation descriptions into a single bulk upload payload
+            jsonable_message = [
+                relation_description.dict() for relation_description in self._relations
+            ]
+
+            # Upload in a single message.
+            resp = self._session.post(
+                base_url,
+                json=jsonable_message,
+            )
+
+            if resp.status_code == 204:
+                for relation_description in self._relations:
+                    print(
+                        f'Stored structure of {relation_description.schema_name}.{relation_description.relation_name} in partition {self._partition_counter}'
+                    )
+            else:
+                print(
+                    f'Failed storing partition {self._partition_counter}: {resp.status_code}, {resp.text}'
+                )
+
+            # Prepare for next partition.
+            self._partition_counter += 1
+            self._relations = []
 
-        return {"Authorization": f"Bearer {jwt}"}
+        if completed_introspection:
+            # Message indicating all done through asking to clear out any stored relation structures
+            # older than when we started.
+            self._session.delete(f"{base_url}?older_than={self._started_at.isoformat()}")
 
 
 def constraints_dataframe(

tests/test_sql_magic_meta_commands.py

@@ -12,7 +12,7 @@
 from noteable_magics.sql.connection import Connection
 from noteable_magics.sql.gate_messaging_types import RelationStructureDescription
 from noteable_magics.sql.meta_commands import (
-    IntrospectAndStoreDatabaseCommand,
+    RelationStructureMessager,
     SchemaStrippingInspector,
     _all_command_classes,
     convert_relation_glob_to_regex,
@@ -586,24 +586,31 @@ def test_nonexistent_schema(self, sql_magic, capsys):
 @pytest.mark.usefixtures("populated_cockroach_database")
 class TestFullIntrospection:
     @pytest.fixture()
-    def patched_jwt(self, tmp_path):
+    def patched_relation_structure_messager(self, tmp_path):
 
-        original_jwt_pathname = IntrospectAndStoreDatabaseCommand.JWT_PATHNAME
+        original_jwt_pathname = RelationStructureMessager.JWT_PATHNAME
         tmp_jwt_path = tmp_path / 'kernel_jwt'
         tmp_jwt_path.write_text('kerneljwtcontents')
 
-        IntrospectAndStoreDatabaseCommand.JWT_PATHNAME = str(tmp_jwt_path)
+        RelationStructureMessager.JWT_PATHNAME = str(tmp_jwt_path)
+
+        # Also make it only buffer at most 2 relations so we'll get some intermediate flushes
+        # when introspecting more than 2 relations.
+
+        original_capacity = RelationStructureMessager.CAPACITY
+
+        RelationStructureMessager.CAPACITY = 2
 
         yield
 
-        IntrospectAndStoreDatabaseCommand.JWT_PATHNAME = original_jwt_pathname
+        RelationStructureMessager.JWT_PATHNAME = original_jwt_pathname
+        RelationStructureMessager.CAPACITY = original_capacity
 
     @pytest.fixture()
-    def patched_requests_mock(self, mocker, patched_jwt, requests_mock):
-        # IntrospectAndStoreDatabaseCommand.inform_gate_relation() gonna try to do a POST
-        # to this URL.
+    def patched_requests_mock(self, mocker, patched_relation_structure_messager, requests_mock):
+        # RelationStructureMessager will do POSTs to this URL.
         requests_mock.post(
-            f"http://gate.default/api/v1/datasources/{COCKROACH_UUID}/schema/relation",
+            f"http://gate.default/api/v1/datasources/{COCKROACH_UUID}/schema/relations",
             status_code=204,
         )
 
@@ -617,27 +624,32 @@ def patched_requests_mock(self, mocker, patched_jwt, requests_mock):
 
     def test_full_introspection(self, sql_magic, capsys, patched_requests_mock):
 
+        # Create one more table, for a total of 5, so that exiting the scope of
+        # the RelationStructureMessager will have a dreg to POST.
+        # (see RelationStructureMessager.__exit__())
+
+        sql_magic.execute(
+            rf'{COCKROACH_HANDLE} create table dregs_table (id int primary key not null, name text)'
+        )
+
+        # Now introspect the whole DB.
         introspection_start = datetime.utcnow()
         sql_magic.execute(rf'{COCKROACH_HANDLE} \introspect')
         introspection_end = datetime.utcnow()
 
         out, err = capsys.readouterr()
 
         assert 'Exception' not in out
-        assert out.startswith('Discovered 4 relations')
+
+        # Expect the core tables plus our dregs_table.
+        # (The 'Done.\n' comes from the 'create table' execution creating dregs_table.)
+        assert out.startswith(f'Done.\nDiscovered {len(KNOWN_TABLES) + 1} relations')
 
         for name, kind in KNOWN_TABLES_AND_KINDS:
             assert f'Introspected {kind} public.{name}' in out
             assert f'Stored structure of public.{name}' in out
 
-        assert 'Done storing discovered table and view structures' in out
-
-        # One POST per relation discovered, and then a single DELETE (as currently implemented)
-        assert len(patched_requests_mock.request_history) == len(KNOWN_TABLES_AND_KINDS) + 1
-        # Remember, booleans are also integers because filthy historical C.
-        assert sum(req.method == 'POST' for req in patched_requests_mock.request_history) == len(
-            KNOWN_TABLES_AND_KINDS
-        )
+        assert 'Done introspecting and messaging gate' in out
 
         # The delete should be the last request, and should specify a since_when timestamp param
         # which should be between when introspection started and ended (in UTC).
@@ -659,38 +671,52 @@ def test_full_introspection(self, sql_magic, capsys, patched_requests_mock):
         had_check_constraints = False
         had_foreign_keys = False
 
+        post_call_count = 0
         for req in patched_requests_mock.request_history:
-            if req.method == 'POST' and req.url.endswith('/schema/relation'):
-
-                # POST body should correspond to RelationStructureDescription describing a single relation.
-                from_json = RelationStructureDescription(**req.json())
-                described_relation_names.add(from_json.relation_name)
-
-                if from_json.indexes:
-                    had_indexes = True
-
-                if from_json.primary_key_name and from_json.primary_key_columns:
-                    had_primary_key_columns = True
-                else:
-                    #
-                    assert (
-                        from_json.primary_key_name is None and from_json.primary_key_columns == []
-                    )
-
-                if from_json.columns:
-                    had_columns = True
-
-                if from_json.unique_constraints:
-                    had_unique_constraints = True
-
-                if from_json.check_constraints:
-                    had_check_constraints = True
-
-                if from_json.foreign_keys:
-                    had_foreign_keys = True
-
-        assert described_relation_names == set(KNOWN_TABLES)
-
+            if req.method == 'POST' and req.url.endswith('/schema/relations'):
+
+                post_call_count += 1
+                # POST body should correspond to a list of RelationStructureDescription each describing a single relation.
+                dict_list = req.json()
+                for member_dict in dict_list:
+                    from_json = RelationStructureDescription(**member_dict)
+                    described_relation_names.add(from_json.relation_name)
+
+                    if from_json.indexes:
+                        had_indexes = True
+
+                    if from_json.primary_key_name and from_json.primary_key_columns:
+                        had_primary_key_columns = True
+                    else:
+                        #
+                        assert (
+                            from_json.primary_key_name is None
+                            and from_json.primary_key_columns == []
+                        )
+
+                    if from_json.columns:
+                        had_columns = True
+
+                    if from_json.unique_constraints:
+                        had_unique_constraints = True
+
+                    if from_json.check_constraints:
+                        had_check_constraints = True
+
+                    if from_json.foreign_keys:
+                        had_foreign_keys = True
+
+        # Expected 3 POST calls for the 5 tables, since we fixture
+        # tuned RelationStructureMessager.CAPACITY down to 2. Two calls of two
+        # relations each, then a trailing dreg call.
+        assert post_call_count == 3
+
+        # Every expected relation should have been described.
+        expected_relation_names = set(KNOWN_TABLES)  # The default table set
+        expected_relation_names.add('dregs_table')  # plus our per-test extra.
+        assert described_relation_names == expected_relation_names
+
+        # Every substructure variation should have been covered across all these tables/view.
         assert had_columns
         assert had_primary_key_columns
         assert had_indexes