feat(csharp/src/Drivers/Databricks): Implement Telemetry Reporting for Databricks

csharp/src/Drivers/Databricks/DatabricksConnection.cs

@@ -28,6 +28,11 @@
 using Apache.Arrow.Adbc.Drivers.Apache.Hive2.Client;
 using Apache.Arrow.Adbc.Drivers.Apache.Spark;
 using Apache.Arrow.Adbc.Drivers.Databricks.Auth;
+using Apache.Arrow.Adbc.Drivers.Databricks.CloudFetch;
+using Apache.Arrow.Adbc.Drivers.Databricks.Telemetry;
+using Apache.Arrow.Adbc.Drivers.Databricks.Telemetry.Enums;
+using Apache.Arrow.Adbc.Drivers.Databricks.Telemetry.Model;
+using Apache.Arrow.Adbc.Tracing;
 using Apache.Arrow.Adbc.Drivers.Databricks.Reader;
 using Apache.Arrow.Ipc;
 using Apache.Hive.Service.Rpc.Thrift;
@@ -67,6 +72,11 @@
         private string _traceParentHeaderName = "traceparent";
         private bool _traceStateEnabled = false;
 
+        // Telemetry
+        private TelemetryHelper? _telemetryHelper;
+        private DatabricksActivityListener _databricksActivityListener;
+        private Guid _guid;
+
         // Identity federation client ID for token exchange
         private string? _identityFederationClientId;
 
@@ -77,15 +87,17 @@
 
         public DatabricksConnection(IReadOnlyDictionary<string, string> properties) : base(properties)
         {
+            _guid = Guid.NewGuid();
             ValidateProperties();
+            InitializeTelemetry();
+            _databricksActivityListener = new DatabricksActivityListener(_telemetryHelper, this.AssemblyName, _guid);
         }
 
         public override IEnumerable<KeyValuePair<string, object?>>? GetActivitySourceTags(IReadOnlyDictionary<string, string> properties)
         {
             IEnumerable<KeyValuePair<string, object?>>? tags = base.GetActivitySourceTags(properties);
-            // TODO: Add any additional tags specific to Databricks connection
-            //tags ??= [];
-            //tags.Concat([new("key", "value")]);
+            tags ??= [];
+            tags = tags.Concat([new(Util.GetDatabricksTag(),_guid)]);
             return tags;
         }
 
@@ -284,6 +296,45 @@
             }
         }
 
+        private void InitializeTelemetry()
+        {
+            var connectionParams = new DriverConnectionParameters();
+            var hostDetails = new HostDetails();
+            var clientContext = new ClientContext();
+            string? token = null;
+
+            if (Properties.TryGetValue(SparkParameters.AuthType, out string? authType))
+            {
+                connectionParams.AuthMech = Util.StringToAuthMech(authType);
+            }
+
+            if (Properties.TryGetValue(SparkParameters.HostName, out string? host))
+            {
+                hostDetails.HostUrl = host;
+            }
+            if (Properties.TryGetValue(SparkParameters.Port, out string? port))
+            {
+                hostDetails.Port = Int32.Parse(port);
+            }
+            connectionParams.HostInfo = hostDetails;
+
+            if (Properties.TryGetValue(SparkParameters.UserAgentEntry, out string? userAgent))
+            {
+                clientContext.UserAgent = userAgent;
+            }
+
+            if(Properties.TryGetValue(SparkParameters.AccessToken, out string? accessToken))
+            {
+                token = accessToken;
+            }   
+            else if(Properties.TryGetValue(SparkParameters.Token, out string? accesstoken))
+            {
+                token = accesstoken;
+            }
+            _telemetryHelper = new TelemetryHelper(token);
+            _telemetryHelper.SetParameters(connectionParams, clientContext);
+        }
+
         /// <summary>
         /// Gets whether server side properties should be applied using queries.
         /// </summary>
@@ -424,6 +475,7 @@
                 }
             }
 
+            _telemetryHelper?.InitializeTelemetryClient(_authHttpClient);
             return baseHandler;
         }
 
