We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Instrument SQLAlchemy engine connection phase
** Is your feature request related to a problem? **
The SQLAlchemy instrumentation does not trace the actual connection to the database

**Describe the solution you'd like**
I want that `connect` function will also be traced

**Describe alternatives you've considered**
Which alternative solutions or features have you considered?

**Additional context**
We are working with SQLAlchemy (snowflake db) and we implemented a solution where we can see the `connect` span also as attached in the screenshot (the span is `database-connect`)
![image](https://user-images.githubusercontent.com/12069200/173789203-b75f42d2-7ae6-4c11-aa0c-4937ede75422.png)




--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py`
Content:
```
1 # Copyright The OpenTelemetry Authors
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """
16 Instrument `sqlalchemy`_ to report SQL queries.
17 
18 There are two options for instrumenting code. The first option is to use
19 the ``opentelemetry-instrument`` executable which will automatically
20 instrument your SQLAlchemy engine. The second is to programmatically enable
21 instrumentation via the following code:
22 
23 .. _sqlalchemy: https://pypi.org/project/sqlalchemy/
24 
25 Usage
26 -----
27 .. code:: python
28 
29     from sqlalchemy import create_engine
30 
31     from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
32     import sqlalchemy
33 
34     engine = create_engine("sqlite:///:memory:")
35     SQLAlchemyInstrumentor().instrument(
36         engine=engine,
37     )
38 
39     # of the async variant of SQLAlchemy
40 
41     from sqlalchemy.ext.asyncio import create_async_engine
42 
43     from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
44     import sqlalchemy
45 
46     engine = create_async_engine("sqlite:///:memory:")
47     SQLAlchemyInstrumentor().instrument(
48         engine=engine.sync_engine
49     )
50 
51 API
52 ---
53 """
54 from collections.abc import Sequence
55 from typing import Collection
56 
57 import sqlalchemy
58 from packaging.version import parse as parse_version
59 from wrapt import wrap_function_wrapper as _w
60 
61 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
62 from opentelemetry.instrumentation.sqlalchemy.engine import (
63     EngineTracer,
64     _get_tracer,
65     _wrap_create_async_engine,
66     _wrap_create_engine,
67 )
68 from opentelemetry.instrumentation.sqlalchemy.package import _instruments
69 from opentelemetry.instrumentation.utils import unwrap
70 
71 
72 class SQLAlchemyInstrumentor(BaseInstrumentor):
73     """An instrumentor for SQLAlchemy
74     See `BaseInstrumentor`
75     """
76 
77     def instrumentation_dependencies(self) -> Collection[str]:
78         return _instruments
79 
80     def _instrument(self, **kwargs):
81         """Instruments SQLAlchemy engine creation methods and the engine
82         if passed as an argument.
83 
84         Args:
85             **kwargs: Optional arguments
86                 ``engine``: a SQLAlchemy engine instance
87                 ``engines``: a list of SQLAlchemy engine instances
88                 ``tracer_provider``: a TracerProvider, defaults to global
89 
90         Returns:
91             An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.
92         """
93         tracer_provider = kwargs.get("tracer_provider")
94         _w("sqlalchemy", "create_engine", _wrap_create_engine(tracer_provider))
95         _w(
96             "sqlalchemy.engine",
97             "create_engine",
98             _wrap_create_engine(tracer_provider),
99         )
100         if parse_version(sqlalchemy.__version__).release >= (1, 4):
101             _w(
102                 "sqlalchemy.ext.asyncio",
103                 "create_async_engine",
104                 _wrap_create_async_engine(tracer_provider),
105             )
106 
107         if kwargs.get("engine") is not None:
108             return EngineTracer(
109                 _get_tracer(tracer_provider),
110                 kwargs.get("engine"),
111                 kwargs.get("enable_commenter", False),
112             )
113         if kwargs.get("engines") is not None and isinstance(
114             kwargs.get("engines"), Sequence
115         ):
116             return [
117                 EngineTracer(
118                     _get_tracer(tracer_provider),
119                     engine,
120                     kwargs.get("enable_commenter", False),
121                 )
122                 for engine in kwargs.get("engines")
123             ]
124 
125         return None
126 
127     def _uninstrument(self, **kwargs):
128         unwrap(sqlalchemy, "create_engine")
129         unwrap(sqlalchemy.engine, "create_engine")
130         if parse_version(sqlalchemy.__version__).release >= (1, 4):
131             unwrap(sqlalchemy.ext.asyncio, "create_async_engine")
132
```

Path: `instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py`
Content:
```
1 # Copyright The OpenTelemetry Authors
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 import os
15 
16 from sqlalchemy.event import listen  # pylint: disable=no-name-in-module
17 
18 from opentelemetry import trace
19 from opentelemetry.instrumentation.sqlalchemy.package import (
20     _instrumenting_module_name,
21 )
22 from opentelemetry.instrumentation.sqlalchemy.version import __version__
23 from opentelemetry.instrumentation.utils import (
24     _generate_opentelemetry_traceparent,
25     _generate_sql_comment,
26 )
27 from opentelemetry.semconv.trace import NetTransportValues, SpanAttributes
28 from opentelemetry.trace import Span
29 from opentelemetry.trace.status import Status, StatusCode
30 
31 
32 def _normalize_vendor(vendor):
33     """Return a canonical name for a type of database."""
34     if not vendor:
35         return "db"  # should this ever happen?
36 
37     if "sqlite" in vendor:
38         return "sqlite"
39 
40     if "postgres" in vendor or vendor == "psycopg2":
41         return "postgresql"
42 
43     return vendor
44 
45 
46 def _get_tracer(tracer_provider=None):
47     return trace.get_tracer(
48         _instrumenting_module_name,
49         __version__,
50         tracer_provider=tracer_provider,
51     )
52 
53 
54 def _wrap_create_async_engine(tracer_provider=None):
55     # pylint: disable=unused-argument
56     def _wrap_create_async_engine_internal(func, module, args, kwargs):
57         """Trace the SQLAlchemy engine, creating an `EngineTracer`
58         object that will listen to SQLAlchemy events.
59         """
60         engine = func(*args, **kwargs)
61         EngineTracer(_get_tracer(tracer_provider), engine.sync_engine)
62         return engine
63 
64     return _wrap_create_async_engine_internal
65 
66 
67 def _wrap_create_engine(tracer_provider=None):
68     # pylint: disable=unused-argument
69     def _wrap_create_engine_internal(func, module, args, kwargs):
70         """Trace the SQLAlchemy engine, creating an `EngineTracer`
71         object that will listen to SQLAlchemy events.
72         """
73         engine = func(*args, **kwargs)
74         EngineTracer(_get_tracer(tracer_provider), engine)
75         return engine
76 
77     return _wrap_create_engine_internal
78 
79 
80 class EngineTracer:
81     def __init__(self, tracer, engine, enable_commenter=False):
82         self.tracer = tracer
83         self.engine = engine
84         self.vendor = _normalize_vendor(engine.name)
85         self.enable_commenter = enable_commenter
86 
87         listen(
88             engine, "before_cursor_execute", self._before_cur_exec, retval=True
89         )
90         listen(engine, "after_cursor_execute", _after_cur_exec)
91         listen(engine, "handle_error", _handle_error)
92 
93     def _operation_name(self, db_name, statement):
94         parts = []
95         if isinstance(statement, str):
96             # otel spec recommends against parsing SQL queries. We are not trying to parse SQL
97             # but simply truncating the statement to the first word. This covers probably >95%
98             # use cases and uses the SQL statement in span name correctly as per the spec.
99             # For some very special cases it might not record the correct statement if the SQL
100             # dialect is too weird but in any case it shouldn't break anything.
101             parts.append(statement.split()[0])
102         if db_name:
103             parts.append(db_name)
104         if not parts:
105             return self.vendor
106         return " ".join(parts)
107 
108     # pylint: disable=unused-argument
109     def _before_cur_exec(
110         self, conn, cursor, statement, params, context, executemany
111     ):
112         attrs, found = _get_attributes_from_url(conn.engine.url)
113         if not found:
114             attrs = _get_attributes_from_cursor(self.vendor, cursor, attrs)
115 
116         db_name = attrs.get(SpanAttributes.DB_NAME, "")
117         span = self.tracer.start_span(
118             self._operation_name(db_name, statement),
119             kind=trace.SpanKind.CLIENT,
120         )
121         with trace.use_span(span, end_on_exit=False):
122             if span.is_recording():
123                 span.set_attribute(SpanAttributes.DB_STATEMENT, statement)
124                 span.set_attribute(SpanAttributes.DB_SYSTEM, self.vendor)
125                 for key, value in attrs.items():
126                     span.set_attribute(key, value)
127 
128         context._otel_span = span
129         if self.enable_commenter:
130             statement = statement + EngineTracer._generate_comment(span=span)
131 
132         return statement, params
133 
134     @staticmethod
135     def _generate_comment(span: Span) -> str:
136         span_context = span.get_span_context()
137         meta = {}
138         if span_context.is_valid:
139             meta.update(_generate_opentelemetry_traceparent(span))
140         return _generate_sql_comment(**meta)
141 
142 
143 # pylint: disable=unused-argument
144 def _after_cur_exec(conn, cursor, statement, params, context, executemany):
145     span = getattr(context, "_otel_span", None)
146     if span is None:
147         return
148 
149     span.end()
150 
151 
152 def _handle_error(context):
153     span = getattr(context.execution_context, "_otel_span", None)
154     if span is None:
155         return
156 
157     if span.is_recording():
158         span.set_status(
159             Status(
160                 StatusCode.ERROR,
161                 str(context.original_exception),
162             )
163         )
164     span.end()
165 
166 
167 def _get_attributes_from_url(url):
168     """Set connection tags from the url. return true if successful."""
169     attrs = {}
170     if url.host:
171         attrs[SpanAttributes.NET_PEER_NAME] = url.host
172     if url.port:
173         attrs[SpanAttributes.NET_PEER_PORT] = url.port
174     if url.database:
175         attrs[SpanAttributes.DB_NAME] = url.database
176     if url.username:
177         attrs[SpanAttributes.DB_USER] = url.username
178     return attrs, bool(url.host)
179 
180 
181 def _get_attributes_from_cursor(vendor, cursor, attrs):
182     """Attempt to set db connection attributes by introspecting the cursor."""
183     if vendor == "postgresql":
184         info = getattr(getattr(cursor, "connection", None), "info", None)
185         if not info:
186             return attrs
187 
188         attrs[SpanAttributes.DB_NAME] = info.dbname
189         is_unix_socket = info.host and info.host.startswith("/")
190 
191         if is_unix_socket:
192             attrs[SpanAttributes.NET_TRANSPORT] = NetTransportValues.UNIX.value
193             if info.port:
194                 # postgresql enforces this pattern on all socket names
195                 attrs[SpanAttributes.NET_PEER_NAME] = os.path.join(
196                     info.host, f".s.PGSQL.{info.port}"
197                 )
198         else:
199             attrs[
200                 SpanAttributes.NET_TRANSPORT
201             ] = NetTransportValues.IP_TCP.value
202             attrs[SpanAttributes.NET_PEER_NAME] = info.host
203             if info.port:
204                 attrs[SpanAttributes.NET_PEER_PORT] = int(info.port)
205     return attrs
206
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py
--- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py
@@ -56,12 +56,14 @@
 
 import sqlalchemy
 from packaging.version import parse as parse_version
+from sqlalchemy.engine.base import Engine
 from wrapt import wrap_function_wrapper as _w
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.sqlalchemy.engine import (
     EngineTracer,
     _get_tracer,
+    _wrap_connect,
     _wrap_create_async_engine,
     _wrap_create_engine,
 )
@@ -97,13 +99,17 @@
             "create_engine",
             _wrap_create_engine(tracer_provider),
         )
