12
12
# See the License for the specific language governing permissions and
13
13
# limitations under the License.
14
14
15
+ import psycopg
15
16
import psycopg2
16
17
from test_psycopg_functional import (
17
18
POSTGRES_DB_NAME ,
21
22
POSTGRES_USER ,
22
23
)
23
24
25
+ from opentelemetry .instrumentation .psycopg import PsycopgInstrumentor
24
26
from opentelemetry .instrumentation .psycopg2 import Psycopg2Instrumentor
25
27
from opentelemetry .test .test_base import TestBase
26
28
27
29
28
- class TestFunctionalPsycopg (TestBase ):
30
+ class TestFunctionalPsycopg2 (TestBase ):
29
31
def setUp (self ):
30
32
super ().setUp ()
31
33
self ._tracer = self .tracer_provider .get_tracer (__name__ )
@@ -52,3 +54,49 @@ def test_commenter_enabled(self):
52
54
self ._cursor .query .decode ("ascii" ),
53
55
r"SELECT 1 /\*db_driver='psycopg2(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),libpq_version=\d*,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;" ,
54
56
)
57
+
58
+
59
+ class TestFunctionalPsycopg (TestBase ):
60
+ def setUp (self ):
61
+ super ().setUp ()
62
+ PsycopgInstrumentor ().instrument (enable_commenter = True )
63
+
64
+ def tearDown (self ):
65
+ PsycopgInstrumentor ().uninstrument ()
66
+ super ().tearDown ()
67
+
68
+ def test_commenter_enabled_root_module (self ):
69
+ connection = psycopg .connect (
70
+ dbname = POSTGRES_DB_NAME ,
71
+ user = POSTGRES_USER ,
72
+ password = POSTGRES_PASSWORD ,
73
+ host = POSTGRES_HOST ,
74
+ port = POSTGRES_PORT ,
75
+ )
76
+ cursor = connection .cursor ()
77
+ cursor .execute ("SELECT 1;" )
78
+ self .assertRegex (
79
+ cursor ._query .query .decode ("ascii" ),
80
+ r"SELECT 1 /\*db_driver='psycopg(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),libpq_version=\d*,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;" ,
81
+ )
82
+
83
+ cursor .close ()
84
+ connection .close ()
85
+
86
+ def test_commenter_enabled_not_root_module (self ):
87
+ connection = psycopg .Connection .connect (
88
+ dbname = POSTGRES_DB_NAME ,
89
+ user = POSTGRES_USER ,
90
+ password = POSTGRES_PASSWORD ,
91
+ host = POSTGRES_HOST ,
92
+ port = POSTGRES_PORT ,
93
+ )
94
+ cursor = connection .cursor ()
95
+ cursor .execute ("SELECT 1;" )
96
+ self .assertRegex (
97
+ cursor ._query .query .decode ("ascii" ),
98
+ r"SELECT 1 /\*db_driver='psycopg(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;" ,
99
+ )
100
+
101
+ cursor .close ()
102
+ connection .close ()
0 commit comments