Skip to content

Conversation

Copy link
Contributor

Copilot AI commented Jan 21, 2026

AzureSearchDriver inherited EFFECTIVE_ENV="MSSentinel" from its parent AzureMonitorDriver, causing the QueryProvider to report the wrong environment name. Additionally, the inherited query() method checked for _query_client (used by parent) instead of _auth_header (used by this driver), causing "Workspace not connected" errors even after successful connection.

Changes

  • Override environment in __init__: Set EFFECTIVE_ENV to "MSSentinelSearch" after parent initialization
  • Override query filter: Include "MSSentinelSearch" in supported environments tuple
  • Extract connection check: Add _ensure_connected() helper to check _auth_header is not None
  • Override query method: Use _ensure_connected() instead of parent's _query_client check

Example

# Before: failed with "Workspace not connected" 
qry_prov = QueryProvider('MSSentinelSearch')
qry_prov.connect(workspace='BasicLogs')
print(qry_prov.environment)  # Showed "MSSentinel" (wrong)
df = qry_prov.exec_query('SyslogBasic_CL | take 1', start=start, end=end)  # Failed

# After: works correctly
qry_prov = QueryProvider('MSSentinelSearch')
qry_prov.connect(workspace='BasicLogs')  
print(qry_prov.environment)  # Shows "MSSentinelSearch" (correct)
df = qry_prov.exec_query('SyslogBasic_CL | take 1', start=start, end=end)  # Succeeds
Original prompt

This section details on the original issue you should resolve

<issue_title>[Bug]: New experimenal MSSentinelSearch data provider doesn't correctly use the AzureSearchDriver</issue_title>
<issue_description>Describe the bug

The MSSentinelSearch query provider / data environment seems to get confused between using the MSSentinel vs MSSentinelSearch data environments and fails to correctly connect the AzureSearchDriver.

To Reproduce

Steps to reproduce the behavior:

  1. Clone from this git repo and checkout main in order to test PR
    Ianhelle/az monitor search driver 2025 02 05 #825 included in main.
  2. Create an editable venv from the source and activate.
  3. Configure msticpyconfig.yaml with a 'Sentinel' workspace that inlcudes a basic table.
  4. Run a test script with the AzureSearchDriver on a table with the 'basic' plan.
import datetime

# Set debug logging
import logging
logging.basicConfig(level=logging.DEBUG)

# Inherit log level
import msticpy
print(f'msticpy version: {msticpy.__version__}')

# Config
msticpy.init_notebook()

# ws_config = msticpy.common.wsconfig.WorkspaceConfig(workspace="MyWorkspace")
# print(f'Workspace config: {ws_config}')
#qry_prov_basic_search = msticpy.QueryProvider(data_environment='MSSentinelSearch', ws_config=ws_config, workspace='BasicLogs')
qry_prov_basic_search = msticpy.QueryProvider('MSSentinelSearch')
qry_prov_basic_search.connect(workspace='BasicLogs')
print(f'Query provider driver: {qry_prov_basic_search.driver_class}')
print(f'Query provider environment: {qry_prov_basic_search.environment}')
print(f'Query provider connections: {qry_prov_basic_search.list_connections()}')

# Prep a small time range to limit basic logs query costs
lookback_period = datetime.timedelta(hours=1)
ingest_grace_period = datetime.timedelta(minutes=15)
end = datetime.datetime.now(datetime.timezone.utc) - ingest_grace_period
start = end - lookback_period
print(f'Start: {start}, End: {end}')

# Test query
df = qry_prov_basic_search.exec_query('SyslogBasic_CL | take 1', start=start, end=end)
print(df)

Expected behavior

AzureSearchDriver is connected and used with the corresponding MSSentinelSearch data environment.

Screenshots and/or Traceback

