diff --git a/MANIFEST.in b/MANIFEST.in index c3426bad..4fa1d272 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,4 +1,5 @@ include ibm_db.h ibm_db.c +include ibm_db_otel.py include CHANGES.md LICENSE README.md include ibm_db_dll.pth include _ibm_db_register_dll.py diff --git a/README.md b/README.md index 0e99ad16..958f4eac 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,45 @@ https://github.com/ibmdb/python-ibmdb/wiki/APIs For the asyncio APIs, please refer to [ASYNC_APIs_WIKI.md](ASYNC_APIs_WIKI.md). +## OpenTelemetry + +Install `opentelemetry-api` and an OpenTelemetry SDK to propagate the active +W3C trace context from Python. The `ibm_db_otel` helper reads the current +context and uses `ibm_db.set_option` to set the CLI connection attributes for +`traceparent`, `tracestate`, and `baggage` before a new connection is +established and refreshes them before SQL execution. This applies to both +`ibm_db` and `ibm_db_dbi` execution APIs, including connections opened before +the span starts. Context attributes explicitly set by the application are not +replaced. + +The connection attribute constants are available as +`ibm_db.SQL_ATTR_OTEL_TRACEPARENT`, `ibm_db.SQL_ATTR_OTEL_TRACESTATE`, and +`ibm_db.SQL_ATTR_OTEL_BAGGAGE`. For example: + +```python +ibm_db.set_option(conn, { + ibm_db.SQL_ATTR_OTEL_TRACEPARENT: traceparent, + ibm_db.SQL_ATTR_OTEL_TRACESTATE: tracestate, + ibm_db.SQL_ATTR_OTEL_BAGGAGE: baggage, +}, 1) +``` + +Automatic propagation is ignored when the installed CLI driver does not +support these attributes. Explicitly setting an attribute with such a driver +raises an exception. + +Run the end-to-end demo with a Db2 connection string: + +```sh +export DB2_CONNSTR='DATABASE=...;HOSTNAME=...;PORT=...;UID=...;PWD=...' +python ibm_db_tests/demo_otel_context.py +``` + +The demo shows automatic propagation at connect time and before direct and +prepared SQL execution, propagation through `ibm_db_dbi`, context changes on +an existing connection (including fixed W3C carriers), clearing stale values +when no context is active, and explicit attributes taking precedence. + ## Timestamp Precision Notes - Python `datetime` stores microseconds, so `ibm_db` returns timestamp values with up to 6 fractional digits in Python objects. @@ -707,6 +746,10 @@ set DB2_PASSWD= ``` +The OpenTelemetry connection-attribute test also accepts a full CLI +connection string through `DB2_CONNSTR`. When set, it takes precedence over +the database and credential values above for that test. + - OR ``` diff --git a/ibm_db.c b/ibm_db.c index d51ebae5..9b271864 100644 --- a/ibm_db.c +++ b/ibm_db.c @@ -27,6 +27,21 @@ #include #include #include "ibm_db.h" + +#ifndef SQL_ATTR_OTEL_TRACEPARENT +#define SQL_ATTR_OTEL_TRACEPARENT 3056 +#endif +#ifndef SQL_ATTR_OTEL_TRACESTATE +#define SQL_ATTR_OTEL_TRACESTATE 3057 +#endif +#ifndef SQL_ATTR_OTEL_BAGGAGE +#define SQL_ATTR_OTEL_BAGGAGE 3058 +#endif + +#define IBM_DB_OTEL_TRACEPARENT_SET 0x01 +#define IBM_DB_OTEL_TRACESTATE_SET 0x02 +#define IBM_DB_OTEL_BAGGAGE_SET 0x04 + #include #include #ifdef _WIN32 @@ -196,6 +211,10 @@ typedef struct _conn_handle_struct SQLSMALLINT errormsg_recno_tracker; int flag_pconnect; /* Indicates that this connection is persistent */ int flag_coordinated; /* Indicates coordinated connection (shared env) */ + unsigned int otel_explicit_attrs; + unsigned int otel_auto_attrs; + int otel_auto_update; + int otel_attrs_unsupported; } conn_handle; static void _python_ibm_db_free_conn_struct(conn_handle *handle); @@ -369,6 +388,7 @@ typedef struct _stmt_handle_struct SQLULEN rowset_size; SQLULEN rows_fetched; SQLUSMALLINT *row_status_array; + conn_handle *conn_res; } stmt_handle; static void _python_ibm_db_free_stmt_struct(stmt_handle *handle); @@ -955,6 +975,8 @@ static stmt_handle *_ibm_db_new_stmt_struct(conn_handle *conn_res) /* Initialize stmt resource so parsing assigns updated options if needed */ stmt_res->hdbc = conn_res->hdbc; + stmt_res->conn_res = conn_res; + Py_INCREF(conn_res); stmt_res->s_bin_mode = conn_res->c_bin_mode; stmt_res->cursor_type = conn_res->c_cursor_type; stmt_res->s_case_mode = conn_res->c_case_mode; @@ -1010,6 +1032,7 @@ static void _python_ibm_db_free_stmt_struct(stmt_handle *handle) LogMsg(DEBUG, messageStr); if (handle != NULL) { + Py_XDECREF((PyObject *)handle->conn_res); Py_TYPE(handle)->tp_free((PyObject *)handle); } LogMsg(INFO, "exit _python_ibm_db_free_stmt_struct()"); @@ -1202,6 +1225,220 @@ static void _python_ibm_db_check_sql_errors(SQLHANDLE handle, SQLSMALLINT hType, } } +static unsigned int _python_ibm_db_otel_attr_bit(long opt_key) +{ + switch (opt_key) + { + case SQL_ATTR_OTEL_TRACEPARENT: + return IBM_DB_OTEL_TRACEPARENT_SET; + case SQL_ATTR_OTEL_TRACESTATE: + return IBM_DB_OTEL_TRACESTATE_SET; + case SQL_ATTR_OTEL_BAGGAGE: + return IBM_DB_OTEL_BAGGAGE_SET; + default: + return 0; + } +} + +static int _python_ibm_db_otel_attr_is_unsupported(SQLHDBC hdbc) +{ + SQLCHAR sqlstate[SQL_SQLSTATE_SIZE + 1] = {0}; + SQLINTEGER native_error = 0; + SQLCHAR message[SQL_MAX_MESSAGE_LENGTH + 1]; + SQLSMALLINT message_length = 0; + SQLRETURN rc; + + rc = SQLGetDiagRec(SQL_HANDLE_DBC, hdbc, 1, sqlstate, &native_error, + message, SQL_MAX_MESSAGE_LENGTH, &message_length); + if (rc != SQL_SUCCESS && rc != SQL_SUCCESS_WITH_INFO) + { + return 0; + } + + return strcmp((const char *)sqlstate, "HY092") == 0 || + strcmp((const char *)sqlstate, "HYC00") == 0 || + strcmp((const char *)sqlstate, "IM001") == 0; +} + +static int _python_ibm_db_set_otel_attr(SQLHDBC hdbc, long opt_key, + PyObject *data) +{ + PyObject *unicode_value = NULL; + SQLRETURN rc; + int isNewBuffer = 0; +#ifdef __MVS__ + SQLCHAR *option_str; +#else + SQLWCHAR *option_str; +#endif + + if (!(PyString_Check(data) || PyUnicode_Check(data))) + { + PyErr_SetString(PyExc_TypeError, + "OpenTelemetry connection attributes must be strings"); + return SQL_ERROR; + } + + unicode_value = PyUnicode_FromObject(data); + if (unicode_value == NULL) + { + return SQL_ERROR; + } +#ifdef __MVS__ + option_str = getUnicodeDataAsSQLCHAR(unicode_value, &isNewBuffer); +#else + option_str = getUnicodeDataAsSQLWCHAR(unicode_value, &isNewBuffer, NULL); +#endif + Py_DECREF(unicode_value); + if (option_str == NULL) + { + return SQL_ERROR; + } + +#ifdef __MVS__ + Py_BEGIN_ALLOW_THREADS; + rc = SQLSetConnectAttr(hdbc, (SQLINTEGER)opt_key, + (SQLPOINTER)option_str, SQL_NTS); + Py_END_ALLOW_THREADS; +#else + Py_BEGIN_ALLOW_THREADS; + rc = SQLSetConnectAttrW(hdbc, (SQLINTEGER)opt_key, + (SQLPOINTER)option_str, SQL_NTS); + Py_END_ALLOW_THREADS; +#endif + if (isNewBuffer) + { + PyMem_Del(option_str); + } + + if (rc == SQL_ERROR) + { + if (_python_ibm_db_otel_attr_is_unsupported(hdbc)) + { + PyErr_SetString(PyExc_Exception, + "The clidriver used with ibm_db does not support OpenTelemetry connection attributes"); + return SQL_ERROR; + } + _python_ibm_db_check_sql_errors(hdbc, SQL_HANDLE_DBC, rc, 1, + NULL, -1, 1); + return SQL_ERROR; + } + + return SQL_SUCCESS; +} + +static int _python_ibm_db_append_otel_attr(PyObject *attr_list, long attr_id) +{ + PyObject *attr_value = PyLong_FromLong(attr_id); + int rc; + + if (attr_value == NULL) + { + return -1; + } + rc = PyList_Append(attr_list, attr_value); + Py_DECREF(attr_value); + return rc; +} + +static int _python_ibm_db_refresh_otel_context(conn_handle *conn_res) +{ + static const long attr_ids[] = { + SQL_ATTR_OTEL_TRACEPARENT, + SQL_ATTR_OTEL_TRACESTATE, + SQL_ATTR_OTEL_BAGGAGE + }; + static const unsigned int attr_bits[] = { + IBM_DB_OTEL_TRACEPARENT_SET, + IBM_DB_OTEL_TRACESTATE_SET, + IBM_DB_OTEL_BAGGAGE_SET + }; + PyObject *module = NULL; + PyObject *apply_context = NULL; + PyObject *explicit_attrs = NULL; + PyObject *clear_attrs = NULL; + PyObject *result = NULL; + unsigned int auto_attrs; + int i; + + if (conn_res == NULL || conn_res->otel_attrs_unsupported) + { + return SQL_SUCCESS; + } + + explicit_attrs = PyList_New(0); + clear_attrs = PyList_New(0); + if (explicit_attrs == NULL || clear_attrs == NULL) + { + Py_XDECREF(explicit_attrs); + Py_XDECREF(clear_attrs); + return SQL_ERROR; + } + + for (i = 0; i < 3; i++) + { + if (conn_res->otel_explicit_attrs & attr_bits[i]) + { + if (_python_ibm_db_append_otel_attr(explicit_attrs, attr_ids[i]) < 0) + { + Py_DECREF(explicit_attrs); + Py_DECREF(clear_attrs); + return SQL_ERROR; + } + } + else if (conn_res->otel_auto_attrs & attr_bits[i]) + { + if (_python_ibm_db_append_otel_attr(clear_attrs, attr_ids[i]) < 0) + { + Py_DECREF(explicit_attrs); + Py_DECREF(clear_attrs); + return SQL_ERROR; + } + } + } + + module = PyImport_ImportModule("ibm_db_otel"); + if (module != NULL) + { + apply_context = PyObject_GetAttrString(module, "apply_otel_context"); + } + if (apply_context != NULL) + { + conn_res->otel_auto_update = 1; + result = PyObject_CallFunctionObjArgs( + apply_context, (PyObject *)conn_res, explicit_attrs, clear_attrs, + NULL); + conn_res->otel_auto_update = 0; + } + + Py_XDECREF(apply_context); + Py_XDECREF(module); + Py_DECREF(explicit_attrs); + Py_DECREF(clear_attrs); + if (result == NULL) + { + return SQL_ERROR; + } + if (result == Py_None) + { + conn_res->otel_attrs_unsupported = 1; + Py_DECREF(result); + return SQL_SUCCESS; + } + + auto_attrs = (unsigned int)PyLong_AsUnsignedLong(result); + Py_DECREF(result); + if (PyErr_Occurred()) + { + return SQL_ERROR; + } + conn_res->otel_auto_attrs = auto_attrs & + (IBM_DB_OTEL_TRACEPARENT_SET | + IBM_DB_OTEL_TRACESTATE_SET | + IBM_DB_OTEL_BAGGAGE_SET); + return SQL_SUCCESS; +} + /* static int _python_ibm_db_assign_options( void *handle, int type, long opt_key, PyObject *data ) */ static int _python_ibm_db_assign_options(void *handle, int type, long opt_key, PyObject *data) { @@ -1217,6 +1454,23 @@ static int _python_ibm_db_assign_options(void *handle, int type, long opt_key, P int isNewBuffer = 0; snprintf(messageStr, sizeof(messageStr), "Parameters - Handle: %p, type: %d, opt_key: %ld, data: %p", handle, type, opt_key, data); LogMsg(DEBUG, messageStr); + if (type == SQL_HANDLE_DBC && _python_ibm_db_otel_attr_bit(opt_key) != 0) + { + unsigned int attr_bit = _python_ibm_db_otel_attr_bit(opt_key); + rc = _python_ibm_db_set_otel_attr( + (SQLHDBC)((conn_handle *)handle)->hdbc, opt_key, data); + if (rc == SQL_ERROR) + { + return -1; + } + if (!((conn_handle *)handle)->otel_auto_update) + { + ((conn_handle *)handle)->otel_explicit_attrs |= attr_bit; + ((conn_handle *)handle)->otel_auto_attrs &= ~attr_bit; + ((conn_handle *)handle)->otel_attrs_unsupported = 0; + } + return SQL_SUCCESS; + } /* First check to see if it is a non-cli attribut */ if (opt_key == ATTR_CASE) { @@ -3428,6 +3682,10 @@ static PyObject *_python_ibm_db_connect_helper(PyObject *self, PyObject *args, i { conn_res->henv = 0; conn_res->hdbc = 0; + conn_res->otel_explicit_attrs = 0; + conn_res->otel_auto_attrs = 0; + conn_res->otel_auto_update = 0; + conn_res->otel_attrs_unsupported = 0; LogMsg(DEBUG, "Created a new connection handle"); } else @@ -3544,6 +3802,19 @@ static PyObject *_python_ibm_db_connect_helper(PyObject *self, PyObject *args, i if (!reused) { + rc = _python_ibm_db_refresh_otel_context(conn_res); + if (rc == SQL_ERROR) + { + LogMsg(ERROR, "Failed to apply automatic OpenTelemetry context before connect"); + Py_BEGIN_ALLOW_THREADS; + SQLFreeHandle(SQL_HANDLE_DBC, conn_res->hdbc); + SQLFreeHandle(SQL_HANDLE_ENV, conn_res->henv); + Py_END_ALLOW_THREADS; + conn_res->hdbc = 0; + conn_res->henv = 0; + break; + } + /* Connect */ /* If the string contains a =, use SQLDriverConnect */ if (NIL_P(databaseObj)) @@ -8788,6 +9059,23 @@ static PyObject *ibm_db_exec(PyObject *self, PyObject *args) LogMsg(DEBUG, messageStr); } + rc = _python_ibm_db_refresh_otel_context(conn_res); + if (rc == SQL_ERROR) + { + Py_BEGIN_ALLOW_THREADS; + SQLFreeHandle(SQL_HANDLE_STMT, stmt_res->hstmt); + Py_END_ALLOW_THREADS; + stmt_res->hstmt = -1; + Py_DECREF((PyObject *)stmt_res); + if (isNewBuffer && stmt != NULL) + { + PyMem_Del(stmt); + } + Py_XDECREF(py_stmt); + PyMem_Del(return_str); + return NULL; + } + Py_BEGIN_ALLOW_THREADS; rc = SQLExecDirectW((SQLHSTMT)stmt_res->hstmt, stmt, SQL_NTS); snprintf(messageStr, sizeof(messageStr), "SQLExecDirectW result: rc = %d", rc); @@ -11845,6 +12133,11 @@ static PyObject *_python_ibm_db_execute_helper1(stmt_handle *stmt_res, PyObject /* No Parameters * We just execute the statement. No additional work needed. */ + rc = _python_ibm_db_refresh_otel_context(stmt_res->conn_res); + if (rc == SQL_ERROR) + { + return NULL; + } Py_BEGIN_ALLOW_THREADS; rc = SQLExecute((SQLHSTMT)stmt_res->hstmt); Py_END_ALLOW_THREADS; @@ -11869,6 +12162,11 @@ static PyObject *_python_ibm_db_execute_helper1(stmt_handle *stmt_res, PyObject } /* Execute Stmt -- All parameters bound */ + rc = _python_ibm_db_refresh_otel_context(stmt_res->conn_res); + if (rc == SQL_ERROR) + { + return NULL; + } Py_BEGIN_ALLOW_THREADS; rc = SQLExecute((SQLHSTMT)stmt_res->hstmt); Py_END_ALLOW_THREADS; @@ -12806,6 +13104,8 @@ static PyObject *ibm_db_next_result(PyObject *self, PyObject *args) new_stmt_res->hstmt = new_hstmt; #endif new_stmt_res->hdbc = stmt_res->hdbc; + new_stmt_res->conn_res = stmt_res->conn_res; + Py_XINCREF((PyObject *)new_stmt_res->conn_res); new_stmt_res->s_bin_mode = stmt_res->s_bin_mode; new_stmt_res->cursor_type = stmt_res->cursor_type; @@ -16195,8 +16495,11 @@ static PyObject *ibm_db_set_option(PyObject *self, PyObject *args) if (rc == SQL_ERROR) { LogMsg(ERROR, "Options array must have string indexes"); - PyErr_SetString(PyExc_Exception, - "Options Array must have string indexes"); + if (!PyErr_Occurred()) + { + PyErr_SetString(PyExc_Exception, + "Options Array must have string indexes"); + } return NULL; } } @@ -16222,8 +16525,11 @@ static PyObject *ibm_db_set_option(PyObject *self, PyObject *args) if (rc == SQL_ERROR) { LogMsg(ERROR, "Options array must have string indexes"); - PyErr_SetString(PyExc_Exception, - "Options Array must have string indexes"); + if (!PyErr_Occurred()) + { + PyErr_SetString(PyExc_Exception, + "Options Array must have string indexes"); + } return NULL; } } @@ -19103,6 +19409,11 @@ static PyObject *ibm_db_execute_many(PyObject *self, PyObject *args) if (error[0] == '\0') { + rc = _python_ibm_db_refresh_otel_context(stmt_res->conn_res); + if (rc == SQL_ERROR) + { + return NULL; + } Py_BEGIN_ALLOW_THREADS; rc = SQLExecute((SQLHSTMT)stmt_res->hstmt); Py_END_ALLOW_THREADS; @@ -21303,6 +21614,10 @@ static PyObject *ibm_db_connect_coordinated(PyObject *self, PyObject *args) conn_res->hdbc = 0; conn_res->flag_pconnect = 0; conn_res->flag_coordinated = 1; /* Mark as coordinated */ + conn_res->otel_explicit_attrs = 0; + conn_res->otel_auto_attrs = 0; + conn_res->otel_auto_update = 0; + conn_res->otel_attrs_unsupported = 0; conn_res->handle_active = 0; conn_res->auto_commit = SQL_AUTOCOMMIT_OFF; conn_res->c_bin_mode = IBM_DB_G(bin_mode); @@ -21710,6 +22025,9 @@ INIT_ibm_db(void) PyModule_AddIntConstant(m, "SQL_ATTR_INFO_WRKSTNNAME", SQL_ATTR_INFO_WRKSTNNAME); PyModule_AddIntConstant(m, "SQL_ATTR_INFO_ACCTSTR", SQL_ATTR_INFO_ACCTSTR); PyModule_AddIntConstant(m, "SQL_ATTR_INFO_APPLNAME", SQL_ATTR_INFO_APPLNAME); + PyModule_AddIntConstant(m, "SQL_ATTR_OTEL_TRACEPARENT", SQL_ATTR_OTEL_TRACEPARENT); + PyModule_AddIntConstant(m, "SQL_ATTR_OTEL_TRACESTATE", SQL_ATTR_OTEL_TRACESTATE); + PyModule_AddIntConstant(m, "SQL_ATTR_OTEL_BAGGAGE", SQL_ATTR_OTEL_BAGGAGE); PyModule_AddIntConstant(m, "SQL_ATTR_USE_TRUSTED_CONTEXT", SQL_ATTR_USE_TRUSTED_CONTEXT); PyModule_AddIntConstant(m, "SQL_ATTR_TRUSTED_CONTEXT_USERID", SQL_ATTR_TRUSTED_CONTEXT_USERID); PyModule_AddIntConstant(m, "SQL_ATTR_TRUSTED_CONTEXT_PASSWORD", SQL_ATTR_TRUSTED_CONTEXT_PASSWORD); diff --git a/ibm_db_dbi.py b/ibm_db_dbi.py index 55fb9e12..1beacf9e 100644 --- a/ibm_db_dbi.py +++ b/ibm_db_dbi.py @@ -129,6 +129,11 @@ async def _async_to_thread_with_logging(name, func, *args): SQL_TXN_SERIALIZABLE = ibm_db.SQL_TXN_SERIALIZABLE SQL_TXN_NO_COMMIT = ibm_db.SQL_TXN_NO_COMMIT +# OpenTelemetry W3C connection attributes. +SQL_ATTR_OTEL_TRACEPARENT = ibm_db.SQL_ATTR_OTEL_TRACEPARENT +SQL_ATTR_OTEL_TRACESTATE = ibm_db.SQL_ATTR_OTEL_TRACESTATE +SQL_ATTR_OTEL_BAGGAGE = ibm_db.SQL_ATTR_OTEL_BAGGAGE + # DRDA fast load (zLOAD) constants are exported only when available in CLI headers. SQL_ATTR_DB2ZLOAD_LOADSTMT = getattr(ibm_db, 'SQL_ATTR_DB2ZLOAD_LOADSTMT', None) SQL_ATTR_DB2ZLOAD_UTILITYID = getattr(ibm_db, 'SQL_ATTR_DB2ZLOAD_UTILITYID', None) @@ -2856,4 +2861,3 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): LogMsg(INFO, "exit AsyncCoordinatedConnection.__aexit__()") return False - diff --git a/ibm_db_otel.py b/ibm_db_otel.py new file mode 100644 index 00000000..0b751eca --- /dev/null +++ b/ibm_db_otel.py @@ -0,0 +1,62 @@ +"""OpenTelemetry context extraction and ibm_db connection attribute updates.""" + +_OTEL_ATTRIBUTES = ( + ("traceparent", "SQL_ATTR_OTEL_TRACEPARENT", 0x01), + ("tracestate", "SQL_ATTR_OTEL_TRACESTATE", 0x02), + ("baggage", "SQL_ATTR_OTEL_BAGGAGE", 0x04), +) + + +def _current_context_headers(): + try: + from opentelemetry.baggage.propagation import W3CBaggagePropagator + from opentelemetry.trace.propagation.tracecontext import ( + TraceContextTextMapPropagator, + ) + except ImportError: + return {} + + headers = {} + TraceContextTextMapPropagator().inject(headers) + W3CBaggagePropagator().inject(headers) + return headers + + +def apply_otel_context(connection, explicit_attrs=(), clear_attrs=()): + """Set current W3C context on a connection using ibm_db.set_option. + + Returns the bit mask of automatically populated attributes, or None when + the CLI driver does not support the OTEL connection attributes. + """ + import ibm_db + + headers = _current_context_headers() + explicit_attrs = set(explicit_attrs) + clear_attrs = set(clear_attrs) + options = {} + active_mask = 0 + + for header, attr_name, bit in _OTEL_ATTRIBUTES: + attr = getattr(ibm_db, attr_name) + if attr in explicit_attrs: + continue + value = headers.get(header) + if value: + options[attr] = value + active_mask |= bit + elif attr in clear_attrs: + options[attr] = "" + + if not options: + return active_mask + + try: + result = ibm_db.set_option(connection, options, 1) + except Exception as error: + if "does not support OpenTelemetry connection attributes" in str(error): + return None + raise + + if not result: + raise RuntimeError("ibm_db.set_option failed to set OpenTelemetry attributes") + return active_mask diff --git a/ibm_db_tests/demo_otel_context.py b/ibm_db_tests/demo_otel_context.py new file mode 100644 index 00000000..4878fcce --- /dev/null +++ b/ibm_db_tests/demo_otel_context.py @@ -0,0 +1,282 @@ +"""Demonstrate automatic W3C context propagation through ibm_db.""" + +import os +import secrets + +import ibm_db +import ibm_db_dbi +from opentelemetry import baggage, context, trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.baggage.propagation import W3CBaggagePropagator +from opentelemetry.trace import ( + NonRecordingSpan, + SpanContext, + TraceFlags, + TraceState, + set_span_in_context, +) +from opentelemetry.trace.propagation.tracecontext import ( + TraceContextTextMapPropagator, +) + + +CONNECTION_STRING = os.environ.get( + "DB2_CONNSTR", + "DATABASE=YOUR_DB;HOSTNAME=YOUR_HOST;PORT=50000;" + "UID=YOUR_USER;PWD=YOUR_PASSWORD", +) + +FIXED_W3C_HEADERS = { + "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", + "tracestate": "congo=t61rcWkgMzE", + "baggage": "userId=alice,serverNode=DF28", +} + +CHANGED_W3C_HEADERS = { + "traceparent": "00-fedcba9876543210fedcba9876543210-fedcba9876543210-00", + "tracestate": "congo=changed", + "baggage": "userId=bob,serverNode=DF29", +} + +OTEL_ATTRIBUTES = ( + ("traceparent", ibm_db.SQL_ATTR_OTEL_TRACEPARENT), + ("tracestate", ibm_db.SQL_ATTR_OTEL_TRACESTATE), + ("baggage", ibm_db.SQL_ATTR_OTEL_BAGGAGE), +) + + +def make_parent_context(): + parent_span = SpanContext( + trace_id=secrets.randbits(128) or 1, + span_id=secrets.randbits(64) or 1, + is_remote=True, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=TraceState([("demo", "ibm_db")]), + ) + parent_context = set_span_in_context(NonRecordingSpan(parent_span)) + return baggage.set_baggage( + "demo", "ibm_db_otel", context=parent_context + ) + + +def current_w3c_headers(): + headers = {} + TraceContextTextMapPropagator().inject(headers) + W3CBaggagePropagator().inject(headers) + return headers + + +def show_connection_attributes(conn, expected_headers, dbi=False): + print(" ibm_db connection attributes:") + for header, attribute in OTEL_ATTRIBUTES: + value = ( + conn.get_option(attribute) + if dbi + else ibm_db.get_option(conn, attribute, 1) + ) + if value is False: + print(f" {header}: unavailable from this CLI driver") + else: + print(f" {header}: {value}") + if expected_headers is not None and value != expected_headers.get( + header, "" + ): + raise AssertionError( + f"{header} did not match the expected W3C context" + ) + + +def run_native_query(conn, label, sql, expected_headers): + stmt = ibm_db.exec_immediate(conn, sql) + if not stmt: + raise RuntimeError(f"Failed to execute query: {label}") + try: + row = ibm_db.fetch_tuple(stmt) + print(f"\n{label}: {row[0] if row else ''}") + finally: + ibm_db.free_result(stmt) + show_connection_attributes(conn, expected_headers) + + +def run_prepared_query(conn, label, sql, parameters, expected_headers): + stmt = ibm_db.prepare(conn, sql) + if not stmt: + raise RuntimeError(f"Failed to prepare query: {label}") + try: + if not ibm_db.execute(stmt, parameters): + raise RuntimeError(f"Failed to execute prepared query: {label}") + row = ibm_db.fetch_tuple(stmt) + print(f"\n{label}: {row[0] if row else ''}") + finally: + ibm_db.free_result(stmt) + show_connection_attributes(conn, expected_headers) + + +def run_carrier_case(conn, case_number, headers): + extracted_context = TraceContextTextMapPropagator().extract(headers) + extracted_context = W3CBaggagePropagator().extract( + headers, extracted_context + ) + token = context.attach(extracted_context) + try: + description = ( + "fixed carrier" + if case_number == 4 + else "changed fixed carrier on the same connection" + ) + print( + f"\nCase {case_number}: W3C context extracted from " + f"the {description}:" + ) + for header, _attribute in OTEL_ATTRIBUTES: + print(f" {header}: {headers[header]}") + run_native_query( + conn, + f"case {case_number} query; current schema", + "SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1", + headers, + ) + finally: + context.detach(token) + + +def main(): + if "YOUR_" in CONNECTION_STRING: + raise SystemExit( + "Set DB2_CONNSTR or edit CONNECTION_STRING with your Db2 connection." + ) + + trace.set_tracer_provider(TracerProvider()) + tracer = trace.get_tracer(__name__) + parent_context_token = context.attach(make_parent_context()) + conn = None + preexisting_conn = None + dbi_conn = None + cursor = None + try: + with tracer.start_as_current_span("ibm_db.otel_demo"): + active_headers = current_w3c_headers() + print( + "\nCase 1: automatically propagate the active span when " + "connecting:" + ) + for header, _attribute in OTEL_ATTRIBUTES: + print(f" {header}: {active_headers.get(header, '')}") + + conn = ibm_db.connect(CONNECTION_STRING, "", "") + print(" attributes immediately after connect, before SQL:") + show_connection_attributes(conn, active_headers) + + print("\nCase 2: propagate context before ibm_db.exec_immediate:") + run_native_query( + conn, + "current schema", + "SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1", + active_headers, + ) + + print("\nCase 3: propagate context before prepared ibm_db.execute:") + run_prepared_query( + conn, + "prepared parameter", + "SELECT CAST(? AS VARCHAR(32)) FROM SYSIBM.SYSDUMMY1", + ("prepared execution",), + active_headers, + ) + + for case_number, headers in ( + (4, FIXED_W3C_HEADERS), + (5, CHANGED_W3C_HEADERS), + ): + run_carrier_case(conn, case_number, headers) + + print("\nCase 6: automatically propagate through ibm_db_dbi:") + with tracer.start_as_current_span("ibm_db_dbi.otel_demo"): + dbi_headers = current_w3c_headers() + dbi_conn = ibm_db_dbi.connect(CONNECTION_STRING, "", "") + cursor = dbi_conn.cursor() + cursor.execute("SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1") + row = cursor.fetchone() + print(f" current schema: {row[0] if row else ''}") + show_connection_attributes(dbi_conn, dbi_headers, dbi=True) + + print( + "\nCase 7: update context on a connection opened before the span:" + ) + empty_token = context.attach(context.Context()) + try: + preexisting_conn = ibm_db.connect(CONNECTION_STRING, "", "") + finally: + context.detach(empty_token) + with tracer.start_as_current_span("ibm_db.otel_existing_connection"): + existing_headers = current_w3c_headers() + run_native_query( + preexisting_conn, + "query on preexisting connection", + "SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1", + existing_headers, + ) + + print("\nCase 8: clear automatically propagated values without a span:") + empty_token = context.attach(context.Context()) + try: + run_native_query( + conn, + "query without active context", + "SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1", + {"traceparent": "", "tracestate": "", "baggage": ""}, + ) + finally: + context.detach(empty_token) + + print("\nCase 9: explicit connection attributes take precedence:") + explicit_headers = { + "traceparent": ( + "00-11111111111111111111111111111111-2222222222222222-01" + ), + "tracestate": "demo=manual", + "baggage": "demo=manual", + } + explicit_attributes = { + attribute: explicit_headers[header] + for header, attribute in OTEL_ATTRIBUTES + } + with tracer.start_as_current_span("ibm_db.otel_explicit"): + try: + result = ibm_db.set_option(conn, explicit_attributes, 1) + except Exception as error: + if "does not support OpenTelemetry connection attributes" not in str( + error + ): + raise + print( + " Explicit attributes are unsupported by this CLI driver; " + "automatic propagation remains optional." + ) + else: + if not result: + raise RuntimeError( + "ibm_db.set_option failed to set explicit OTEL attributes" + ) + run_native_query( + conn, + "query with explicit attributes", + "SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1", + explicit_headers, + ) + finally: + try: + if cursor is not None: + cursor.close() + if dbi_conn is not None: + dbi_conn.close() + if preexisting_conn: + ibm_db.close(preexisting_conn) + if conn: + ibm_db.close(conn) + finally: + context.detach(parent_context_token) + + +if __name__ == "__main__": + main() diff --git a/ibm_db_tests/test_331_OtelConnectionAttrs.py b/ibm_db_tests/test_331_OtelConnectionAttrs.py new file mode 100644 index 00000000..a7651c10 --- /dev/null +++ b/ibm_db_tests/test_331_OtelConnectionAttrs.py @@ -0,0 +1,274 @@ +import os +import unittest +from unittest.mock import patch + +import ibm_db +import ibm_db_dbi + + +TRACEPARENT = "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01" +TRACESTATE = "congo=t61rcWkgMzE" +BAGGAGE = "userId=alice,serverNode=DF28" +UPDATED_TRACEPARENT = "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01" +UPDATED_TRACESTATE = "rojo=00f067aa0ba902b7" +UPDATED_BAGGAGE = "userId=bob,serverNode=DF29" +OTEL_ATTRS = { + ibm_db.SQL_ATTR_OTEL_TRACEPARENT: TRACEPARENT, + ibm_db.SQL_ATTR_OTEL_TRACESTATE: TRACESTATE, + ibm_db.SQL_ATTR_OTEL_BAGGAGE: BAGGAGE, +} +UPDATED_OTEL_ATTRS = { + ibm_db.SQL_ATTR_OTEL_TRACEPARENT: UPDATED_TRACEPARENT, + ibm_db.SQL_ATTR_OTEL_TRACESTATE: UPDATED_TRACESTATE, + ibm_db.SQL_ATTR_OTEL_BAGGAGE: UPDATED_BAGGAGE, +} + + +class IbmDbTestCase(unittest.TestCase): + + def test_331_OtelContextHelper(self): + try: + from opentelemetry import context + from opentelemetry.baggage.propagation import W3CBaggagePropagator + from opentelemetry.trace.propagation.tracecontext import ( + TraceContextTextMapPropagator, + ) + except ImportError: + self.skipTest("OpenTelemetry API packages are not installed") + + from ibm_db_otel import apply_otel_context + + headers = { + "traceparent": TRACEPARENT, + "tracestate": TRACESTATE, + "baggage": BAGGAGE, + } + active = TraceContextTextMapPropagator().extract(headers) + active = W3CBaggagePropagator().extract(headers, active) + token = context.attach(active) + try: + with patch.object(ibm_db, "set_option", return_value=True) as setter: + mask = apply_otel_context(object()) + self.assertEqual(mask, 7) + self.assertEqual(setter.call_args.args[1], OTEL_ATTRS) + + with patch.object(ibm_db, "set_option", return_value=True) as setter: + mask = apply_otel_context( + object(), + explicit_attrs=(ibm_db.SQL_ATTR_OTEL_TRACEPARENT,), + ) + self.assertEqual(mask, 6) + self.assertEqual( + setter.call_args.args[1], + { + ibm_db.SQL_ATTR_OTEL_TRACESTATE: TRACESTATE, + ibm_db.SQL_ATTR_OTEL_BAGGAGE: BAGGAGE, + }, + ) + finally: + context.detach(token) + + token = context.attach(context.Context()) + try: + with patch.object(ibm_db, "set_option", return_value=True) as setter: + mask = apply_otel_context( + object(), + explicit_attrs=(ibm_db.SQL_ATTR_OTEL_TRACEPARENT,), + clear_attrs=tuple(OTEL_ATTRS), + ) + self.assertEqual(mask, 0) + self.assertEqual( + setter.call_args.args[1], + { + ibm_db.SQL_ATTR_OTEL_TRACESTATE: "", + ibm_db.SQL_ATTR_OTEL_BAGGAGE: "", + }, + ) + finally: + context.detach(token) + + def test_331_OtelConnectionAttrs(self): + self.assertEqual(ibm_db.SQL_ATTR_OTEL_TRACEPARENT, 3056) + self.assertEqual(ibm_db.SQL_ATTR_OTEL_TRACESTATE, 3057) + self.assertEqual(ibm_db.SQL_ATTR_OTEL_BAGGAGE, 3058) + self.assertEqual(ibm_db_dbi.SQL_ATTR_OTEL_TRACEPARENT, 3056) + self.assertEqual(ibm_db_dbi.SQL_ATTR_OTEL_TRACESTATE, 3057) + self.assertEqual(ibm_db_dbi.SQL_ATTR_OTEL_BAGGAGE, 3058) + self._test_manual_and_automatic_context() + + def _test_manual_and_automatic_context(self): + try: + from opentelemetry import context + from opentelemetry.baggage.propagation import W3CBaggagePropagator + from opentelemetry.trace.propagation.tracecontext import ( + TraceContextTextMapPropagator, + ) + except ImportError: + context = None + W3CBaggagePropagator = None + TraceContextTextMapPropagator = None + + db2_connstr = os.environ.get("DB2_CONNSTR") + if not db2_connstr: + import config + + def connect_ibm_db(options=None): + if db2_connstr: + args = (db2_connstr, "", "") + else: + args = (config.database, config.user, config.password) + if options is None: + return ibm_db.connect(*args) + return ibm_db.connect(*args, options) + + def connect_ibm_db_dbi(conn_options=None): + if db2_connstr: + args = (db2_connstr, "", "") + else: + args = (config.database, config.user, config.password) + return ibm_db_dbi.connect(*args, conn_options=conn_options) + + try: + probe = connect_ibm_db() + except Exception: + if db2_connstr: + self.fail("Db2 connection failed using DB2_CONNSTR") + self.skipTest("Db2 database connection is unavailable") + else: + ibm_db.close(probe) + + otel_attrs_supported = False + conn = None + try: + conn = connect_ibm_db(OTEL_ATTRS) + except Exception as error: + self.assertIn( + "clidriver used with ibm_db does not support OpenTelemetry", + str(error), + ) + else: + self.assertTrue(conn) + self.assertTrue(ibm_db.set_option(conn, OTEL_ATTRS, 1)) + for attr, expected in OTEL_ATTRS.items(): + self.assertEqual(ibm_db.get_option(conn, attr, 1), expected) + otel_attrs_supported = True + finally: + if conn: + ibm_db.close(conn) + + try: + dbi_manual_conn = connect_ibm_db_dbi( + conn_options=OTEL_ATTRS.copy() + ) + except Exception as error: + self.assertIn( + "clidriver used with ibm_db does not support OpenTelemetry", + str(error), + ) + else: + try: + for attr, expected in OTEL_ATTRS.items(): + self.assertEqual(dbi_manual_conn.get_option(attr), expected) + otel_attrs_supported = True + finally: + dbi_manual_conn.close() + + conn = connect_ibm_db() + self.assertTrue(conn) + dbi_conn = connect_ibm_db_dbi() + self.assertTrue(dbi_conn) + token = None + cursor = None + try: + if context is not None: + headers = { + "traceparent": TRACEPARENT, + "tracestate": TRACESTATE, + "baggage": BAGGAGE, + } + otel_context = TraceContextTextMapPropagator().extract(headers) + otel_context = W3CBaggagePropagator().extract( + headers, otel_context + ) + token = context.attach(otel_context) + + context_conn = connect_ibm_db() + try: + for attr, expected in OTEL_ATTRS.items(): + actual = ibm_db.get_option(context_conn, attr, 1) + if otel_attrs_supported or actual is not False: + self.assertEqual(actual, expected) + finally: + ibm_db.close(context_conn) + + stmt = ibm_db.exec_immediate( + conn, "SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1" + ) + self.assertTrue(stmt) + if context is not None: + for attr, expected in OTEL_ATTRS.items(): + actual = ibm_db.get_option(conn, attr, 1) + if otel_attrs_supported or actual is not False: + self.assertEqual(actual, expected) + + cursor = dbi_conn.cursor() + cursor.execute("SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1") + if context is not None: + for attr, expected in OTEL_ATTRS.items(): + actual = dbi_conn.get_option(attr) + if otel_attrs_supported or actual is not False: + self.assertEqual(actual, expected) + + context.detach(token) + token = None + updated_headers = { + "traceparent": UPDATED_TRACEPARENT, + "tracestate": UPDATED_TRACESTATE, + "baggage": UPDATED_BAGGAGE, + } + updated_context = TraceContextTextMapPropagator().extract( + updated_headers + ) + updated_context = W3CBaggagePropagator().extract( + updated_headers, updated_context + ) + token = context.attach(updated_context) + + stmt = ibm_db.exec_immediate( + conn, "SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1" + ) + self.assertTrue(stmt) + for attr, expected in UPDATED_OTEL_ATTRS.items(): + actual = ibm_db.get_option(conn, attr, 1) + if otel_attrs_supported or actual is not False: + self.assertEqual(actual, expected) + + cursor.execute( + "SELECT CURRENT_SCHEMA FROM SYSIBM.SYSDUMMY1" + ) + for attr, expected in UPDATED_OTEL_ATTRS.items(): + actual = dbi_conn.get_option(attr) + if otel_attrs_supported or actual is not False: + self.assertEqual(actual, expected) + finally: + if token is not None: + context.detach(token) + if cursor is not None: + cursor.close() + dbi_conn.close() + ibm_db.close(conn) + + +if __name__ == "__main__": + unittest.main() + + +#__END__ +#__LUW_EXPECTED__ +#test_331_OtelConnectionAttrs (test_331_OtelConnectionAttrs.IbmDbTestCase.test_331_OtelConnectionAttrs) ... ok +#__ZOS_EXPECTED__ +#test_331_OtelConnectionAttrs (test_331_OtelConnectionAttrs.IbmDbTestCase.test_331_OtelConnectionAttrs) ... ok +#__SYSTEMI_EXPECTED__ +# +#__IDS_EXPECTED__ +# diff --git a/setup.py b/setup.py index f78ecb50..e93795ce 100644 --- a/setup.py +++ b/setup.py @@ -497,7 +497,7 @@ def print_exception( e, url): (get_python_lib(), ['./CHANGES.md']), (get_python_lib(), ['./LICENSE'])] -modules = ['ibm_db_dbi', 'testfunctions', 'ibmdb_tests', 'ibm_db_ctx', '_ibm_db_register_dll'] +modules = ['ibm_db_dbi', 'ibm_db_otel', 'testfunctions', 'ibmdb_tests', 'ibm_db_ctx', '_ibm_db_register_dll'] if 'zos' == sys.platform: ext_modules = _ext_modules(os.path.join(os.getcwd(), include_dir), library, ibm_db_lib, ibm_db_lib_runtime)