[release-12.0.5] Azure: Fix logs editor rendering (#109667)

* Azure: Fix logs editor rendering (#109491)

* Fix logs editor rendering

- Add test

* Type fixes

* Fix schema conditions and add test

* Fix lint

* Update public/app/plugins/datasource/azuremonitor/components/LogsQueryEditor/TimeManagement.tsx

Co-authored-by: Adam Yeats <16296989+adamyeats@users.noreply.github.com>

* Lint

---------

Co-authored-by: Adam Yeats <16296989+adamyeats@users.noreply.github.com>
(cherry picked from commit b34642b188)

# Conflicts:
#	public/app/plugins/datasource/azuremonitor/components/LogsQueryEditor/LogsQueryEditor.test.tsx

* Mock error and warn

public/app/plugins/datasource/azuremonitor/components/LogsQueryBuilder/LogsQueryBuilder.tsx

@@ -38,7 +38,7 @@ interface LogsQueryBuilderProps {
   query: AzureMonitorQuery;
   basicLogsEnabled: boolean;
   onQueryChange: (newQuery: AzureMonitorQuery) => void;
-  schema: EngineSchema;
+  schema?: EngineSchema;
   templateVariableOptions: SelectableValue<string>;
   datasource: Datasource;
   timeRange?: TimeRange;

public/app/plugins/datasource/azuremonitor/components/LogsQueryEditor/LogsQueryEditor.test.tsx

@@ -6,6 +6,7 @@ import { dateTime, LoadingState } from '@grafana/data';
 import createMockDatasource from '../../__mocks__/datasource';
 import createMockQuery from '../../__mocks__/query';
 import { ResultFormat } from '../../dataquery.gen';
+import { EngineSchema } from '../../types/types';
 import { createMockResourcePickerData } from '../MetricsQueryEditor/MetricsQueryEditor.test';
 
 import LogsQueryEditor from './LogsQueryEditor';
@@ -628,4 +629,130 @@ describe('LogsQueryEditor', () => {
       expect(onChange).toHaveBeenCalledWith(newQuery);
     });
   });
+
+  describe('schema loading and auto-completion', () => {
+    it('loads schema and table plans when resources change and builder mode is set', async () => {
+      // Mock these as we expect Kusto to complain about workers
+      jest.spyOn(console, 'warn').mockImplementation();
+      jest.spyOn(console, 'error').mockImplementation();
+      const mockSchema: EngineSchema = {
+        clusterType: 'Engine',
+        cluster: {
+          connectionString:
+            '/subscriptions/subscriptionId/resourceGroups/resourceGroup/providers/Microsoft.OperationalInsights/workspaces/la-workspace',
+          databases: [
+            {
+              name: '/subscriptions/subscriptionId/resourceGroups/resourceGroup/providers/Microsoft.OperationalInsights/workspaces/la-workspace',
+              tables: [
+                {
+                  columns: [
+                    {
+                      description: '',
+                      isPreferredFacet: false,
+                      name: 'TenantId',
+                      type: 'string',
+                    },
+                    {
+                      description: 'Date and time when dependency call was recorded.',
+                      isPreferredFacet: false,
+                      name: 'TimeGenerated',
+                      type: 'datetime',
+                    },
+                  ],
+                  description: 'Application Insights dependencies.',
+                  id: 'AppDependencies',
+                  name: 'AppDependencies',
+                  timespanColumn: 'TimeGenerated',
+                  hasData: true,
+                  related: {
+                    solutions: [],
+                    functions: [],
+                    categories: [],
+                  },
+                },
+              ],
+              functions: [],
+              majorVersion: 0,
+              minorVersion: 0,
+              entityGroups: [],
+            },
+          ],
+        },
+        database: {
+          name: '/subscriptions/subscriptionId/resourceGroups/resourceGroup/providers/Microsoft.OperationalInsights/workspaces/la-workspace',
+          tables: [
+            {
+              columns: [
+                {
+                  description: '',
+                  isPreferredFacet: false,
+                  name: 'TenantId',
+                  type: 'string',
+                },
+                {
+                  description: 'Date and time when dependency call was recorded.',
+                  isPreferredFacet: false,
+                  name: 'TimeGenerated',
+                  type: 'datetime',
+                },
+              ],
+              description: 'Application Insights dependencies.',
+              id: 'AppDependencies',
+              name: 'AppDependencies',
+              timespanColumn: 'TimeGenerated',
+              hasData: true,
+              related: {
+                solutions: [],
+                functions: [],
+                categories: [],
+              },
+            },
+          ],
+          functions: [],
+          majorVersion: 0,
+          minorVersion: 0,
+          entityGroups: [],
+        },
+      };
+      const mockDatasource = createMockDatasource();
+      mockDatasource.azureLogAnalyticsDatasource.getKustoSchema = jest.fn().mockResolvedValue(mockSchema);
+      // @ts-ignore: forcibly attach for test
+      mockDatasource.azureMonitorDatasource.getWorkspaceTablePlan = jest.fn().mockResolvedValue('plan');
+      const query = createMockQuery({
+        azureLogAnalytics: {
+          resources: [
+            '/subscriptions/def-456/resourceGroups/dev-3/providers/microsoft.operationalinsights/workspaces/la-workspace',
+          ],
+          mode: require('../../dataquery.gen').LogsEditorMode.Builder,
+        },
+      });
+      const onChange = jest.fn();
+      const onQueryChange = jest.fn();
+
+      await act(async () => {
+        render(
+          <LogsQueryEditor
+            query={query}
+            datasource={mockDatasource}
+            variableOptionGroup={variableOptionGroup}
+            onChange={onChange}
+            onQueryChange={onQueryChange}
+            setError={() => {}}
+            basicLogsEnabled={true}
+          />
+        );
+      });
+
+      await waitFor(() => {
+        expect(mockDatasource.azureLogAnalyticsDatasource.getKustoSchema).toHaveBeenCalledWith(
+          query.azureLogAnalytics?.resources?.[0]
+        );
+      });
+      expect(mockDatasource.azureMonitorDatasource.getWorkspaceTablePlan).toHaveBeenCalledTimes(1);
+      expect(mockDatasource.azureMonitorDatasource.getWorkspaceTablePlan).toHaveBeenCalledWith(
+        query.azureLogAnalytics?.resources,
+        'AppDependencies'
+      );
+    });
+  });
 });

