diff --git a/src/dsdk/mongo.py b/src/dsdk/mongo.py index 8166300..d0ca934 100644 --- 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 index fbdef3f..413b2a6 100644 --- 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)