async def execute(self,
query,
params=None,
connection=None,
transaction=False):
if isinstance(query, (dict, list)):
# build zql query
query, params = build(query,
dialect=self.backend.dialect,
combine=True)
pool = await self.pool
connection = connection or self._connection
connection = aecho(connection) if connection else pool.acquire()
pquery = print_query(query, params)
async with connection as conn:
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}{pquery}{SEPN}",
True):
raise Exception(f"{self}: execute aborted")
else:
self.log(f"{self}: execute{SEP}{pquery}{SEPN}")
transaction = conn.transaction() if transaction else aecho()
async with transaction:
try:
return await self.backend.execute(conn, query, params)
except Exception as e:
err = f"{self}: execute failed; {e.__class__.__name__}: {e}"
err += f"\nQuery:{SEP}{pquery}{SEPN}"
raise Exception(err)
async def stream(self,
query,
params=None,
transaction=True,
connection=None):
if isinstance(query, (dict, list)):
# build zql
queries = build(query, dialect=self.backend.dialect)
else:
queries = [(query, params)]
pool = await self.pool
connection = connection or self._connection
connection = aecho(connection) if connection else pool.acquire()
async with connection as conn:
transaction = conn.transaction() if transaction else aecho()
async with transaction:
for query, params in queries:
pquery = print_query(query, params)
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}{pquery}{SEPN}",
True):
raise Exception(f"{self}: stream aborted")
else:
# self.log(f"{self}: stream{SEP}{pquery}{SEPN}")
async for row in self.backend.cursor(
conn, query, params):
yield row
async def query(self,
query,
params=None,
connection=None,
many=True,
columns=True,
transaction=False):
pool = await self.pool
connection = connection or self._connection
connection = aecho(connection) if connection else pool.acquire()
if isinstance(query, (dict, list)):
# build zql
queries = build(query, dialect=self.backend.dialect)
else:
queries = [(query, params)]
one = len(queries) == 1
all_results = []
async with connection as conn:
transaction = conn.transaction() if transaction else aecho()
async with transaction:
for query, params in queries:
pquery = print_query(query, params)
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}{pquery}{SEPN}",
True):
raise Exception(f"{self}: query aborted")
else:
self.log(f"{self}: query{SEP}{pquery}{SEPN}")
try:
results = await self.backend.fetch(conn, query, params)
except Exception as e:
err = f"{self}: query failed; {e.__class__.__name__}: {e}"
err += f"\nQuery:{SEP}{pquery}{SEPN}"
raise Exception(err)
if many:
all_results.append(
results if columns else [r[0] for r in results])
continue
else:
num = len(results)
if num == 0:
# no results -> return None
all_results.append(None)
continue
if num != 1:
raise Exception(
f"{self}: query failed; expecting <=1 row, got {num}\n"
f"Query:{SEP}{query}{SEPN}")
result = results[0]
all_results.append(result if columns else result[0])
return all_results[0] if one else all_results
async def copy_from(self, **kwargs):
pool = await self.pool
table_name = kwargs.pop("table_name", None)
schema_name = kwargs.get('schema_name', None)
transaction = kwargs.pop("transaction", False)
connection = kwargs.pop("connection", self._connection)
connection = aecho(connection) if connection else pool.acquire()
close = kwargs.pop("close", False)
query = kwargs.pop("query", None)
params = kwargs.pop('params', None)
if table_name:
target_label = f"{schema_name}.{table_name}" if schema_name else table_name
else:
if not query:
raise NotImplementedError("table or query is required")
if isinstance(query, (list, dict)):
# compile from zQL
query, params = build(query,
dialect=self.backend.dialect,
combine=True)
target_label = print_query(query, params)
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}copy from {target_label}{SEPN}",
True):
raise Exception(f"{self}: copy_from aborted")
else:
self.log(f"{self}: copy_from{SEP}{target_label}{SEPN}")
async with connection as conn:
transaction = conn.transaction() if transaction else aecho()
async with transaction:
result = None
if table_name:
result = await self.backend.copy_from_table(
conn, table_name, **kwargs)
else:
result = await self.backend.copy_from_query(
conn, query, params, **kwargs)
if close:
if hasattr(close, "close"):
# close passed in object
output = close
else:
# close output object
output = kwargs.get("output")
if getattr(output, "close"):
output.close()
return result
async def copy_to(self, **kwargs):
pool = await self.pool
table_name = kwargs.pop("table_name", None)
transaction = kwargs.pop("transaction", False)
schema_name = kwargs.get('schema_name', None)
columns = ', '.join(kwargs.get('columns', []))
columns = f' ({columns})' if columns else ''
target_label = f"{schema_name}.{table_name}{columns}" if schema_name else table_name
connection = kwargs.pop("connection", None) or self._connection
connection = aecho(connection) if connection else pool.acquire()
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}copy to {target_label}{SEPN}",
True):
raise Exception(f"{self}: copy_to aborted")
else:
self.log(f"{self}: copy_to{SEP}{target_label}{SEPN}")
async with connection as conn:
transaction = conn.transaction() if transaction else aecho()
async with transaction:
return await self.backend.copy_to_table(
conn, table_name, **kwargs)