def dropTable(conn, tableName, dbName=None, mustExist=True):
"""
Drop table <tableName> in database <dbName>.
@param tableName Table name.
@param dbName Database name.
@param mustExist Flag indicating what to do if the database does not exist.
If <dbName> is not set, the current database name will be used.
Raises NoSuchTableError if the table does not exist and the mustExist flag
is set to True.
Raises sqlalchemy exceptions.
"""
if conn.engine.url.get_backend_name() == "mysql":
dbNameStr = "`%s`." % dbName if dbName is not None else ""
try:
conn.execute("DROP TABLE %s`%s`" % (dbNameStr, tableName))
except DBAPIError as e:
if e.orig.args[0] == MySqlErr.ER_BAD_TABLE_ERROR:
if mustExist:
raise NoSuchTableError(dbNameStr + tableName)
return
raise
else:
raise NoSuchModuleError(conn.engine.url.get_backend_name())
def createTableLike(conn, dbName, tableName, templDb, templTable):
"""
Create table <dbName>.<tableName> like <templDb>.<templTable>
@param conn Database connection or engine.
@param dbName Name of the database where the tables should be created
@param tableName Name of the table to create
@param templDb Name of the database where the template table is
@param templTable Name of the template table
Raises TableExistsError if the to-be-created table already exists.
Raises NoSuchTableError if the template table does not exists.
Raises sqlalchemy exceptions.
"""
if conn.engine.url.get_backend_name() == "mysql":
query = "CREATE TABLE {0}.{1} LIKE {2}.{3}".format(
dbName, tableName, templDb, templTable)
try:
conn.execute(query)
except OperationalError as e:
if e.orig.args[0] == MySqlErr.ER_TABLE_EXISTS_ERROR:
raise TableExistsError("CREATE TABLE LIKE",
dbName + '.' + tableName, e.orig)
raise
except ProgrammingError as e:
if e.orig.args[0] == MySqlErr.ER_NO_SUCH_TABLE:
raise NoSuchTableError("CREATE TABLE LIKE", templTable, e.orig)
raise
else:
raise NoSuchModuleError(conn.engine.url.get_backend_name())
def _get_table(self, connection, table_name):
project, dataset, table_name = self._split_table_name(table_name)
table = connection.connection._client.dataset(
dataset, project=project).table(table_name)
try:
table.reload()
except NotFound as e:
raise NoSuchTableError(table_name)
return table
def get_columns(self, connection, table_name, schema=None, **kw):
raw_connection = self._raw_connection(connection)
schema = schema if schema else raw_connection.schema_name
query = """
SELECT
table_schema,
table_name,
column_name,
data_type,
is_nullable,
column_default,
ordinal_position,
comment
FROM information_schema.columns
WHERE table_schema = '{schema}'
AND table_name = '{table}'
""".format(
schema=schema, table=table_name
)
retry_config = raw_connection.retry_config
retry = tenacity.Retrying(
retry=retry_if_exception(
lambda exc: self._retry_if_data_catalog_exception(
exc, schema, table_name
)
),
stop=stop_after_attempt(retry_config.attempt),
wait=wait_exponential(
multiplier=retry_config.multiplier,
max=retry_config.max_delay,
exp_base=retry_config.exponential_base,
),
reraise=True,
)
try:
return [
{
"name": row.column_name,
"type": _TYPE_MAPPINGS.get(
self._get_column_type(row.data_type), NULLTYPE
),
"nullable": True if row.is_nullable == "YES" else False,
"default": row.column_default
if not self._is_nan(row.column_default)
else None,
"ordinal_position": row.ordinal_position,
"comment": row.comment,
}
for row in retry(connection.execute, query).fetchall()
]
except OperationalError as e:
if not self._retry_if_data_catalog_exception(e, schema, table_name):
raise_from(NoSuchTableError(table_name), e)
else:
raise e
def _get_query_data(query_func: Callable[[Engine, SiteList], pd.DataFrame],
epd_engine: Engine, site_numbers: SiteList) -> pd.DataFrame:
"""Run query spefified by `query_func` against the database."""
try:
query_data = query_func(epd_engine, site_numbers)
except NoSuchTableError as e:
logging.error('Query could not find table. Check EPD has been '
'restored')
raise NoSuchTableError(e)
return query_data
def _get_table(self, connection, table_name, schema=None):
if isinstance(connection, Engine):
connection = connection.connect()
client = connection.connection._client
table_ref = self._table_reference(schema, table_name, client.project)
try:
table = client.get_table(table_ref)
except NotFound:
raise NoSuchTableError(table_name)
return table
def _get_table(self, connection, table_name, schema=None, **kw):
raw_connection = self._raw_connection(connection)
schema = schema if schema else raw_connection.schema_name
with raw_connection.connection.cursor() as cursor:
try:
return cursor.get_table_metadata(table_name,
schema_name=schema)
except pyathena.error.OperationalError as exc:
cause = exc.__cause__
if (isinstance(cause, botocore.exceptions.ClientError)
and cause.response["Error"]["Code"]
== "MetadataException"):
raise NoSuchTableError(table_name) from exc
raise
def _get_table(self, connection, table_name, schema=None):
if isinstance(connection, Engine):
connection = connection.connect()
project, dataset, table_name_prepared = self._split_table_name(table_name)
if dataset is None and schema is not None:
dataset = schema
table_name_prepared = table_name
table = connection.connection._client.dataset(dataset, project=project).table(table_name_prepared)
try:
t = connection.connection._client.get_table(table)
except NotFound as e:
raise NoSuchTableError(table_name)
return t
def get_columns(self, connection, table_name, schema=None, **kw):
schema = schema if schema else connection.connection.schema_name
query = """
SELECT
table_schema,
table_name,
column_name,
data_type,
is_nullable,
column_default,
ordinal_position,
comment
FROM information_schema.columns
WHERE table_schema = '{schema}'
AND table_name = '{table}'
""".format(schema=schema, table=table_name)
retry = tenacity.Retrying(
retry=retry_if_exception(
lambda exc: self._retry_if_data_catalog_exception(
exc, schema, table_name)),
stop=stop_after_attempt(connection.connection.retry_attempt),
wait=wait_exponential(
multiplier=connection.connection.retry_multiplier,
max=connection.connection.retry_max_delay,
exp_base=connection.connection.retry_exponential_base),
reraise=True)
try:
return [{
'name':
row.column_name,
'type':
_TYPE_MAPPINGS.get(self._get_column_type(row.data_type),
NULLTYPE),
'nullable':
True if row.is_nullable == 'YES' else False,
'default':
row.column_default,
'ordinal_position':
row.ordinal_position,
'comment':
row.comment,
} for row in retry(connection.execute, query).fetchall()]
except OperationalError as e:
if not self._retry_if_data_catalog_exception(
e, schema, table_name):
raise_from(NoSuchTableError(table_name), e)
else:
raise e
def _get_table(self, connection, table_name, schema=None):
if isinstance(connection, Engine):
connection = connection.connect()
client = connection.connection._client
project_id, dataset_id, table_id = self._split_table_name(table_name)
project_id = project_id or client.project
dataset_id = dataset_id or schema or self.dataset_id
table_ref = TableReference.from_string("{}.{}.{}".format(
project_id, dataset_id, table_id))
try:
table = client.get_table(table_ref)
except NotFound:
raise NoSuchTableError(table_name)
return table
def get_columns(self, connection, table_name, schema=None, **kw):
full_table = table_name
if schema:
full_table = schema + '.' + table_name
conn = self._get_dbapi_connection(connection)
table = conn.odps.get_table(full_table)
result = []
try:
for col in table.schema.columns:
col_type = _odps_type_to_sqlalchemy_type[type(col.type)]
result.append({
'name': col.name,
'type': col_type,
'nullable': True,
'default': None,
})
except NoSuchObject as e:
# convert ODPSError to SQLAlchemy NoSuchTableError
raise NoSuchTableError(str(e))
return result
def reflect(self, table_names):
"""Get table metadata through reflection.
sqlalchemy already provides a reflect method, but it will stop at the
first failure, while this method will try to get as much as possible.
:param table_names: Table names to inspect
:type table_names: list(str)
"""
inspector = inspect(self.engine)
for table_name in table_names:
columns = []
for column_data in inspector.get_columns(table_name):
# Rename 'type' to 'type_' to create column object
column_type = column_data.pop('type', None)
column_data['type_'] = column_type
columns.append(Column(**column_data))
if not columns:
raise NoSuchTableError(table_name)
Table(table_name, self.metadata, *columns)
def get_or_create_records_in_table(self, tablename, records):
self.meta.reflect()
try:
table = self.meta.tables[tablename]
except KeyError:
raise NoSuchTableError(tablename)
res_list = []
try:
for record in records:
res = self.session.query(table).filter_by(**record).first()
if not res:
self.session.execute(table.insert(), [record])
res = self.session.query(table).filter_by(**record).first()
res_list.append((res, True))
else:
res_list.append((res, False))
except IntegrityError as e:
self.rollback()
raise e
return res_list
def __getattr__(self, name):
try:
return self.meta.tables[name]
except KeyError:
raise NoSuchTableError(name)
def test_health_check_failed(client, db):
with patch('cachito.web.app.db.session.execute') as mock_execute:
mock_execute.side_effect = NoSuchTableError()
rv = client.get('/healthcheck')
assert rv.status_code == 500
assert rv.json.keys() == {'error'}
