Update and delete support

[mjgarton]

Parts of datafusion and its ecosystem support DELETE and UPDATE queries (for example the sql parser), but other parts only really support SELECT and INSERT (for example, the TableProvider trait, the DefaultPhysicalPlanner).

Overall this means that datafusion by default does not support updates or deletes.

datafusion-postgres currently reflects this by supporting SELECT and INSERT queries. Supporting INSERT includes mapping from the returned table structure and correctly returning the number of affected rows to the postgresql client. The format it needs to map from is from DataSinkExec:

/// Create a output record batch with a count
///
/// ```text
/// +-------+,
/// | count |,
/// +-------+,
/// | 6     |,
/// +-------+,
/// ```
fn make_count_batch(count: u64) -> RecordBatch {
    let array = Arc::new(UInt64Array::from(vec![count])) as ArrayRef;

    RecordBatch::try_from_iter_with_nullable(vec![("count", array, false)]).unwrap()
}

It is actually possible today for applications using datafusion to support DELETE and UPDATE queries by using a custom QueryPlanner implementation. However, this will not by correctly handled by datafusion-postgres yet because it doesn't map the number of rows affected.

As an experiment I made my custom QueryPlanner implementation with support for update and delete queries, and it returns the same "count" structure as DataSinkExec returns, and I hacked datafusion-postgres to map the returned count correctly. This all works fine.

My question is: How should this be supported properly in datafusion-postgres? It seems wrong to add support for update & delete mapping based on assuming a specific table structure when nothing is defined in datafusion itself for updates & deletes.

Maybe datafusion-postgres could expose a way to customise or override how it maps results to pgwire responses?

I'd be happy to work on this, but I'm not sure what the best approach is.

(Longer term - I also wonder whether one data datafusions TableProvider trait might add support for updates and deletes, and whether working on that is the right approach?)

[tonyalaribe]

@mjgarton I'm curious about how you got datafusion-postgres to pass on the update/delete queries to your logic for actual handling.

I'm currently exploring approaches for this, and so far, extended datafusion-postgres to support user-defined hooks passed into the DfSessionService struct as a trait. That way, users can handle special queries however they see fit. In my case, that means catching update and delete queries, parsing them into statements, and manually executing them without Data Fusion.

My hook looks like this:

#[async_trait]
pub trait QueryHook: Send + Sync {
    async fn handle_query(
        &self,
        query: &str,
        session_context: &SessionContext,
        client: &dyn ClientInfo,
    ) -> Option<PgWireResult<Vec<Response<'static>>>>;
}

Then in the PgWire trait implementations

#[async_trait]
impl SimpleQueryHandler for DfSessionService {
    async fn do_query<'a, C>(&self, client: &mut C, query: &str) -> PgWireResult<Vec<Response<'a>>>
    where
        C: ClientInfo + Unpin + Send + Sync,
    {
        let query_lower = query.to_lowercase().trim().to_string();
        ...

        // Check query hook first
        if let Some(hook) = &self.query_hook {
            if let Some(result) = hook.handle_query(query, &self.session_context, client).await {
                return result;
            }
        }

        let df_result = self.session_context.sql(query).await;

        ...
    }
    ...
}

But i'm also unsure if this is the right way to go about it, I might create a PR regardless.

[tonyalaribe]

Your comment helped me look at things again, and I realised there's no need for a hook in datafusion-postgres. I created a custom QueryPlanner implementation like you mentioned, and it works fine. I had to implement a custom QueryPlanner and then a custom ExecutionPlan.

[mjgarton]

@tonyalaribe Yes, that's exactly how I did it too. The only part that's needed in datafusion-postgres is to correctly return the number of affected rows to the client.

Something like this would work for me, but as explained above, I don't think it's the right approach:

0cfafc0

[sunng87]

@tonyalaribe @mjgarton I still think the QueryHook idea is a good one. We can implement all those implementation of set, show statements, or even RBAC as optional hooks. But i think we can accept a Statement instead of string for the query argument.

