Intermediate changes

commit_hash:4b0cad6620396d4b5422b264f4aa162104d45e8d

Author: robot-piglet

yql/essentials/udfs/common/python/bindings/py_cast.cpp

@@ -916,7 +916,19 @@ TPyObjectPtr ToPyArgs(
     TPyObjectPtr tuple(PyTuple_New(argsCount));
 
     for (ui32 i = 0; i < argsCount; i++) {
-        auto arg = ToPyObject(ctx, inspector.GetArgType(i), args[i]);
+        const auto argType = inspector.GetArgType(i);
+        auto arg = ToPyObject(ctx, argType, args[i]);
+        // PyTuple_SET_ITEM doesn't handle the case if nullptr is
+        // given as a payload to be set (unlike PyList_Append or
+        // PyDict_SetItem do), so we have to explicitly handle
+        // the failed export from UnboxedValue to PyObject here.
+        if (!arg) {
+            ::TStringBuilder sb;
+            sb << "Failed to export ";
+            NUdf::TTypePrinter(*ctx->PyCtx->TypeInfoHelper, argType).Out(sb.Out);
+            sb << " given as args[" << i << "]: ";
+            UdfTerminate((sb << ctx->PyCtx->Pos << GetLastErrorAsString()).data());
+        }
         PyTuple_SET_ITEM(tuple.Get(), i, arg.Release());
     }
 

yql/essentials/udfs/common/python/bindings/py_cast_ut.cpp

@@ -124,4 +124,29 @@ Y_UNIT_TEST_SUITE(TPyCastTest) {
     Y_UNIT_TEST(BadFromPythonJson) {
         TestBadUtf8Encode<NUdf::TJson>();
     }
+
+    Y_UNIT_TEST(BadToPythonJson) {
+        TPythonTestEngine engine;
+        UNIT_ASSERT_EXCEPTION_CONTAINS(
+            engine.UnsafeCall<void(NUdf::TJson)>(
+                [](const TType*, const NUdf::IValueBuilder& builder) {
+                    // XXX: The value below is built with the
+                    // following expression:
+                    // $query = "a=1&t%EDb=2";
+                    // $qdict = Url::QueryStringToDict($query);
+                    // $qyson = Yson::From($qdict);
+                    // $badJson = Yson::SerializeJson($qyson);
+                    //
+                    // For more info, see YQL-20231 and YQL-20220.
+                    constexpr TStringBuf badJson = "\x7b\x22\x61\x22\x3a\x5b\x22\x31\x22\x5d\x2c\x22\x74\xed\x62\x22\x3a\x5b\x22\x32\x22\x5d\x7d";
+                    return builder.NewString(badJson);
+                },
+                "def Test(arg):\n"
+                "   pass",
+                [](const NUdf::TUnboxedValuePod&) {
+                    Y_UNREACHABLE();
+                }
+            ),
+            yexception, "Failed to export Json given as args[0]");
+    }
 } // Y_UNIT_TEST_SUITE(TPyCastTest)

yql/essentials/udfs/common/python/bindings/py_test_engine.h

@@ -115,6 +115,32 @@ class TPythonTestEngine {
         ToMiniKQLWithArg<TChecker>(type, argValue, script, std::move(checker));
     }
 
+    template <typename FunctionType,
+              typename TMiniKQLValueBuilder,
+              typename TChecker>
+    void UnsafeCall(TMiniKQLValueBuilder&& builder,
+                    const TStringBuf& script,
+                    TChecker&& checker)
+    {
+        TPyObjectPtr function = CompilePythonFunction(script);
+        const auto functionType = GetTypeBuilder().SimpleSignatureType<FunctionType>();
+        NUdf::TCallableTypeInspector inspector(*CastCtx_->PyCtx->TypeInfoHelper, functionType);
+        Y_ENSURE(inspector.GetArgsCount() == 1);
+        const TType* argType = static_cast<const TType*>(inspector.GetArgType(0));
+        NUdf::TUnboxedValue value = builder(argType, GetValueBuilder());
+        TPyObjectPtr pyArgs = ToPyArgs(CastCtx_, functionType, &value, inspector);
+        TPyObjectPtr resultObj = PyObject_CallObject(function.Get(), pyArgs.Get());
+
+        if (!resultObj) {
+            return checker(NUdf::TUnboxedValuePod::Invalid());
+        }
+
+        const auto returnType = inspector.GetReturnType();
+        Y_ENSURE(CastCtx_->PyCtx->TypeInfoHelper->GetTypeKind(returnType) != NUdf::ETypeKind::Stream);
+
+        checker(FromPyObject(CastCtx_, returnType, resultObj.Get()));
+    }
+
     template <typename TMiniKQLValueBuilder>
     TPyObjectPtr ToPython(
                 NUdf::TType* udfType,
@@ -189,9 +215,7 @@ class TPythonTestEngine {
     }
 
 private:
-    TPyObjectPtr RunPythonFunction(
-            const TStringBuf& script, PyObject* args = nullptr)
-    {
+    TPyObjectPtr CompilePythonFunction(const TStringBuf& script) {
         TString filename(TStringBuf("embedded:test.py"));
         TPyObjectPtr code(Py_CompileString(script.data(), filename.data(), Py_file_input));
         if (!code) {
@@ -211,6 +235,13 @@ class TPythonTestEngine {
             PyErr_Print();
             UNIT_FAIL("function 'Test' is not found in module");
         }
+        return function;
+    }
+
+    TPyObjectPtr RunPythonFunction(
+            const TStringBuf& script, PyObject* args = nullptr)
+    {
+        TPyObjectPtr function(CompilePythonFunction(script));
         return PyObject_CallObject(function.Get(), args);
     }