Azure Data Studio Extension Development: help with query execution

[dfalso]

Hi all, not sure if this is the best place to ask this, but I don't see much discussion elsewhere.
I'm starting to develop extensions for internal use at my workplace, and can't find any documentation on proper ways of executing SQL. I've only been able to find one Q&A, which references the QueryProvider.runQueryAndReturn method. This doesn't work for me for two reasons:

1. runQueryAndReturn doesn't capture messages (error or otherwise).
2. runQueryAndReturn only works with SQL that returns result sets. So, a stored procedure that doesn't return a result set throws an error.

I've started to play around with runQueryString, and subscribing to notifications via registerOnMessage and registerOnQueryComplete. This isn't nearly as easy as firing sql through runQueryAndReturn and AWAITing the response, but I have made some guesses. However, I'm confused as to where to only subscribe to queries my process is executing, how to unsubscribe, etc. I'm currently seeing some bizarre results when opening up a new query window and simply running SELECT 1 after my process runs.
Documentation and examples would be great if anyone has them.

[Charles-Gagnon]

Are you able to provide the code that you have currently or describe what you'd like to do in more detail (and specifically what issues you're having)? There's not a lot of direct documentation out there for this stuff unfortunately but I can try to unblock you and get something started for that.

[dfalso]

Thanks for responding so quickly!
I registered for messages like this:

this.queryProvider.registerOnMessage(
    (message: azdata.QueryExecuteMessageParams) => this.onMessageCallback(message)
    );

and for completion like this:

this.queryProvider.registerOnQueryComplete(
    async (result: azdata.QueryExecuteCompleteNotificationResult) => 
    {
<stuff>
    }
);

*this is async because inside I need to AWAIT this.queryProvider.getQueryRows.

I wasn't sure if I should be calling disposeQuery() or not, but here's the line I have:

this.queryProvider.disposeQuery(this.queryResult.ownerUri);

• queryResult is of type QueryExecuteCompleteNotificationResult
• I decided this must be correct since before I had this, opening a new query window would show the the results of my extension process and then hang.

Ultimately, I didn't know what exactly was going on in the registration methods (do they support async callbacks? mixing async/await with callbacks is tricky to me), and since I'm newer to Javascript/Typescript, I was trying to achieve something more like runQueryAndReturn() - no callbacks, just an await with a return bundle that had the info I needed. But I also noticed that I didn't know how to unsubscribe from other queries on the provider, and was getting noise from Object Explorer itself.

I'm happy to create an Issue related to documentation if you think that's where this belongs. But it'd be nice to have more out there for everyone.

[Charles-Gagnon]

1. The "callbacks" in this case are just event notifications, the caller never awaits on them or anything. So they support both async and non-async functions just fine since in either case the calling code doesn't care about the response
2. The callbacks are indeed for the entire provider itself, you would need to use the ownerUri field to filter out the ones that didn't apply to you.

How are you currently running the queries?

There seems to be some issues currently it looks like - for example I can get the notifications to fire when I use runQueryString but not when I use runQueryAndReturn.

Do you want this to be a background thing or do you want to actually open up an editor and then run the query in there? And then are the messages something you actually plan on using yourself?

[dfalso]

Thanks for the clarifications on the callbacks.

In regards to the URI for filtering, my process adds a button to the dashboard and also an entry in the Object Explorer database nodes menu. I've been using the IConnectionProfile when providing a callback toazdata.tasks.registerTask(). That's why I end up getting a lot of noise - every time the Object Explorer refreshes, I get query notifications. Another edge case issue I found is that I can sometimes hit a race condition of sorts - I register my registerOnQueryComplete, then try to runQueryString, but if the connection is in use I get a notification about the currently running query, not the one I'm intending to run.
Do you recommend spawning a new connection then? If so, I'll have to look into how to clone/recreate a connection. (I'm only dealing with MSSQL)

In my testing, the runQueryAndReturn method doesn't fire any other events. That's why I started using runQueryString.

