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.Ipc;
 using Apache.Hive.Service.Rpc.Thrift;
 using Thrift.Protocol;
@@ -66,14 +71,21 @@ internal class DatabricksConnection : SparkHttpConnection
         private string _traceParentHeaderName = "traceparent";
         private bool _traceStateEnabled = false;
 
+        // Telemetry
+        private TelemetryHelper? _telemetryHelper;
+        private DatabricksActivityListener _databricksActivityListener;
+        private Guid _guid;
+
         // Default namespace
         private TNamespace? _defaultNamespace;
 
         private HttpClient? _authHttpClient;
 
         public DatabricksConnection(IReadOnlyDictionary<string, string> properties) : base(properties)
         {
+            _guid = Guid.NewGuid();
             ValidateProperties();
+            _databricksActivityListener = new DatabricksActivityListener(_telemetryHelper, this.AssemblyName, _guid);
         }
 
         protected override TCLIService.IAsync CreateTCLIServiceClient(TProtocol protocol)
@@ -264,6 +276,43 @@ private void ValidateProperties()
                 // Default QueryTimeSeconds in Hive2Connection is only 60s, which is too small for lots of long running query
                 QueryTimeoutSeconds = DefaultQueryTimeSeconds;
             }
+
+            //Telemetry
+            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(hostDetails.HostUrl, token);
+            _telemetryHelper.SetParameters(connectionParams, clientContext);
         }
 
         /// <summary>
@@ -403,6 +452,7 @@ protected override HttpMessageHandler CreateHttpHandler()
                 }
             }
 
+            _telemetryHelper?.InitializeTelemetryClient(_authHttpClient);
             return baseHandler;
         }
 
@@ -737,6 +787,16 @@ protected override void Dispose(bool disposing)
             if (disposing)
             {
                 _authHttpClient?.Dispose();
+                _databricksActivityListener?.Dispose();
+
+                try
+                {
+                    _telemetryHelper?.ForceFlushAsync().Wait(TimeSpan.FromSeconds(2));
+                }
+                catch (Exception ex)
+                {
+                    Debug.WriteLine($"Failed to flush telemetry on dispose: {ex.Message}");
+                }
             }
             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 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
+{
+    public class DatabricksActivityListener : IDisposable
+    {
+        private readonly ActivityListener _activityListener;
+        private TelemetryHelper? _telemetryHelper;
+
+        public DatabricksActivityListener(TelemetryHelper? telemetryHelper, string sourceName)
+        {
+            this._telemetryHelper = telemetryHelper;
+            this._activityListener = new ActivityListener
+            {
+                ShouldListenTo = (activitySource) => activitySource.Name == sourceName,
+                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"))
+            {
+                var sqlExecutionEvent = new SqlExecutionEvent();
+                var operationDetail = new OperationDetail();
+                operationDetail.OperationType = Util.StringToOperationType("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;
+
+public class DatabricksConnectionConfig
+{
+    public static readonly int MAX_BATCH_SIZE = 1;
+    public static readonly int FLUSH_INTERVAL_MILLIS = 300000; // 5 minutes
+}

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

@@ -0,0 +1,3 @@
+# Databricks Telemetry
+
+Uses `DatabricksActivityListener.cs` to record events.

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

@@ -0,0 +1,124 @@
+/*
+* 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
+{
+    public 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" : $"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;
+            }
+        }
+
+        /// <summary>
+        /// Sends a single telemetry event asynchronously
+        /// </summary>
+        /// <param name="telemetryEvent">Single telemetry event to send</param>
+        /// <returns>Task representing the async operation</returns>
+        public async Task<bool> SendTelemetryAsync(TelemetryFrontendLog telemetryEvent)
+        {
+            if (string.IsNullOrEmpty(_telemetryUrl))
+            {
+                return false;
+            }
+
+            try
+            {
+                var request = new HttpRequestMessage(HttpMethod.Post, _telemetryUrl);
+
+                // Serialize the event to JSON
+                var telemetryRequest = new TelemetryRequest();
+                telemetryRequest.UploadTime = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
+                telemetryRequest.ProtoLogs = new List<string> { JsonSerializer.Serialize(telemetryEvent) };
+                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;
+            }
+        }
+    }
+}