@@ -724,6 +776,15 @@
         {
             if (disposing)
             {
+                _databricksActivityListener?.Dispose();
+                try
+                {
+                    _telemetryHelper?.ForceFlushAsync().Wait(TimeSpan.FromSeconds(2));
+                }
+                catch (Exception ex)
+                {
+                    Debug.WriteLine($"Failed to flush telemetry on dispose: {ex.Message}");
+                }
                 _authHttpClient?.Dispose();
             }
             base.Dispose(disposing);

csharp/src/Drivers/Databricks/Telemetry/DatabricksActivityListener.cs

@@ -0,0 +1,68 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements.  See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You under the Apache License, Version 2.0
+* (the "License"); you may not use this file except in compliance with
+* the License.  You may obtain a copy of the License at
+*
+*    http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System;
+using System.Diagnostics;
+using System.IO;
+using System.Linq;
+using Apache.Arrow.Adbc.Tracing;
+using Apache.Arrow.Adbc.Drivers.Apache;
+using Apache.Arrow.Adbc.Drivers.Databricks.Telemetry.Model;
+using Apache.Arrow.Adbc.Drivers.Databricks.Telemetry.Enums;
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks.Telemetry
+{
+    internal class DatabricksActivityListener : IDisposable
+    {
+        private readonly ActivityListener _activityListener;
+        private TelemetryHelper? _telemetryHelper;
+
+        public DatabricksActivityListener(TelemetryHelper? telemetryHelper, string sourceName, Guid guid)
+        {
+            this._telemetryHelper = telemetryHelper;
+            this._activityListener = new ActivityListener
+            {
+                ShouldListenTo = (activitySource) => activitySource.Tags?.Any(kvp => kvp.Key == Util.GetDatabricksTag() && kvp.Value?.Equals(guid) == true) == true,
+                Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData,
+                ActivityStopped = OnActivityStopped,
+            };
+            ActivitySource.AddActivityListener(_activityListener);
+        }
+
+        private void OnActivityStopped(Activity activity)
+        {       
+            if(_telemetryHelper == null)
+            {
+                return;
+            }
+
+            if(activity.OperationName?.EndsWith("ExecuteStatementAsync") == true)
+            {
+                var sqlExecutionEvent = new SqlExecutionEvent();
+                var operationDetail = new OperationDetail();
+                operationDetail.OperationType = OperationType.EXECUTE_STATEMENT_ASYNC;
+                sqlExecutionEvent.OperationDetail = operationDetail;
+                _telemetryHelper.AddSqlExecutionEvent(sqlExecutionEvent, Convert.ToInt64(activity.Duration.TotalMilliseconds));
+            }
+        }
+
+        public void Dispose()
+        {   
+            this._activityListener.Dispose();
+        }
+    }
+}

csharp/src/Drivers/Databricks/Telemetry/DatabricksConnectionConfig.cs

@@ -0,0 +1,24 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements.  See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You under the Apache License, Version 2.0
+* (the "License"); you may not use this file except in compliance with
+* the License.  You may obtain a copy of the License at
+*
+*    http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks.Telemetry;
+
+internal class DatabricksConnectionConfig
+{
+    public static readonly int MAX_BATCH_SIZE = 200;
+    public static readonly int FLUSH_INTERVAL_MILLIS = 300000; // 5 minutes
+}

csharp/src/Drivers/Databricks/Telemetry/README.md

@@ -0,0 +1,7 @@
+# Databricks Telemetry
+
+Uses `DatabricksActivityListener.cs` to listen and record events to `TelemetryHelper.cs` using ActivityListener.
+
+`TelemetryHelper.cs` stores telemetry logs and telemetry parameters.
+
+`TelemtryClient.cs` sends logs to `/telemetry` or `/telemetry-unauth` depending on authentication level.

csharp/src/Drivers/Databricks/Telemetry/TelemetryClient.cs

@@ -0,0 +1,83 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements.  See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You under the Apache License, Version 2.0
+* (the "License"); you may not use this file except in compliance with
+* the License.  You may obtain a copy of the License at
+*
+*    http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Net.Http;
+using System.Net.Http.Headers;
+using System.Text.Json;
+using System.Threading.Tasks;
+using Apache.Arrow.Adbc.Drivers.Databricks.Telemetry.Model;
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks.Telemetry
+{
+    internal class TelemetryClient
+    {
+        private readonly HttpClient _httpClient;
+        private readonly string? _telemetryUrl;
+        private readonly string? _accessToken;
+
+        public TelemetryClient(HttpClient httpClient, string? hostUrl, string? accessToken)
+        {
+            _httpClient = httpClient;
+            _accessToken = accessToken;
+            _telemetryUrl = !string.IsNullOrEmpty(hostUrl) ? accessToken != null ? $"https://{hostUrl}/telemetry-ext" : $"https://{hostUrl}/telemetry-unauth" : null;
+        }
+
+        /// <summary>
+        /// Sends a batch of telemetry events asynchronously
+        /// </summary>
+        /// <param name="telemetryBatch">List of telemetry events to send</param>
+        /// <returns>Task representing the async operation</returns>
+        public async Task<bool> SendTelemetryBatchAsync(List<TelemetryFrontendLog> telemetryBatch)
+        {
+            if (string.IsNullOrEmpty(_telemetryUrl) || telemetryBatch.Count == 0)
+            {
+                return false;
+            }
+
+            try
+            {
+                var request = new HttpRequestMessage(HttpMethod.Post, _telemetryUrl);
+
+                // Serialize the batch to JSON
+                var telemetryRequest = new TelemetryRequest();
+                telemetryRequest.UploadTime = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
+                telemetryRequest.ProtoLogs = telemetryBatch.Select(x => JsonSerializer.Serialize(x)).ToList();
+                request.Content = new StringContent(JsonSerializer.Serialize(telemetryRequest));
+
+                // Set headers
+                request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json"));
+                request.Content.Headers.ContentType = new MediaTypeHeaderValue("application/json");
+                if(_accessToken != null)
+                {
+                    request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _accessToken);
+                }
+
+                var response = await _httpClient.SendAsync(request);
+                return response.IsSuccessStatusCode;
+            }
+            catch (Exception ex)
+            {
+                // Log the exception but don't throw to prevent telemetry failures from affecting main functionality
+                System.Diagnostics.Debug.WriteLine($"Failed to send telemetry: {ex.Message}");
+                return false;
+            }
+        }
+    }
+}