From my perspective, any code run against the SQL Server could potentially fail, so I want to be able to trap that failure and obtain the error messages, to either display to the user or at least to log. The registerOnMessage method callback can use the returned QueryExecuteMessageParams instance, which has the message property (of type IResultMessage), which has an isError property.

In answer to your last question, I think the current model makes sense if you want to get notifications about all the queries being run - I've seen a query history extension that does that. But I want my process to run a SQL statement, get the results, and run another one, and that's more awkward in the current model. In other words, if someone wants to do something like:

await runMyQuery("Create table ...");
if (failed) { do something with error }
await runMyQuery("Create index ...");
if (failed) { do something with error }
await runMyQuery("Create index ...");
...

it becomes a mess of callbacks and making sure the callback instance corresponds to where you think your code left off.

Can you speak to disposeQuery? Is it for cleaning up the BatchSummary\ResultSetSummary from the QueryExecuteCompleteNotificationResult? I wasn't sure.

Again, thanks for chatting and trying to help. As much as I have questions, I'm also trying to document as much of what I've learned as possible on this thread, since like you said, there's not much out there.

[Charles-Gagnon]

You should definitely make a new connection, re-using an existing connection that's tied to a document or something can easily fail if that connection is closed (for example if the editor is closed). You can use azdata.connection.connect to create the new connection from a profile (either entering your own or using something like getCurrentConnection)

For the dispose query - that is sent to the SQL Tools Service service that handles executing all the queries, and as the name implies it just calls Dispose on the query object which goes through and disposes everything associated with that query (result sets and such). From what I can tell looking at it this should not close the connection associated with the query, just clean up the resources from that query. So when you're finished getting the results or whatever of the query then calling dispose would be a good idea to ensure you don't have a bunch of stale data sitting around.

If you call runQueryAndReturn does it not throw an error if an error occurs executing the query? That's what I'd typically expect the behavior to be which should fit your case to do something like this

try {
    const results = await runMyQuery("Create table ...");
} catch (err) {
   // Handle error here
}

I'll try and find some time to sit down and try this out myself to see if I can get a working sample to show you too, feel free to ask any more direct questions though if you have them!

[dfalso]

This is great.
Good information about the connection - the few examples out there reuse existing ones, so this is good to know.

For the dispose: do you know if it unsubscribes the callbacks from the register* methods? Or is there anything that does?
Given a situation in which I want to run 3 statements (and hand-waving over handling errors for the moment), my assumption was that the following code was a bad idea:

...
prov.registerOnQueryComplete((prov, uri) => callback1(prov, uri));
prov.runQueryString(uri, "create table ...");

function callback1(prov, uri)
{
    prov.registerOnQueryComplete((prov, uri) => callback2(prov, uri));
    prov.runQueryString(uri, "create index index1..."); //now both callback1 and callback2 are going to fire, correct?
}
function callback2(prov, uri)
{
    prov.registerOnQueryComplete((prov, uri) => callback3(prov, uri));
    prov.runQueryString(uri, "create index index2..."); //now callback1, callback2, and callback3 are going to fire, correct?
}

function callback3(prov, uri)
{
    <do stuff>
}

At the moment, I'm glomming everything into one statement and hoping for the best:

