Update logging

pennsignals · Jun 2, 2020 · fbebef3 · fbebef3
1 parent 7dc7fc8
commit fbebef3
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 21 deletions.
diff --git a/src/dsdk/mongo.py b/src/dsdk/mongo.py
@@ -253,13 +253,7 @@ def insert_many(collection: Collection, docs: Sequence[Dict[str, Any]]):
 UPDATE_ONE = "".join(
     (
         "{",
-        ", ".join(
-            (
-                '"key": "mongo.update_one"',
-                '"collection": "%s.%s"',
-                '"id": "%s"',
-            )
-        ),
+        ", ".join(('"key": "mongo.update_one"', '"collection": "%s.%s"',)),
         "}",
     )
 )
@@ -272,9 +266,6 @@ def update_one(
     """Update one with retry."""
     result = collection.update_one(key, doc)
     logger.info(
-        UPDATE_ONE,
-        collection.database.name,
-        collection.name,
-        result.inserted_id,
+        UPDATE_ONE, collection.database.name, collection.name,
     )
     return result
diff --git a/src/dsdk/mssql.py b/src/dsdk/mssql.py
@@ -73,21 +73,32 @@ def _inject_mssql_uri(mssql_uri: str) -> str:
             type=_inject_mssql_uri,
         )
 
+    OPEN = "".join(("{", ", ".join(('"key": "mssql.open"')), "}"))
+
+    CLOSE = "".join(("{", ", ".join(('"key": "mssql.close"')), "}"))
+
+    CONNECT = """
+select 1 as n
+"""
+
     @contextmanager
     def open_mssql(self) -> Generator:
         """Open mssql."""
         with self._mssql.connect() as con:
-            yield con
-            logger.info('"action": "connect"')
+            # force lazy connection open.
+            cur = con.execute(self.CONNECT)
+            for _ in cur.fetchall():
+                pass
+            logger.info(self.OPEN)
+            try:
+                yield con
+            finally:
+                logger.info(self.CLOSE)
 
 
 class CheckTablePrivileges(Task):  # pylint: disable=too-few-public-methods
     """Check table privileges."""
 
-    CONNECT = """
-select 1 as n
-"""
-
     EXTANT = """
 select 1 as n where exists (select 1 as n from {table})
 """
@@ -122,10 +133,6 @@ def __call__(self, batch, service):
         """__call__."""
         logger.info(self.ON)
         with service.open_mssql() as con:
-            # force lazy connection open.
-            cur = con.execute(self.CONNECT)
-            for _ in cur.fetchall():
-                pass
             errors = []
             for table in self.tables:
                 sql = self.EXTANT.format(table=table)