turbaszek commented on a change in pull request #10933:
URL: https://github.com/apache/airflow/pull/10933#discussion_r488016665



##########
File path: airflow/utils/orm_event_handlers.py
##########
@@ -52,6 +55,40 @@ def set_mysql_timezone(dbapi_connection, connection_record):
             cursor.execute("SET time_zone = '+00:00'")
             cursor.close()
 
+        is_mariadb = engine.dialect.server_version_info and "MariaDB" in 
engine.dialect.server_version_info
+        if idle_session_timeout > 0:
+            if is_mariadb:
+                # 
https://mariadb.com/kb/en/mariadb-1032-release-notes/#variables
+                if engine.dialect.server_version_info > (10, 3):
+                    setting = "idle_write_transaction_timeout"
+                else:
+                    setting = "idle_readwrite_transaction_timeout"
+
+                @event.listens_for(engine, "connect")
+                def set_mysql_transaction_idle_timeout(dbapi_connection, 
connection_record):
+                    cursor = dbapi_connection.cursor()
+                    cursor.execute("SET @@SESSION.%s = %s", (setting, 
int(idle_session_timeout),))
+                    cursor.close()
+            else:
+                # MySQL doesn't have the equivalent of postgres' 
idle_in_transaction_session_timeout -- the
+                # best we can do is set the time we wait just set the timeout 
for any idle connection
+                @event.listens_for(engine, "connect")
+                def set_mysql_transaction_idle_timeout(dbapi_connection, 
connection_record):
+                    cursor = dbapi_connection.cursor()
+                    cursor.execute("SET @@SESSION.wait_timeout = %s", 
(int(idle_session_timeout),))
+                    cursor.close()
+
+    if engine.dialect.name == "postgresql":
+        if idle_session_timeout > 0:

Review comment:
       Those ifs can be combined into one




----------------------------------------------------------------
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.

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


Reply via email to