...
this.globalstate = "step1";
prov.registerOnQueryComplete((prov, uri, state) => queryCompleteCallback(prov, uri, state));
prov.runQueryString(uri, "create table ...;
create index index1...;
create index index2...;");

function queryCompleteCallback(prov, uri, state)
{
    <do stuff>
}

Re: runQueryAndReturn: I would have expected it to behave that way too, but it doesn't seem to:

try{
    const testresult: azdata.SimpleExecuteResult = await prov.runQueryAndReturn(uri, "select 1/0");
}
catch (error)
{
    if (error)
    {
        throw error;
    }
}

Expected: Divide by zero error encountered.
Result: No error, but the SimpleExecuteResult has 0 rows.

try{
    const testresult2: azdata.SimpleExecuteResult = await prov.runQueryAndReturn(uri, "select 1 from dbo.MyTableWhichDoesNotExist");
}
catch (error)
{
    if (error)
    {
        throw error;
    }
}

Expected: Invalid object name 'dbo.MyTableWhichDoesNotExist'.
Result: error.message equals: "Query has no results to return"

[Charles-Gagnon]

DisposeQuery does not unsubscribe any callbacks that were registered (those callbacks are a global thing for the provider, not tied to a specific query).

There currently isn't a way to unsubscribe to those callbacks either - once you've set one up it will continue to be called until ADS is restarted. Probably behavior we should change but for now you'll just have to ignore the "extra" events that you'll be getting if you use any of them.

What you have with registering the callback multiple times isn't what you want - since the first one will also be triggered from any subsequent queries (or any queries being completed at all). Instead you'd want to set up a single callback and then somehow in there determine which of your callback methods you want to call based on the info passed to you (uri and state)

That's unfortunate about the error not throwing - I'll see if I can dig into this stuff and find a way to get the behavior you want (or fix it so it behaves correctly). But for now you'll just have to find another way I guess.

A real hacky solution would be to have the TRY/CATCH in the query itself, and then return a "special" value if it fails. Then you could just check the return value for that. Probably easier than trying to hook into the messages callback (especially given that runQueryAndReturn doesn't trigger those)

[dfalso]

Re: registering callback1,2,3: that's what I thought, which is why I don't do it. I think the real best way forward, for me, right now, is to:

1. register callbacks for registerOnMessage and registerOnQueryComplete. QueryExecuteCompleteNotificationResult has the ownerURI property, and I can pass the connection ID's URI and the QueryProvider as a parameter in the closure for the callback.
2. create a new connection for each SQL statement (or batch), or store global state somewhere (since any state variables would be passed in at registration time, I can't do that at runQueryString time)
3. big switch statement in the callbacks
4. disposeQuery when I'm done

Re: hacky solution: I'm doing something in essence like that to pass output parameters back to the controller. I'm assuming there's currently no parameterization capabilities that I overlooked.

Re: errors not throwing. Please do check my examples, but that's what I get. I'd imagine that there'd be more appetite in creating a new method than changing the operations of an existing one. Do you think it's worth me opening an Issue for discussion, or is the team overloaded right now? If so, I'd be interested in getting that discussion going.

[Charles-Gagnon]

If you could open an issue that'd be great! I can't promise we'll be able to get to it immediately but getting it out there would let us have all the context in one place and help other people who might try to use that same API.

And no, I don't believe we have any support for parameters at this point unfortunately.

[dfalso]

OK, thanks. I'll start up an issue with everything I can think of summarized.

[Charles-Gagnon]

#19509 hopefully that's a lot clearer!

[dfalso]

Yes it is. Thanks!

[dfalso]

It looks like azdata.connection.ConnectionProfile does not implement azdata.IConnectionProfile. Therefore, I can't use the results of getCurrentConnection() (or anything from getConnections()) as a parameter into azdata.connection.connect(). Is there a way to get from the one to the other using built-in methods that I'm missing? The property names look close enough that I could probably build my own conversion mapper, but was hoping to use a more official way.
Aside: I would have assumed that ConnectionProfile would implement IConnectionProfile, given the name.

[Charles-Gagnon]

Yeah I was wondering that too, I have no idea why that's the case (it was done before my time here). There aren't any conversion methods currently (but we should probably add one...if not just having ConnectionProfile outright implement IConnectionProfile...) but for the time being this should work too

const connProfile = await azdata.connection.getCurrentConnection();
const iConnProfile = { ...connProfile, providerName: connProfile.providerId, id: '' };
const result = await azdata.connection.connect(iConnProfile, false, false);

Since the providerName and id are the only two values missing. Provider name is actually just the ID, and for "connect" we actually overwrite the ID regardless with a new one so it doesn't matter what you put there.

[dfalso]

I'm happy to add that as an Issue also. I'll write that up today sometime.

[dfalso]

#19523
I've got a few more questions that I'm going to post on different Discussion threads, so as to not clutter this one up.
But I wanted to write some actual code to summarize what I have understood from above, as well as document for anyone getting started.

    _messages: azdata.QueryExecuteMessageParams[] = [];
    accumulateMessages(message: azdata.QueryExecuteMessageParams, uri: string)
    {
        if (message.ownerUri === uri)
        {
            this._messages.push(message);
        }
    }
    //async because of potential "await prov.getQueryRows()" call
    async onResult(result: azdata.QueryExecuteCompleteNotificationResult, prov: azdata.QueryProvider, uri: string)
    {
        if (result.ownerUri === uri)
        {
            if (result.batchSummaries.some(x => x.hasError))
            {
                //example showing usage of messages:
                const errormessagesummary: string = this._messages.filter(x => x.message.isError).map(x => x.message.message).join('\n');
            }
            else
            {
                //do stuff

                //cleanup when done
                prov.disposeQuery(uri);
            }
        }
    }
    async createNewConnectionFromExisting(connectionID: string): Promise<string>
    {
        const profiles: azdata.connection.ConnectionProfile[] = await azdata.connection.getConnections(true);
        let foundprofile: azdata.connection.ConnectionProfile | undefined = profiles.find(x => x.connectionId === connectionID);
        if (foundprofile === undefined ) { throw new Error(); }
        else
        {
            const iConnProfile: azdata.IConnectionProfile = { ...foundprofile, providerName: foundprofile.providerId, id: '' };
            const result: azdata.ConnectionResult = await azdata.connection.connect(iConnProfile, false, false);
            return result.connectionId;
        }

    }
    async startRunMyQuery(prov: azdata.QueryProvider, connectionID: string, query: string)
    {
        let connectionIDUsed: string;
        const queryNeedsNewConnection: boolean = true;
        if (queryNeedsNewConnection)
        {
            connectionIDUsed = await this.createNewConnectionFromExisting(connectionID);
        }
        else
        {
            connectionIDUsed = connectionID;
        }
        const uri: string = await azdata.connection.getUriForConnection(connectionID);

        prov.registerOnMessage((message: azdata.QueryExecuteMessageParams) => this.accumulateMessages(message, uri));
        prov.registerOnQueryComplete(async (result: azdata.QueryExecuteCompleteNotificationResult) => this.onResult(result, prov, uri));
        prov.runQueryString(uri, query);
    }

Please let me know if I messed anything up.

[Charles-Gagnon]

You probably shouldn't be registering the callbacks every time you call startRunMyQuery - instead just registering them once when your extension is activated and then keeping track of the "active" URIs you care about at the same level as _messages and having the callbacks do a lookup there to determine if you care about that message.

Other than that nothing else stands out aside from some general cleanup (for example queryNeedsNewConnection is always true so that if statement isn't useful atm)

[dfalso]

Good point as to the callbacks. I hardcoded queryNeedsNewConnection for example purposes, but I'll make that a parameter to be more explicit.
I'll make some corrections here:

    _messages: azdata.QueryExecuteMessageParams[] = [];

    //state variables for current process
    _currentProcessURI: string = "";
    _currentProcessProvider: azdata.QueryProvider;

    accumulateMessages(message: azdata.QueryExecuteMessageParams)
    {
        this._messages.push(message);
    }
    //async because of potential "await prov.getQueryRows()" call
    async onResult(result: azdata.QueryExecuteCompleteNotificationResult)
    {
        if (result.ownerUri === _currentProcessURI)
        {
            if (result.batchSummaries.some(x => x.hasError))
            {
                //example showing usage of messages:
                const errormessagesummary: string = this._messages.filter(x => x.ownerUri === _currentProcessURI && x.message.isError).map(x => x.message.message).join('\n');
            }
            else
            {
                //do stuff

                //cleanup when done
                _currentProcessProvider.disposeQuery(uri);
            }
        }
    }
    async createNewConnectionFromExisting(connectionID: string): Promise<string>
    {
        const profiles: azdata.connection.ConnectionProfile[] = await azdata.connection.getConnections(true);
        let foundprofile: azdata.connection.ConnectionProfile | undefined = profiles.find(x => x.connectionId === connectionID);
        if (foundprofile === undefined ) { throw new Error(); }
        else
        {
            const iConnProfile: azdata.IConnectionProfile = { ...foundprofile, providerName: foundprofile.providerId, id: '' };
            const result: azdata.ConnectionResult = await azdata.connection.connect(iConnProfile, false, false);
            return result.connectionId;
        }

    }

    doThisBeforeRunningQueries(prov: azdata.QueryProvider)
    {
        prov.registerOnMessage((message: azdata.QueryExecuteMessageParams) => this.accumulateMessages(message));
        prov.registerOnQueryComplete(async (result: azdata.QueryExecuteCompleteNotificationResult) => this.onResult(result));
    }

    async startRunMyQuery(prov: azdata.QueryProvider, connectionID: string, query: string, queryNeedsNewConnection: boolean)
    {
        let connectionIDUsed: string;
        if (queryNeedsNewConnection)
        {
            connectionIDUsed = await this.createNewConnectionFromExisting(connectionID);
        }
        else
        {
            connectionIDUsed = connectionID;
        }
        _currentProcessURI = await azdata.connection.getUriForConnection(connectionID);
        _currentProcessProvider = prov;

        prov.runQueryString(_currentProcessURI, query);
    }

(I'm not sure if I can back into the connection from on the URI, so I opted to preserve the current QueryProvider as well)
Otherwise, I think this is it.

[dfalso]

I don't think there's much more to add, so I'm marking this as answered.
Thanks for all your help!

[dfalso]

Hi, opening this up again because I'm experiencing some problems.
I've created a barebones WebView project which demonstrates the issue. These are code excerpts from that, so please ignore the lack of error handling, etc. The WebView itself is a simple page with a button, which relays a message to the controller to run a query. The receiveMessage function is below:

    async receiveMessageAsync(connectionid: string, panel: vscode.WebviewPanel)
    {
        const profile: azdata.connection.ConnectionProfile = await controller.getConnectionProfileAsync_ConnectionID(connectionid);

        const newConnectionID: string = await controller.cloneConnection(connectionid, profile);
        const newConnectionProfile: azdata.connection.ConnectionProfile = await controller.getConnectionProfileAsync_ConnectionID(newConnectionID);
        const newConnectionURI: string = await azdata.connection.getUriForConnection(newConnectionID);

        const query: string = 
            "select '" + connectionid + "' as [orig_ConnectionID], '" + newConnectionID + "' as [new_ConnectionID]"
            ;
        const queryProvider: azdata.QueryProvider = azdata.dataprotocol.getProvider<azdata.QueryProvider>(newConnectionProfile.providerId, azdata.DataProviderType.QueryProvider);
        queryProvider.registerOnMessage(
            (message: azdata.QueryExecuteMessageParams) => this.accumulateMessages(message, newConnectionURI)
            );
        queryProvider.registerOnQueryComplete(
            async (result: azdata.QueryExecuteCompleteNotificationResult) => this.accumulateResults(result, newConnectionURI, panel, queryProvider)
            );
        queryProvider.runQueryString(newConnectionURI, query);
    }

You can see that I'm creating a new connection from the original, then using that connection for the query registration and execution.

Here's the accumulateResults:

    async accumulateResults(result: azdata.QueryExecuteCompleteNotificationResult, connectionURI: string, panel: vscode.WebviewPanel, queryProvider: azdata.QueryProvider)
    {
        if (result.ownerUri === connectionURI)
        {
            const batchSummary: azdata.BatchSummary = result.batchSummaries[0];
            if (batchSummary.resultSetSummaries === null) { throw new Error(); }
            const resultSetSummary: azdata.ResultSetSummary = batchSummary.resultSetSummaries[0];

            const executeresult: azdata.QueryExecuteSubsetResult = await queryProvider.getQueryRows(
                {
                  ownerUri: connectionURI
                  , batchIndex: resultSetSummary.batchId
                  , resultSetIndex: resultSetSummary.id
                  , rowsStartIndex: 0
                  , rowsCount: resultSetSummary.rowCount
                }
                );

            const message: any = {
                origid: executeresult.resultSubset.rows[0][0].displayValue
                , newid: executeresult.resultSubset.rows[0][1].displayValue
                };

            queryProvider.disposeQuery(connectionURI);
            const profile: azdata.connection.ConnectionProfile = await controller.getConnectionProfileAsync_ConnectionURI(connectionURI);
            const connectionProvider: azdata.ConnectionProvider = await azdata.dataprotocol.getProvider<azdata.ConnectionProvider>(
                profile.providerId
                , azdata.DataProviderType.ConnectionProvider
                );
            connectionProvider.disconnect(connectionURI);

            await panel.webview.postMessage(message);
        }
    }

This webview works fine, but the issue is the side effects. Here's the script:

1. Debug this extension (or alternately, package it and run ADS). In this case, I've attached a hook into the Object Explorer node menu, and that's where it gets the original connection. Call the extension.
2. Click the button and wait for the connection IDs to be received by the webview form.
3. Open a new query window. Run this: "select 1". You'll find that the query window's Results shows the results from the webview SQL call, and not the "select 1" call. Also, you'll find that the query window is stuck in Running state, but the Cancel button doesn't work.
4. In testing this, you can keep the webview open, or close it before opening the new query window - doesn't seem to matter. You can open the new query window via Ctrl+N, or right-click on a server node. You can even script out an object and replace the SQL. Doesn't seem to matter.

I've attached a picture of the new query window state. I've also attached a ZIP file of the extension code, though I stripped out the .\nodes folder to make it compact enough. I'm also unmarking the previous answer, since it doesn't seem to be complete.

issuedemonstrator.zip

[Charles-Gagnon]

Unfortunately you've hit a long-standing issue we have where the register... functions only allow a single handler to be registered

microsoft/sqlops-dataprotocolclient#62

So by registering your own you're overwriting the one that the MSSQL extension registered which completely breaks everything.

There isn't any way around this, as things stand right now you're completely unable to use any of these functions.

One thing you can look to see if it'd prove useful is the query events - I recently updated them to include the range of text that was executed as well as the messages (it already included the query document URI) (#19777 and #19823). See https://github.com/microsoft/azuredatastudio/blob/main/extensions/query-history/src/queryHistoryProvider.ts#L27 for an example of how I'm using that.

Note that you'll need to bring in the azdata.proposed.d.ts typings file to get the typing support for this since this isn't promoted to the stable API yet. Just copy https://raw.githubusercontent.com/microsoft/azuredatastudio/main/src/sql/azdata.proposed.d.ts into your project src directory and it should be picked up right away.

[dfalso]

Ugh. It's a shame they marked it as "engineering-focused" since this is very "customer scenario-facing" for extension developers.

Is it possible to instantiate another instance of the query provider (in my case, the MSSQL one)? If so, I could register extension callbacks against the new instance. I'm already needing to create a new connection anyway - I could create the new connection in the new provider.

Otherwise, I'll take a look at the query events and report back. I'll have to see if one can get result sets from those events or not. If no, that's a big problem for me (and all extension developers).

[Charles-Gagnon]

No, only one provider for every connection type is currently supported. If you made your own that would be a completely separate implementation (each provider is supposed to be for different data sources, so PGSQL, MSSQL, MySQL, etc.) - so you'd need to implement all the things like making connections. Definitely not the path you want to be taking.

The events don't currently send the query results (that I'm seeing at least) but you could try running the query and getting the results from that, and then using the event to gather the messages separately.