[sunng87]

@mjgarton for the update/delete operation, is it required in your datafusion based application? In this library we should focus on the integration of postgres protocol and datafusion, instead of building a full-featured database. I think it's still very early for our library to support it before full support landed in datafusion. If it's really necessary to add, it would be nice to be opt-in with the QueryHook framework. Let me know what do you think.

[mjgarton]

@sunng87 Yes my application does support updates and deletes. I fully agree that datafusion-postgres should not care about these things because even datafusion itself doesn't support it.

To be clear though: I was not proposing actually supporting the update & delete logic itself in datafusion-postgres - only the logic to correctly return the "affected rows" information in the tag in pgwire. But that too could be done externally and wired up using something like a query hook mechanism, so I agree hooks and external configuration are the best way to go.

[sunng87]

Sounds perfect @mjgarton !

@tonyalaribe Will you be working on this QueryHook framework? This will be the major feature for extensibility of datafusion-postgres

[tonyalaribe]

@sunng87 sure. I'll create a PR.
Other than accepting Statement instead of query string, do you have any other nice to haves or expectations?

[sunng87]

Thank you @tonyalaribe . I think it should be good. We can refactor our current query handler implementation to move most preprocessing logic into the new framework, and we will see what are other potential requirements.

[mjgarton]

I notice the QueryHook support is now implemented. Thanks @tonyalaribe and @sunng87 for getting this done so quickly.

However, I don't think it's sufficient for the use case that I originally opened this issue for.

At first, I thought I could use a QueryHook, although I would have to copy and paste a lot of the code from DfSessionService.

In the end though, I realised that I don't think its possible because for my use case I need to execute the query inside the QueryHook, but I don't have access to the query parameters when using extended query. (Also, I can't know if I am doing an simple or an extended query as far as I can tell)

It's quite possible that I've misunderstood how to use the QueryHook mechanism. Do you have any thoughts?

I will find some time in the next 1 or 2 days to see what would need modifying to support my use case.

[sunng87]

@mjgarton you are right. I think we focused too much on simple query scenario. I will pass the portal to query hook which contains the query parameters. Let me know if there is any other input you need for the delete/update hook implementation.

[sunng87]

My straight-forward idea is to change QueryHook to cover three most important APIs in pgwire:

#[async_trait]
pub trait QueryHook: Send + Sync {
    /// called in simple query handler to return response directly
    async fn handle_simple_query(
        &self,
        statement: &Statement,
        session_context: &SessionContext,
        client: &dyn ClientInfo,
    ) -> Option<PgWireResult<Response>>;

    /// called at extended query parse phase, for generating `LogicalPlan` from statement
    async fn handle_extended_parse_query(
        &self,
        statement: &Statement,
        session_context: &SessionContext,
        client: &dyn ClientInfo,
    ) -> Option<PgWireResult<LogicalPlan>>;

    /// called at extended query execute phase, for query execution
    async fn handle_extended_query(
        &self,
        statement: &Statement,
        logical_plan: &LogicalPlan,
        params: &ParamValues,
        session_context: &SessionContext,
        client: &dyn ClientInfo,
    ) -> Option<PgWireResult<Response>>;
}

[mjgarton]

I wanted to be sure it worked for my use case before replying, so I've implemented a subset of this in #208

I added handle_extended_query but not handle_extended_parse_query because I don't personally have a need for it yet.

I also did some other minor changes in handlers.rs to avoid parsing the query multiple times.

I can confirm that this works for my use case, but I did have to duplicate quite a lot of code into my query hook implementation. I am still thinking about possibly future changes to avoid that.

Also, FYI once we're happy with QueryHook, I would like to do some more general clean-up and fix some todos in handlers.rs

[mjgarton]

I think I will close this now, because I have everything I need to support my original requirement. (It would be great to see a release containing QueryHook support if possible though @sunng87 ?)