dstandish commented on code in PR #59183:
URL: https://github.com/apache/airflow/pull/59183#discussion_r2627464589


##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -58,6 +59,63 @@
 log = structlog.get_logger(__name__)
 
 
+@contextmanager
+def _lock_asset_model(
+    *,
+    session: Session,
+    asset_id: int,
+    max_retries: int = 10,
+    retry_delay: float = 0.1,
+):
+    """
+    Context manager to acquire a lock for AssetPartitionDagRun creation.
+
+    - SQLite: Use a no-op ORM update to trigger a write-transaction and 
acquire SQLite's global writer lock.
+    - Postgres/MySQL: uses row-level lock on AssetModel.
+    """
+    if get_dialect_name(session) == "sqlite":
+        import time
+
+        from sqlalchemy import update
+
+        # no-op update
+        # This is used to acquire SQLite's global writer lock.
+        stmt = update(AssetModel).where(AssetModel.id == 
asset_id).values(id=AssetModel.id)
+        for _ in range(max_retries):
+            try:
+                session.execute(stmt)
+                session.flush()
+            except exc.OperationalError as err:
+                err_msg = str(err).lower()
+                if "locked" in err_msg or "busy" in err_msg:
+                    session.rollback()
+                    time.sleep(retry_delay)
+                    continue
+
+            try:
+                # lock acquired
+                yield
+            finally:
+                session.flush()

Review Comment:
   i am not a sqlalchemy expert but i am not sure why flushing when catching an 
error.  if there's an error what happens to the transaction?  is it ever 
committed?  if something else commits, there's no point to do a flush here.  if 
something else doesn't commit, the same is true.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to