Restore Doctor search after AdvisoryAI cold-start race

This commit is contained in:
master
2026-03-11 21:19:42 +02:00
parent 08006100a5
commit 66e67f1a97
5 changed files with 166 additions and 17 deletions

View File

@@ -10,6 +10,7 @@ namespace StellaOps.AdvisoryAI.KnowledgeSearch;
internal sealed class PostgresKnowledgeSearchStore : IKnowledgeSearchStore, IKnowledgeSearchCorpusAvailabilityStore, IAsyncDisposable
{
private static readonly JsonDocument EmptyJsonDocument = JsonDocument.Parse("{}");
private const string SchemaLockKey = "advisoryai_knowledge_schema";
private readonly KnowledgeSearchOptions _options;
private readonly ILogger<PostgresKnowledgeSearchStore> _logger;
@@ -38,6 +39,8 @@ internal sealed class PostgresKnowledgeSearchStore : IKnowledgeSearchStore, IKno
await using var connection = await GetDataSource().OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
await using var transaction = await connection.BeginTransactionAsync(cancellationToken).ConfigureAwait(false);
await AcquireSchemaLockAsync(connection, transaction, cancellationToken).ConfigureAwait(false);
const string createSchemaSql = "CREATE SCHEMA IF NOT EXISTS advisoryai;";
await ExecuteNonQueryAsync(connection, transaction, createSchemaSql, cancellationToken).ConfigureAwait(false);
@@ -1123,6 +1126,19 @@ internal sealed class PostgresKnowledgeSearchStore : IKnowledgeSearchStore, IKno
await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
}
private static async Task AcquireSchemaLockAsync(
NpgsqlConnection connection,
NpgsqlTransaction transaction,
CancellationToken cancellationToken)
{
await using var command = connection.CreateCommand();
command.Transaction = transaction;
command.CommandText = "SELECT pg_advisory_xact_lock(hashtext($1));";
command.CommandTimeout = ToCommandTimeoutSeconds(TimeSpan.FromSeconds(30));
command.Parameters.AddWithValue(SchemaLockKey);
await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
}
private static int ToCommandTimeoutSeconds(TimeSpan timeout)
{
if (timeout <= TimeSpan.Zero)

View File

@@ -372,6 +372,60 @@ public sealed class UnifiedSearchLiveAdapterIntegrationTests
(await CountDomainChunksAsync(connection, "policy")).Should().Be(4);
}
[Fact]
public async Task PostgresKnowledgeSearchStore_EnsureSchemaAsync_IsSafeUnderConcurrentStartupCalls()
{
await using var fixture = await StartPostgresOrSkipAsync();
var options = Options.Create(new KnowledgeSearchOptions
{
Enabled = true,
ConnectionString = fixture.ConnectionString,
FtsLanguageConfig = "simple"
});
await using var resetConnection = new NpgsqlConnection(fixture.ConnectionString);
await resetConnection.OpenAsync();
await ExecuteSqlAsync(resetConnection, "DROP SCHEMA IF EXISTS advisoryai CASCADE;");
var stores = Enumerable.Range(0, 6)
.Select(_ => new PostgresKnowledgeSearchStore(options, NullLogger<PostgresKnowledgeSearchStore>.Instance))
.ToArray();
var gate = new ManualResetEventSlim(false);
try
{
var tasks = stores
.Select(store => Task.Run(async () =>
{
gate.Wait();
await store.EnsureSchemaAsync(CancellationToken.None);
}))
.ToArray();
gate.Set();
await Task.WhenAll(tasks);
await using var verifyConnection = new NpgsqlConnection(fixture.ConnectionString);
await verifyConnection.OpenAsync();
(await ScalarAsync<string?>(
verifyConnection,
"SELECT to_regclass('advisoryai.kb_chunk')::text;")).Should().Be("advisoryai.kb_chunk");
(await ScalarAsync<int>(
verifyConnection,
"SELECT COUNT(*) FROM advisoryai.__migration_history;")).Should().BeGreaterThan(0);
}
finally
{
foreach (var store in stores)
{
await store.DisposeAsync();
}
gate.Dispose();
}
}
[Fact]
public async Task UnifiedSearchIndexer_RebuildAllAsync_PopulatesEnglishFtsColumns_AndRecallsUnifiedDomains()
{
@@ -981,6 +1035,19 @@ public sealed class UnifiedSearchLiveAdapterIntegrationTests
return Convert.ToInt32(scalar, System.Globalization.CultureInfo.InvariantCulture);
}
private static async Task<T> ScalarAsync<T>(NpgsqlConnection connection, string sql)
{
await using var command = connection.CreateCommand();
command.CommandText = sql;
var scalar = await command.ExecuteScalarAsync();
if (scalar is null || scalar is DBNull)
{
return default!;
}
return (T)Convert.ChangeType(scalar, typeof(T), CultureInfo.InvariantCulture);
}
private static async Task<int> CountEnglishTsvRowsAsync(NpgsqlConnection connection, string domain)
{
await using var command = connection.CreateCommand();

View File

@@ -141,19 +141,27 @@ async function resolveLink(page, options, timeoutMs = ELEMENT_WAIT_MS) {
const deadline = Date.now() + timeoutMs;
while (Date.now() < deadline) {
if (options.hrefIncludes) {
const candidates = page.locator(`a[href*="${options.hrefIncludes}"]`);
const count = await candidates.count();
for (let index = 0; index < count; index += 1) {
const candidate = candidates.nth(index);
const text = ((await candidate.innerText().catch(() => '')) || '').trim();
if (!options.name || text === options.name || text.includes(options.name)) {
return candidate;
}
const anchors = page.locator('a');
const anchorCount = await anchors.count();
for (let index = 0; index < anchorCount; index += 1) {
const candidate = anchors.nth(index);
const href = ((await candidate.getAttribute('href').catch(() => '')) || '').trim();
const text = ((await candidate.innerText().catch(() => '')) || '').trim();
if (options.hrefIncludes && !href.includes(options.hrefIncludes)) {
continue;
}
if (options.name && !(text === options.name || text.includes(options.name))) {
continue;
}
if (href || text) {
return candidate;
}
}
if (options.name) {
if (!options.hrefIncludes && options.name) {
const roleLocator = page.getByRole('link', { name: options.name }).first();
if ((await roleLocator.count()) > 0) {
return roleLocator;
@@ -278,7 +286,7 @@ async function main() {
{
route: '/mission-control/board',
actions: [
{ action: 'link:View all', name: 'View all', expectedPath: '/releases/runs' },
{ action: 'link:View all', name: 'View all', hrefIncludes: '/releases/runs', expectedPath: '/releases/runs' },
{ action: 'link:Review', name: 'Review', expectedPath: '/releases/approvals' },
{ action: 'link:Risk detail', name: 'Risk detail', expectedPath: '/security' },
{ action: 'link:Ops detail', name: 'Ops detail', expectedPath: '/ops/operations/data-integrity' },
@@ -286,24 +294,21 @@ async function main() {
{
action: 'link:Stage detail',
name: 'Detail',
hrefIncludes:
'/setup/topology/environments/stage/posture?tenant=demo-prod&regions=us-east&environments=stage&timeWindow=7d&region=us-east&environment=stage',
hrefIncludes: '/setup/topology/environments/stage/posture',
expectedPath: '/setup/topology/environments/stage/posture',
expectQuery: { environment: 'stage', region: 'us-east' },
},
{
action: 'link:Stage findings',
name: 'Findings',
hrefIncludes:
'/security/findings?tenant=demo-prod&regions=us-east&environments=stage&timeWindow=7d&region=us-east&environment=stage',
hrefIncludes: '/security/findings?tenant=demo-prod&regions=us-east&environments=stage',
expectedPath: '/security/findings',
expectQuery: { environment: 'stage', region: 'us-east' },
},
{
action: 'link:Risk table open stage',
name: 'Open',
hrefIncludes:
'/setup/topology/environments/stage/posture?tenant=demo-prod&regions=us-east&environments=stage&timeWindow=7d&region=us-east&environment=stage',
hrefIncludes: '/setup/topology/environments/stage/posture',
expectedPath: '/setup/topology/environments/stage/posture',
expectQuery: { environment: 'stage', region: 'us-east' },
},