INFO:msticpy.data.drivers.azure_monitor_driver:AzureMonitorDriver loaded. connect_str  None, kwargs: {'data_environment': <DataEnvironment.MSSentinelSearch: 25>}
INFO:msticpy.data.core.data_providers:Using data environment MSSentinel
INFO:msticpy.data.core.data_providers:Driver class: AzureSearchDriver
...
INFO:msticpy.data.core.data_providers:Calling connect on driver
INFO:msticpy.data.drivers.azure_monitor_driver:WorkspaceConfig created from workspace name BasicLogs
...
INFO:msticpy.data.drivers.azure_monitor_driver:WorkspaceConfig created from workspace name BasicLogs
INFO:msticpy.data.drivers.azure_search_driver:Created HTTP-based query client using /search endpoint.
connected
INFO:msticpy.data.core.data_providers:Adding query pivot functions
Query provider driver: <class 'msticpy.data.drivers.azure_search_driver.AzureSearchDriver'>
Query provider environment: MSSentinel
Query provider connections: ['Default: BasicLogs']
Start: 2025-02-16 19:20:23.193644+00:00, End: 2025-02-16 20:20:23.193644+00:00
INFO:msticpy.data.core.query_provider_connections_mixin:Executing query 'SyslogBasic_CL | take 1...'
DEBUG:msticpy.data.core.query_provider_connections_mixin:Full query: SyslogBasic_CL | take 1
DEBUG:msticpy.data.core.query_provider_connections_mixin:Query options: {'start': datetime.datetime(2025, 2, 16, 19, 20, 23, 193644, tzinfo=datetime.timezone.utc), 'end': datetime.datetime(2025, 2, 16, 20, 20, 23, 193644, tzinfo=datetime.timezone.utc)}
Traceback (most recent call last):
  File "/home/***/SecurityNotebooks/demo/_bugs/msticpy/msticpy_cannot_query_azure_monitor_basic_logs/test_data_prov_ms_sentinel_search.py", line 39, in <module>
    df = qry_prov_basic_search.exec_query('SyslogBasic_CL | take 1', start=start, end=end)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/***/SecurityNotebooks/demo/_bugs/msticpy/msticpy_cannot_query_azure_monitor_basic_logs/src/msticpy/msticpy/data/core/query_provider_connections_mixin.py", line 99, in exec_query
    return self._query_provider.query(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/***/SecurityNotebooks/demo/_bugs/msticpy/msticpy_cannot_query_azure_monitor_basic_logs/src/msticpy/msticpy/data/drivers/azure_monitor_driver.py", line 294, in query
    raise MsticpyNotConnectedError(
msticpy.common.exceptions.MsticpyNotConnectedError: ('Workspace not connected.', 'Please run connect() to connect to the workspace', 'before running a query.', 'https://msticpy.readthedocs.io/en/latest/data_acquisition/DataProv-MSSentinel.html')

Environment (please complete the following information):

  • Python Version: 3.111
  • OS: Ubuntu
  • Python environment: editable venv from cloned msticpy source with main branch...

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.

Copilot AI and others added 4 commits January 21, 2026 17:37
…d override query method

Co-authored-by: ianhelle <13070017+ianhelle@users.noreply.github.com>
…ation

Co-authored-by: ianhelle <13070017+ianhelle@users.noreply.github.com>
Co-authored-by: ianhelle <13070017+ianhelle@users.noreply.github.com>
Co-authored-by: ianhelle <13070017+ianhelle@users.noreply.github.com>
Copilot AI changed the title [WIP] Fix MSSentinelSearch data provider connection issue with AzureSearchDriver Fix MSSentinelSearch environment and connection handling in AzureSearchDriver Jan 21, 2026
Copilot AI requested a review from ianhelle January 21, 2026 17:46
Copilot AI changed the title Fix MSSentinelSearch environment and connection handling in AzureSearchDriver Fix MSSentinelSearch environment name and connection check in AzureSearchDriver Jan 21, 2026
@ianhelle ianhelle marked this pull request as ready for review January 22, 2026 20:56
@ianhelle ianhelle enabled auto-merge January 30, 2026 18:30
- Add drop_duplicates(subset=['query']) before merge in get_whois_df to prevent
  row multiplication from duplicate whois results
- Change net_df fixture scope from module to function for test isolation with
  random sampling
- Add autouse fixture to clear LRU caches (get_whois_info, _whois_lookup) between
  tests to prevent state leakage
@ianhelle ianhelle requested a review from FlorianBracq February 2, 2026 16:38
@ianhelle
Copy link
Contributor

ianhelle commented Feb 2, 2026

Thx for the review @FlorianBracq -
These 3 PRs were all Co-pilot authored from issues - a bit of an experiment but not bad responses. However, I should have looked at them with a more critical eye than I did. I took your suggestions on this.

Copy link
Collaborator

@FlorianBracq FlorianBracq left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, a few more changes would probably help IMO.
Feel free to correct me if you feel otherwise!

Copy link
Contributor

@ianhelle ianhelle left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changes caused a bit of test upheaval but all good calls!

- Add connected property override in AzureSearchDriver, AzureMonitorDriver, and PrismaCloudDriver to validate both connection flag and client/header presence
- Remove redundant query() method override from AzureSearchDriver (now uses parent implementation with proper connection checking)
- Update DriverBase._ensure_connected() to check self.connected property instead of self._connected to enable child class customization
- Simplify _ensure_connected() calls across all drivers (remove provider name parameter)
- Use string literals for query filter data_environments instead of enum .name (fixes enum alias issues)
- Update tests to match stricter connection checking requirements

These changes ensure robust connection state validation across all drivers and prevent queries from running with incomplete connection state.

Related to PR #871
@ianhelle ianhelle requested a review from FlorianBracq February 7, 2026 18:55
Copy link
Collaborator

@FlorianBracq FlorianBracq left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That looks great!

- Change SubscriptionClient import from azure.mgmt.resource.subscriptions
  to azure.mgmt.subscription (removed in azure-mgmt-resource v25)
- Use SubscriptionClient directly instead of generic client() call to
  satisfy mypy type narrowing in _check_client and _legacy_auth
- Change Items.sku and Items.identity types to Any to match SDK types
- Add api_versions None guard in _get_api
- Use split('@', maxsplit=1) in query_source.py (ruff PLC0207)
- Remove stale azure.mgmt.resource.subscriptions mock from docs conf.py
- Update expected error messages from 'not connected to Splunk/OpenObserve'
  to 'not connected to SplunkDriver/OpenObserveDriver' to match the
  _ensure_connected implementation using class names
- Replace _create_not_connected_err('Splunk') calls with _ensure_connected()
  in splunk_driver.py for consistency with other drivers
- Update test assertions to match _ensure_connected class name format
  ('not connected to SplunkDriver/OpenObserveDriver')
@ianhelle ianhelle added this pull request to the merge queue Feb 9, 2026
Merged via the queue into main with commit cfeceef Feb 9, 2026
11 checks passed
@ianhelle ianhelle deleted the copilot/fix-azure-search-driver-connection branch February 9, 2026 22:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: New experimenal MSSentinelSearch data provider doesn't correctly use the AzureSearchDriver

3 participants