public/app/plugins/datasource/azuremonitor/components/LogsQueryEditor/LogsQueryEditor.tsx

@@ -112,11 +112,9 @@ const LogsQueryEditor = ({
             if (schema.database?.tables) {
               schema.database.tables = t;
             }
-            setSchema(schema);
           });
-        } else {
-          setSchema(schema);
         }
+        setSchema(schema);
         setIsLoadingSchema(false);
       });
     }
@@ -285,7 +283,7 @@ const LogsQueryEditor = ({
         !!config.featureToggles.azureMonitorLogsBuilderEditor ? (
           <LogsQueryBuilder
             query={query}
-            schema={schema!}
+            schema={schema}
             basicLogsEnabled={basicLogsEnabled}
             onQueryChange={onQueryChange}
             templateVariableOptions={templateVariableOptions}

public/app/plugins/datasource/azuremonitor/components/LogsQueryEditor/TimeManagement.test.tsx

@@ -60,6 +60,25 @@ describe('LogsQueryEditor.TimeManagement', () => {
     );
   });
 
+  it('should render correctly even if no tables are in the schema', async () => {
+    const mockDatasource = createMockDatasource();
+    const query = createMockQuery({ azureLogAnalytics: { timeColumn: undefined } });
+    const onChange = jest.fn();
+
+    render(
+      <TimeManagement
+        query={query}
+        datasource={mockDatasource}
+        variableOptionGroup={variableOptionGroup}
+        onQueryChange={onChange}
+        setError={() => {}}
+        schema={FakeSchemaData.getLogAnalyticsFakeEngineSchema([])}
+      />
+    );
+
+    expect(screen.getByText('Time-range')).toBeInTheDocument();
+  });
+
   it('should render the default value if no time columns exist', async () => {
     const mockDatasource = createMockDatasource();
     const query = createMockQuery();

public/app/plugins/datasource/azuremonitor/components/LogsQueryEditor/TimeManagement.tsx

@@ -20,7 +20,7 @@ export function TimeManagement({ query, onQueryChange: onChange, schema }: Azure
       const timeColumnsSet: Set<string> = new Set();
       const defaultColumnsMap: Map<string, SelectableValue> = new Map();
       const db = schema.database;
-      if (db) {
+      if (db && db?.tables?.length > 0) {
         for (const table of db.tables) {
           const cols = table.columns.reduce<SelectableValue[]>((prev, curr, i) => {
             if (curr.type === 'datetime') {
@@ -39,11 +39,9 @@ export function TimeManagement({ query, onQueryChange: onChange, schema }: Azure
             });
           }
         }
-      }
         setTimeColumns(timeColumnOptions);
         const defaultColumns = Array.from(defaultColumnsMap.values());
         setDefaultTimeColumns(defaultColumns);
-
         // Set default value
         if (
           !query.azureLogAnalytics.timeColumn ||
@@ -64,6 +62,7 @@ export function TimeManagement({ query, onQueryChange: onChange, schema }: Azure
           }
         }
       }
+    }
   }, [schema, query.azureLogAnalytics?.dashboardTime, query.azureLogAnalytics?.timeColumn, setDefaultColumn]);
 
   const handleTimeColumnChange = useCallback(