+        _w(
+            "sqlalchemy.engine.base",
+            "Engine.connect",
+            _wrap_connect(tracer_provider),
+        )
         if parse_version(sqlalchemy.__version__).release >= (1, 4):
             _w(
                 "sqlalchemy.ext.asyncio",
                 "create_async_engine",
                 _wrap_create_async_engine(tracer_provider),
             )
-
         if kwargs.get("engine") is not None:
             return EngineTracer(
                 _get_tracer(tracer_provider),
@@ -127,5 +133,6 @@
     def _uninstrument(self, **kwargs):
         unwrap(sqlalchemy, "create_engine")
         unwrap(sqlalchemy.engine, "create_engine")
+        unwrap(Engine, "connect")
         if parse_version(sqlalchemy.__version__).release >= (1, 4):
             unwrap(sqlalchemy.ext.asyncio, "create_async_engine")
diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py
--- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py
+++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py
@@ -77,6 +77,23 @@
     return _wrap_create_engine_internal
 
 
+def _wrap_connect(tracer_provider=None):
+    tracer = trace.get_tracer(
+        _instrumenting_module_name,
+        __version__,
+        tracer_provider=tracer_provider,
+    )
+
+    # pylint: disable=unused-argument
+    def _wrap_connect_internal(func, module, args, kwargs):
+        with tracer.start_as_current_span(
+            "connect", kind=trace.SpanKind.CLIENT
+        ):
+            return func(*args, **kwargs)
+
+    return _wrap_connect_internal
+
+
 class EngineTracer:
     def __init__(self, tracer, engine, enable_commenter=False):
         self.tracer = tracer

{"golden_diff": "diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py\n--- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py\n+++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py\n@@ -56,12 +56,14 @@\n \n import sqlalchemy\n from packaging.version import parse as parse_version\n+from sqlalchemy.engine.base import Engine\n from wrapt import wrap_function_wrapper as _w\n \n from opentelemetry.instrumentation.instrumentor import BaseInstrumentor\n from opentelemetry.instrumentation.sqlalchemy.engine import (\n     EngineTracer,\n     _get_tracer,\n+    _wrap_connect,\n     _wrap_create_async_engine,\n     _wrap_create_engine,\n )\n@@ -97,13 +99,17 @@\n             \"create_engine\",\n             _wrap_create_engine(tracer_provider),\n         )\n+        _w(\n+            \"sqlalchemy.engine.base\",\n+            \"Engine.connect\",\n+            _wrap_connect(tracer_provider),\n+        )\n         if parse_version(sqlalchemy.__version__).release >= (1, 4):\n             _w(\n                 \"sqlalchemy.ext.asyncio\",\n                 \"create_async_engine\",\n                 _wrap_create_async_engine(tracer_provider),\n             )\n-\n         if kwargs.get(\"engine\") is not None:\n             return EngineTracer(\n                 _get_tracer(tracer_provider),\n@@ -127,5 +133,6 @@\n     def _uninstrument(self, **kwargs):\n         unwrap(sqlalchemy, \"create_engine\")\n         unwrap(sqlalchemy.engine, \"create_engine\")\n+        unwrap(Engine, \"connect\")\n         if parse_version(sqlalchemy.__version__).release >= (1, 4):\n             unwrap(sqlalchemy.ext.asyncio, \"create_async_engine\")\ndiff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py\n--- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py\n+++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py\n@@ -77,6 +77,23 @@\n     return _wrap_create_engine_internal\n \n \n+def _wrap_connect(tracer_provider=None):\n+    tracer = trace.get_tracer(\n+        _instrumenting_module_name,\n+        __version__,\n+        tracer_provider=tracer_provider,\n+    )\n+\n+    # pylint: disable=unused-argument\n+    def _wrap_connect_internal(func, module, args, kwargs):\n+        with tracer.start_as_current_span(\n+            \"connect\", kind=trace.SpanKind.CLIENT\n+        ):\n+            return func(*args, **kwargs)\n+\n+    return _wrap_connect_internal\n+\n+\n class EngineTracer:\n     def __init__(self, tracer, engine, enable_commenter=False):\n         self.tracer = tracer\n", "issue": "Instrument SQLAlchemy engine connection phase\n** Is your feature request related to a problem? **\r\nThe SQLAlchemy instrumentation does not trace the actual connection to the database\r\n\r\n**Describe the solution you'd like**\r\nI want that `connect` function will also be traced\r\n\r\n**Describe alternatives you've considered**\r\nWhich alternative solutions or features have you considered?\r\n\r\n**Additional context**\r\nWe are working with SQLAlchemy (snowflake db) and we implemented a solution where we can see the `connect` span also as attached in the screenshot (the span is `database-connect`)\r\n![image](https://user-images.githubusercontent.com/12069200/173789203-b75f42d2-7ae6-4c11-aa0c-4937ede75422.png)\r\n\r\n\r\n\n", "before_files": [{"content": "# Copyright The OpenTelemetry Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nInstrument `sqlalchemy`_ to report SQL queries.\n\nThere are two options for instrumenting code. The first option is to use\nthe ``opentelemetry-instrument`` executable which will automatically\ninstrument your SQLAlchemy engine. The second is to programmatically enable\ninstrumentation via the following code:\n\n.. _sqlalchemy: https://pypi.org/project/sqlalchemy/\n\nUsage\n-----\n.. code:: python\n\n    from sqlalchemy import create_engine\n\n    from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor\n    import sqlalchemy\n\n    engine = create_engine(\"sqlite:///:memory:\")\n    SQLAlchemyInstrumentor().instrument(\n        engine=engine,\n    )\n\n    # of the async variant of SQLAlchemy\n\n    from sqlalchemy.ext.asyncio import create_async_engine\n\n    from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor\n    import sqlalchemy\n\n    engine = create_async_engine(\"sqlite:///:memory:\")\n    SQLAlchemyInstrumentor().instrument(\n        engine=engine.sync_engine\n    )\n\nAPI\n---\n\"\"\"\nfrom collections.abc import Sequence\nfrom typing import Collection\n\nimport sqlalchemy\nfrom packaging.version import parse as parse_version\nfrom wrapt import wrap_function_wrapper as _w\n\nfrom opentelemetry.instrumentation.instrumentor import BaseInstrumentor\nfrom opentelemetry.instrumentation.sqlalchemy.engine import (\n    EngineTracer,\n    _get_tracer,\n    _wrap_create_async_engine,\n    _wrap_create_engine,\n)\nfrom opentelemetry.instrumentation.sqlalchemy.package import _instruments\nfrom opentelemetry.instrumentation.utils import unwrap\n\n\nclass SQLAlchemyInstrumentor(BaseInstrumentor):\n    \"\"\"An instrumentor for SQLAlchemy\n    See `BaseInstrumentor`\n    \"\"\"\n\n    def instrumentation_dependencies(self) -> Collection[str]:\n        return _instruments\n\n    def _instrument(self, **kwargs):\n        \"\"\"Instruments SQLAlchemy engine creation methods and the engine\n        if passed as an argument.\n\n        Args:\n            **kwargs: Optional arguments\n                ``engine``: a SQLAlchemy engine instance\n                ``engines``: a list of SQLAlchemy engine instances\n                ``tracer_provider``: a TracerProvider, defaults to global\n\n        Returns:\n            An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.\n        \"\"\"\n        tracer_provider = kwargs.get(\"tracer_provider\")\n        _w(\"sqlalchemy\", \"create_engine\", _wrap_create_engine(tracer_provider))\n        _w(\n            \"sqlalchemy.engine\",\n            \"create_engine\",\n            _wrap_create_engine(tracer_provider),\n        )\n        if parse_version(sqlalchemy.__version__).release >= (1, 4):\n            _w(\n                \"sqlalchemy.ext.asyncio\",\n                \"create_async_engine\",\n                _wrap_create_async_engine(tracer_provider),\n            )\n\n        if kwargs.get(\"engine\") is not None:\n            return EngineTracer(\n                _get_tracer(tracer_provider),\n                kwargs.get(\"engine\"),\n                kwargs.get(\"enable_commenter\", False),\n            )\n        if kwargs.get(\"engines\") is not None and isinstance(\n            kwargs.get(\"engines\"), Sequence\n        ):\n            return [\n                EngineTracer(\n                    _get_tracer(tracer_provider),\n                    engine,\n                    kwargs.get(\"enable_commenter\", False),\n                )\n                for engine in kwargs.get(\"engines\")\n            ]\n\n        return None\n\n    def _uninstrument(self, **kwargs):\n        unwrap(sqlalchemy, \"create_engine\")\n        unwrap(sqlalchemy.engine, \"create_engine\")\n        if parse_version(sqlalchemy.__version__).release >= (1, 4):\n            unwrap(sqlalchemy.ext.asyncio, \"create_async_engine\")\n", "path": "instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py"}, {"content": "# Copyright The OpenTelemetry Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nfrom sqlalchemy.event import listen  # pylint: disable=no-name-in-module\n\nfrom opentelemetry import trace\nfrom opentelemetry.instrumentation.sqlalchemy.package import (\n    _instrumenting_module_name,\n)\nfrom opentelemetry.instrumentation.sqlalchemy.version import __version__\nfrom opentelemetry.instrumentation.utils import (\n    _generate_opentelemetry_traceparent,\n    _generate_sql_comment,\n)\nfrom opentelemetry.semconv.trace import NetTransportValues, SpanAttributes\nfrom opentelemetry.trace import Span\nfrom opentelemetry.trace.status import Status, StatusCode\n\n\ndef _normalize_vendor(vendor):\n    \"\"\"Return a canonical name for a type of database.\"\"\"\n    if not vendor:\n        return \"db\"  # should this ever happen?\n\n    if \"sqlite\" in vendor:\n        return \"sqlite\"\n\n    if \"postgres\" in vendor or vendor == \"psycopg2\":\n        return \"postgresql\"\n\n    return vendor\n\n\ndef _get_tracer(tracer_provider=None):\n    return trace.get_tracer(\n        _instrumenting_module_name,\n        __version__,\n        tracer_provider=tracer_provider,\n    )\n\n\ndef _wrap_create_async_engine(tracer_provider=None):\n    # pylint: disable=unused-argument\n    def _wrap_create_async_engine_internal(func, module, args, kwargs):\n        \"\"\"Trace the SQLAlchemy engine, creating an `EngineTracer`\n        object that will listen to SQLAlchemy events.\n        \"\"\"\n        engine = func(*args, **kwargs)\n        EngineTracer(_get_tracer(tracer_provider), engine.sync_engine)\n        return engine\n\n    return _wrap_create_async_engine_internal\n\n\ndef _wrap_create_engine(tracer_provider=None):\n    # pylint: disable=unused-argument\n    def _wrap_create_engine_internal(func, module, args, kwargs):\n        \"\"\"Trace the SQLAlchemy engine, creating an `EngineTracer`\n        object that will listen to SQLAlchemy events.\n        \"\"\"\n        engine = func(*args, **kwargs)\n        EngineTracer(_get_tracer(tracer_provider), engine)\n        return engine\n\n    return _wrap_create_engine_internal\n\n\nclass EngineTracer:\n    def __init__(self, tracer, engine, enable_commenter=False):\n        self.tracer = tracer\n        self.engine = engine\n        self.vendor = _normalize_vendor(engine.name)\n        self.enable_commenter = enable_commenter\n\n        listen(\n            engine, \"before_cursor_execute\", self._before_cur_exec, retval=True\n        )\n        listen(engine, \"after_cursor_execute\", _after_cur_exec)\n        listen(engine, \"handle_error\", _handle_error)\n\n    def _operation_name(self, db_name, statement):\n        parts = []\n        if isinstance(statement, str):\n            # otel spec recommends against parsing SQL queries. We are not trying to parse SQL\n            # but simply truncating the statement to the first word. This covers probably >95%\n            # use cases and uses the SQL statement in span name correctly as per the spec.\n            # For some very special cases it might not record the correct statement if the SQL\n            # dialect is too weird but in any case it shouldn't break anything.\n            parts.append(statement.split()[0])\n        if db_name:\n            parts.append(db_name)\n        if not parts:\n            return self.vendor\n        return \" \".join(parts)\n\n    # pylint: disable=unused-argument\n    def _before_cur_exec(\n        self, conn, cursor, statement, params, context, executemany\n    ):\n        attrs, found = _get_attributes_from_url(conn.engine.url)\n        if not found:\n            attrs = _get_attributes_from_cursor(self.vendor, cursor, attrs)\n\n        db_name = attrs.get(SpanAttributes.DB_NAME, \"\")\n        span = self.tracer.start_span(\n            self._operation_name(db_name, statement),\n            kind=trace.SpanKind.CLIENT,\n        )\n        with trace.use_span(span, end_on_exit=False):\n            if span.is_recording():\n                span.set_attribute(SpanAttributes.DB_STATEMENT, statement)\n                span.set_attribute(SpanAttributes.DB_SYSTEM, self.vendor)\n                for key, value in attrs.items():\n                    span.set_attribute(key, value)\n\n        context._otel_span = span\n        if self.enable_commenter:\n            statement = statement + EngineTracer._generate_comment(span=span)\n\n        return statement, params\n\n    @staticmethod\n    def _generate_comment(span: Span) -> str:\n        span_context = span.get_span_context()\n        meta = {}\n        if span_context.is_valid:\n            meta.update(_generate_opentelemetry_traceparent(span))\n        return _generate_sql_comment(**meta)\n\n\n# pylint: disable=unused-argument\ndef _after_cur_exec(conn, cursor, statement, params, context, executemany):\n    span = getattr(context, \"_otel_span\", None)\n    if span is None:\n        return\n\n    span.end()\n\n\ndef _handle_error(context):\n    span = getattr(context.execution_context, \"_otel_span\", None)\n    if span is None:\n        return\n\n    if span.is_recording():\n        span.set_status(\n            Status(\n                StatusCode.ERROR,\n                str(context.original_exception),\n            )\n        )\n    span.end()\n\n\ndef _get_attributes_from_url(url):\n    \"\"\"Set connection tags from the url. return true if successful.\"\"\"\n    attrs = {}\n    if url.host:\n        attrs[SpanAttributes.NET_PEER_NAME] = url.host\n    if url.port:\n        attrs[SpanAttributes.NET_PEER_PORT] = url.port\n    if url.database:\n        attrs[SpanAttributes.DB_NAME] = url.database\n    if url.username:\n        attrs[SpanAttributes.DB_USER] = url.username\n    return attrs, bool(url.host)\n\n\ndef _get_attributes_from_cursor(vendor, cursor, attrs):\n    \"\"\"Attempt to set db connection attributes by introspecting the cursor.\"\"\"\n    if vendor == \"postgresql\":\n        info = getattr(getattr(cursor, \"connection\", None), \"info\", None)\n        if not info:\n            return attrs\n\n        attrs[SpanAttributes.DB_NAME] = info.dbname\n        is_unix_socket = info.host and info.host.startswith(\"/\")\n\n        if is_unix_socket:\n            attrs[SpanAttributes.NET_TRANSPORT] = NetTransportValues.UNIX.value\n            if info.port:\n                # postgresql enforces this pattern on all socket names\n                attrs[SpanAttributes.NET_PEER_NAME] = os.path.join(\n                    info.host, f\".s.PGSQL.{info.port}\"\n                )\n        else:\n            attrs[\n                SpanAttributes.NET_TRANSPORT\n            ] = NetTransportValues.IP_TCP.value\n            attrs[SpanAttributes.NET_PEER_NAME] = info.host\n            if info.port:\n                attrs[SpanAttributes.NET_PEER_PORT] = int(info.port)\n    return attrs\n", "path": "instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py"}], "after_files": [{"content": "# Copyright The OpenTelemetry Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nInstrument `sqlalchemy`_ to report SQL queries.\n\nThere are two options for instrumenting code. The first option is to use\nthe ``opentelemetry-instrument`` executable which will automatically\ninstrument your SQLAlchemy engine. The second is to programmatically enable\ninstrumentation via the following code:\n\n.. _sqlalchemy: https://pypi.org/project/sqlalchemy/\n\nUsage\n-----\n.. code:: python\n\n    from sqlalchemy import create_engine\n\n    from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor\n    import sqlalchemy\n\n    engine = create_engine(\"sqlite:///:memory:\")\n    SQLAlchemyInstrumentor().instrument(\n        engine=engine,\n    )\n\n    # of the async variant of SQLAlchemy\n\n    from sqlalchemy.ext.asyncio import create_async_engine\n\n    from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor\n    import sqlalchemy\n\n    engine = create_async_engine(\"sqlite:///:memory:\")\n    SQLAlchemyInstrumentor().instrument(\n        engine=engine.sync_engine\n    )\n\nAPI\n---\n\"\"\"\nfrom collections.abc import Sequence\nfrom typing import Collection\n\nimport sqlalchemy\nfrom packaging.version import parse as parse_version\nfrom sqlalchemy.engine.base import Engine\nfrom wrapt import wrap_function_wrapper as _w\n\nfrom opentelemetry.instrumentation.instrumentor import BaseInstrumentor\nfrom opentelemetry.instrumentation.sqlalchemy.engine import (\n    EngineTracer,\n    _get_tracer,\n    _wrap_connect,\n    _wrap_create_async_engine,\n    _wrap_create_engine,\n)\nfrom opentelemetry.instrumentation.sqlalchemy.package import _instruments\nfrom opentelemetry.instrumentation.utils import unwrap\n\n\nclass SQLAlchemyInstrumentor(BaseInstrumentor):\n    \"\"\"An instrumentor for SQLAlchemy\n    See `BaseInstrumentor`\n    \"\"\"\n\n    def instrumentation_dependencies(self) -> Collection[str]:\n        return _instruments\n\n    def _instrument(self, **kwargs):\n        \"\"\"Instruments SQLAlchemy engine creation methods and the engine\n        if passed as an argument.\n\n        Args:\n            **kwargs: Optional arguments\n                ``engine``: a SQLAlchemy engine instance\n                ``engines``: a list of SQLAlchemy engine instances\n                ``tracer_provider``: a TracerProvider, defaults to global\n\n        Returns:\n            An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.\n        \"\"\"\n        tracer_provider = kwargs.get(\"tracer_provider\")\n        _w(\"sqlalchemy\", \"create_engine\", _wrap_create_engine(tracer_provider))\n        _w(\n            \"sqlalchemy.engine\",\n            \"create_engine\",\n            _wrap_create_engine(tracer_provider),\n        )\n        _w(\n            \"sqlalchemy.engine.base\",\n            \"Engine.connect\",\n            _wrap_connect(tracer_provider),\n        )\n        if parse_version(sqlalchemy.__version__).release >= (1, 4):\n            _w(\n                \"sqlalchemy.ext.asyncio\",\n                \"create_async_engine\",\n                _wrap_create_async_engine(tracer_provider),\n            )\n        if kwargs.get(\"engine\") is not None:\n            return EngineTracer(\n                _get_tracer(tracer_provider),\n                kwargs.get(\"engine\"),\n                kwargs.get(\"enable_commenter\", False),\n            )\n        if kwargs.get(\"engines\") is not None and isinstance(\n            kwargs.get(\"engines\"), Sequence\n        ):\n            return [\n                EngineTracer(\n                    _get_tracer(tracer_provider),\n                    engine,\n                    kwargs.get(\"enable_commenter\", False),\n                )\n                for engine in kwargs.get(\"engines\")\n            ]\n\n        return None\n\n    def _uninstrument(self, **kwargs):\n        unwrap(sqlalchemy, \"create_engine\")\n        unwrap(sqlalchemy.engine, \"create_engine\")\n        unwrap(Engine, \"connect\")\n        if parse_version(sqlalchemy.__version__).release >= (1, 4):\n            unwrap(sqlalchemy.ext.asyncio, \"create_async_engine\")\n", "path": "instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py"}, {"content": "# Copyright The OpenTelemetry Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nfrom sqlalchemy.event import listen  # pylint: disable=no-name-in-module\n\nfrom opentelemetry import trace\nfrom opentelemetry.instrumentation.sqlalchemy.package import (\n    _instrumenting_module_name,\n)\nfrom opentelemetry.instrumentation.sqlalchemy.version import __version__\nfrom opentelemetry.instrumentation.utils import (\n    _generate_opentelemetry_traceparent,\n    _generate_sql_comment,\n)\nfrom opentelemetry.semconv.trace import NetTransportValues, SpanAttributes\nfrom opentelemetry.trace import Span\nfrom opentelemetry.trace.status import Status, StatusCode\n\n\ndef _normalize_vendor(vendor):\n    \"\"\"Return a canonical name for a type of database.\"\"\"\n    if not vendor:\n        return \"db\"  # should this ever happen?\n\n    if \"sqlite\" in vendor:\n        return \"sqlite\"\n\n    if \"postgres\" in vendor or vendor == \"psycopg2\":\n        return \"postgresql\"\n\n    return vendor\n\n\ndef _get_tracer(tracer_provider=None):\n    return trace.get_tracer(\n        _instrumenting_module_name,\n        __version__,\n        tracer_provider=tracer_provider,\n    )\n\n\ndef _wrap_create_async_engine(tracer_provider=None):\n    # pylint: disable=unused-argument\n    def _wrap_create_async_engine_internal(func, module, args, kwargs):\n        \"\"\"Trace the SQLAlchemy engine, creating an `EngineTracer`\n        object that will listen to SQLAlchemy events.\n        \"\"\"\n        engine = func(*args, **kwargs)\n        EngineTracer(_get_tracer(tracer_provider), engine.sync_engine)\n        return engine\n\n    return _wrap_create_async_engine_internal\n\n\ndef _wrap_create_engine(tracer_provider=None):\n    # pylint: disable=unused-argument\n    def _wrap_create_engine_internal(func, module, args, kwargs):\n        \"\"\"Trace the SQLAlchemy engine, creating an `EngineTracer`\n        object that will listen to SQLAlchemy events.\n        \"\"\"\n        engine = func(*args, **kwargs)\n        EngineTracer(_get_tracer(tracer_provider), engine)\n        return engine\n\n    return _wrap_create_engine_internal\n\n\ndef _wrap_connect(tracer_provider=None):\n    tracer = trace.get_tracer(\n        _instrumenting_module_name,\n        __version__,\n        tracer_provider=tracer_provider,\n    )\n\n    # pylint: disable=unused-argument\n    def _wrap_connect_internal(func, module, args, kwargs):\n        with tracer.start_as_current_span(\n            \"connect\", kind=trace.SpanKind.CLIENT\n        ):\n            return func(*args, **kwargs)\n\n    return _wrap_connect_internal\n\n\nclass EngineTracer:\n    def __init__(self, tracer, engine, enable_commenter=False):\n        self.tracer = tracer\n        self.engine = engine\n        self.vendor = _normalize_vendor(engine.name)\n        self.enable_commenter = enable_commenter\n\n        listen(\n            engine, \"before_cursor_execute\", self._before_cur_exec, retval=True\n        )\n        listen(engine, \"after_cursor_execute\", _after_cur_exec)\n        listen(engine, \"handle_error\", _handle_error)\n\n    def _operation_name(self, db_name, statement):\n        parts = []\n        if isinstance(statement, str):\n            # otel spec recommends against parsing SQL queries. We are not trying to parse SQL\n            # but simply truncating the statement to the first word. This covers probably >95%\n            # use cases and uses the SQL statement in span name correctly as per the spec.\n            # For some very special cases it might not record the correct statement if the SQL\n            # dialect is too weird but in any case it shouldn't break anything.\n            parts.append(statement.split()[0])\n        if db_name:\n            parts.append(db_name)\n        if not parts:\n            return self.vendor\n        return \" \".join(parts)\n\n    # pylint: disable=unused-argument\n    def _before_cur_exec(\n        self, conn, cursor, statement, params, context, executemany\n    ):\n        attrs, found = _get_attributes_from_url(conn.engine.url)\n        if not found:\n            attrs = _get_attributes_from_cursor(self.vendor, cursor, attrs)\n\n        db_name = attrs.get(SpanAttributes.DB_NAME, \"\")\n        span = self.tracer.start_span(\n            self._operation_name(db_name, statement),\n            kind=trace.SpanKind.CLIENT,\n        )\n        with trace.use_span(span, end_on_exit=False):\n            if span.is_recording():\n                span.set_attribute(SpanAttributes.DB_STATEMENT, statement)\n                span.set_attribute(SpanAttributes.DB_SYSTEM, self.vendor)\n                for key, value in attrs.items():\n                    span.set_attribute(key, value)\n\n        context._otel_span = span\n        if self.enable_commenter:\n            statement = statement + EngineTracer._generate_comment(span=span)\n\n        return statement, params\n\n    @staticmethod\n    def _generate_comment(span: Span) -> str:\n        span_context = span.get_span_context()\n        meta = {}\n        if span_context.is_valid:\n            meta.update(_generate_opentelemetry_traceparent(span))\n        return _generate_sql_comment(**meta)\n\n\n# pylint: disable=unused-argument\ndef _after_cur_exec(conn, cursor, statement, params, context, executemany):\n    span = getattr(context, \"_otel_span\", None)\n    if span is None:\n        return\n\n    span.end()\n\n\ndef _handle_error(context):\n    span = getattr(context.execution_context, \"_otel_span\", None)\n    if span is None:\n        return\n\n    if span.is_recording():\n        span.set_status(\n            Status(\n                StatusCode.ERROR,\n                str(context.original_exception),\n            )\n        )\n    span.end()\n\n\ndef _get_attributes_from_url(url):\n    \"\"\"Set connection tags from the url. return true if successful.\"\"\"\n    attrs = {}\n    if url.host:\n        attrs[SpanAttributes.NET_PEER_NAME] = url.host\n    if url.port:\n        attrs[SpanAttributes.NET_PEER_PORT] = url.port\n    if url.database:\n        attrs[SpanAttributes.DB_NAME] = url.database\n    if url.username:\n        attrs[SpanAttributes.DB_USER] = url.username\n    return attrs, bool(url.host)\n\n\ndef _get_attributes_from_cursor(vendor, cursor, attrs):\n    \"\"\"Attempt to set db connection attributes by introspecting the cursor.\"\"\"\n    if vendor == \"postgresql\":\n        info = getattr(getattr(cursor, \"connection\", None), \"info\", None)\n        if not info:\n            return attrs\n\n        attrs[SpanAttributes.DB_NAME] = info.dbname\n        is_unix_socket = info.host and info.host.startswith(\"/\")\n\n        if is_unix_socket:\n            attrs[SpanAttributes.NET_TRANSPORT] = NetTransportValues.UNIX.value\n            if info.port:\n                # postgresql enforces this pattern on all socket names\n                attrs[SpanAttributes.NET_PEER_NAME] = os.path.join(\n                    info.host, f\".s.PGSQL.{info.port}\"\n                )\n        else:\n            attrs[\n                SpanAttributes.NET_TRANSPORT\n            ] = NetTransportValues.IP_TCP.value\n            attrs[SpanAttributes.NET_PEER_NAME] = info.host\n            if info.port:\n                attrs[SpanAttributes.NET_PEER_PORT] = int(info.port)\n    return attrs\n", "path": "instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py"}]}