Skip to content

Commit

Permalink
Fix faulty executor config serialization logic (apache#26191)
Browse files Browse the repository at this point in the history
The bind processor logic had the effect of applying serialization logic multiple times, which produced an incorrect serialization output.

Resolves apache#26101.

(cherry picked from commit 87108d7)
  • Loading branch information
dstandish authored and jedcunningham committed Oct 4, 2022
1 parent 51bcc11 commit 6831e78
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 12 deletions.
9 changes: 5 additions & 4 deletions airflow/utils/sqlalchemy.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.

import copy
import datetime
import json
import logging
Expand Down Expand Up @@ -164,9 +164,10 @@ def bind_processor(self, dialect):
super_process = super().bind_processor(dialect)

def process(value):
if isinstance(value, dict) and 'pod_override' in value:
value['pod_override'] = BaseSerialization()._serialize(value['pod_override'])
return super_process(value)
val_copy = copy.copy(value)
if isinstance(val_copy, dict) and 'pod_override' in val_copy:
val_copy['pod_override'] = BaseSerialization._serialize(val_copy['pod_override'])
return super_process(val_copy)

return process

Expand Down
24 changes: 16 additions & 8 deletions tests/utils/test_sqlalchemy.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
import datetime
import pickle
import unittest
from copy import copy
from unittest import mock
from unittest.mock import MagicMock

Expand All @@ -31,7 +30,7 @@

from airflow import settings
from airflow.models import DAG
from airflow.serialization.enums import Encoding
from airflow.serialization.enums import DagAttributeTypes, Encoding
from airflow.serialization.serialized_objects import BaseSerialization
from airflow.settings import Session
from airflow.utils.sqlalchemy import ExecutorConfigType, nowait, prohibit_commit, skip_locked, with_row_locks
Expand Down Expand Up @@ -239,10 +238,21 @@ def tearDown(self):

class TestExecutorConfigType:
@pytest.mark.parametrize(
'input',
['anything', {'pod_override': TEST_POD}],
'input, expected',
[
('anything', 'anything'),
(
{'pod_override': TEST_POD},
{
"pod_override": {
"__var": {"spec": {"containers": [{"name": "base"}]}},
"__type": DagAttributeTypes.POD,
}
},
),
],
)
def test_bind_processor(self, input):
def test_bind_processor(self, input, expected):
"""
The returned bind processor should pickle the object as is, unless it is a dictionary with
a pod_override node, in which case it should run it through BaseSerialization.
Expand All @@ -251,10 +261,8 @@ def test_bind_processor(self, input):
mock_dialect = MagicMock()
mock_dialect.dbapi = None
process = config_type.bind_processor(mock_dialect)
expected = copy(input)
if 'pod_override' in input:
expected['pod_override'] = BaseSerialization()._serialize(input['pod_override'])
assert pickle.loads(process(input)) == expected
assert pickle.loads(process(input)) == expected, "should should not mutate variable"

@pytest.mark.parametrize(
'input',
Expand Down

0 comments on commit 6831e78

Please sign in to comment.