Using explain to get the approximate count with distributed tables

[jnoortheen]

I am using explain on queries to get the approximate count of a query. When the tables are distributed with citus, I am getting this following error

Traceback (most recent call last):
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/backends/utils.py", line 85, in _execute
    return self.cursor.execute(sql, params)
psycopg2.errors.InternalError_: cannot switch local execution status from local execution required to local execution disabled since it can cause visibility problems in the current transaction


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/src/main/queryset_fn.py", line 191, in estimated_count
    return qs.count(), False
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/models/query.py", line 412, in count
    return self.query.get_count(using=self.db)
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/models/sql/query.py", line 528, in get_count
    number = obj.get_aggregation(using, ['__count'])['__count']
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/models/sql/query.py", line 513, in get_aggregation
    result = compiler.execute_sql(SINGLE)
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/models/sql/compiler.py", line 1175, in execute_sql
    cursor.execute(sql, params)
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/backends/utils.py", line 66, in execute
    return self._execute_with_wrappers(sql, params, many=False, executor=self._execute)
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/backends/utils.py", line 75, in _execute_with_wrappers
    return executor(sql, params, many, context)
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/backends/utils.py", line 85, in _execute
    return self.cursor.execute(sql, params)
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/utils.py", line 90, in __exit__
    raise dj_exc_value.with_traceback(traceback) from exc_value
  File "/Users/noor/.virtualenvs/app/lib/python3.8/site-packages/django/db/backends/utils.py", line 85, in _execute
    return self.cursor.execute(sql, params)
django.db.utils.InternalError: cannot switch local execution status from local execution required to local execution disabled since it can cause visibility problems in the current transaction 

the estimated count function look like this

def _estimate_count(queryset: models.QuerySet) -> Optional[int]:
    lines = queryset.only("pk").explain(format="json")
    if isinstance(lines, str):
        lines = ast.literal_eval(lines)  # turn to dict
    if isinstance(lines, list) and lines:
        return lines[0]["Plan"]["Plan Rows"]
    return None

The queryset can be something as simple as Model.objects.filter() . It seems when estimate is used with other queries, this error is raised. It is also happening during tests.