[skywalking-python] branch master updated: Fix sw_psycopg2 register_type() (#211)

[yihaochen]

This is an automated email from the ASF dual-hosted git repository.

yihaochen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/skywalking-python.git


The following commit(s) were added to refs/heads/master by this push:
     new bf0db59  Fix sw_psycopg2 register_type() (#211)
bf0db59 is described below

commit bf0db59f6cfe804fdc0195079694da0e5b693c99
Author: Tomasz Pytel <tompy...@gmail.com>
AuthorDate: Fri Jun 3 15:13:03 2022 -0300

    Fix sw_psycopg2 register_type() (#211)
---
 setup.cfg                         |  2 +-
 skywalking/plugins/sw_psycopg2.py | 50 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/setup.cfg b/setup.cfg
index 0e12735..285feaa 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -30,7 +30,7 @@ extend-ignore =
 
 
 max-line-length = 120
-max-complexity = 20
+max-complexity = 32
 exclude = venv*,*egg_info,skywalking/protocol
 count = True
 show-source = True
diff --git a/skywalking/plugins/sw_psycopg2.py 
b/skywalking/plugins/sw_psycopg2.py
index 621cb5a..7d7c80c 100644
--- a/skywalking/plugins/sw_psycopg2.py
+++ b/skywalking/plugins/sw_psycopg2.py
@@ -128,3 +128,53 @@ def install():
 
     _connect = psycopg2.connect
     psycopg2.connect = connect
+
+    try:  # try to instrument register_type which will fail if it gets a 
wrapped cursor or connection
+        from psycopg2._psycopg import register_type as _register_type
+
+        def register_type(c, conn_or_curs):
+            if isinstance(conn_or_curs, ProxyConnection):
+                conn_or_curs = conn_or_curs._self_conn
+            elif isinstance(conn_or_curs, ProxyCursor):
+                conn_or_curs = conn_or_curs._self_cur
+
+            return _register_type(c, conn_or_curs)
+
+        try:
+            import psycopg2._ipaddress
+
+            if psycopg2._ipaddress.register_type is _register_type:
+                psycopg2._ipaddress.register_type = register_type
+
+        except Exception:
+            pass
+
+        try:
+            import psycopg2._ipaddress
+
+            if psycopg2._json.register_type is _register_type:
+                psycopg2._json.register_type = register_type
+
+        except Exception:
+            pass
+
+        try:
+            import psycopg2._ipaddress
+
+            if psycopg2._range.register_type is _register_type:
+                psycopg2._range.register_type = register_type
+
+        except Exception:
+            pass
+
+        try:
+            import psycopg2._ipaddress
+
+            if psycopg2.extensions.register_type is _register_type:
+                psycopg2.extensions.register_type = register_type
+
+        except Exception:
+            pass
+
+    except Exception:
+        pass