def __init__(self,
credentials,
table_name,
key_columns,
fixed_length_key=True):
super().__init__(fixed_length_key=fixed_length_key)
if not sqlalchemy:
raise ge_exceptions.DataContextError(
"ModuleNotFoundError: No module named 'sqlalchemy'")
if not self.fixed_length_key:
raise ge_exceptions.InvalidConfigError(
"DatabaseStoreBackend requires use of a fixed-length-key")
drivername = credentials.pop("drivername")
options = URL(drivername, **credentials)
self.engine = create_engine(options)
meta = MetaData()
self.key_columns = key_columns
# Dynamically construct a SQLAlchemy table with the name and column names we'll use
cols = []
for column in key_columns:
if column == "value":
raise ge_exceptions.InvalidConfigError(
"'value' cannot be used as a key_element name")
cols.append(Column(column, String, primary_key=True))
cols.append(Column("value", String))
try:
table = Table(table_name,
meta,
autoload=True,
autoload_with=self.engine)
# We do a "light" check: if the columns' names match, we will proceed, otherwise, create the table
if {str(col.name).lower()
for col in table.columns
} != (set(key_columns) | {"value"}):
raise ge_exceptions.StoreBackendError(
f"Unable to use table {table_name}: it exists, but does not have the expected schema."
)
except NoSuchTableError:
table = Table(table_name, meta, *cols)
try:
meta.create_all(self.engine)
except SQLAlchemyError as e:
raise ge_exceptions.StoreBackendError(
f"Unable to connect to table {table_name} because of an error. It is possible your table needs to be migrated to a new schema. SqlAlchemyError: {str(e)}"
)
self._table = table
def _set(self, key, value, allow_update=True):
cols = {k: v for (k, v) in zip(self.key_columns, key)}
cols["value"] = value
if allow_update:
if self.has_key(key):
ins = (
self._table.update()
.where(getattr(self._table.columns, self.key_columns[0]) == key[0])
.values(**cols)
)
else:
ins = self._table.insert().values(**cols)
else:
ins = self._table.insert().values(**cols)
try:
self.engine.execute(ins)
except IntegrityError as e:
if self._get(key) == value:
logger.info(f"Key {str(key)} already exists with the same value.")
else:
raise ge_exceptions.StoreBackendError(
f"Integrity error {str(e)} while trying to store key"
)
def remove_key(self, key):
delete_statement = self._table.delete().where(
and_(*(getattr(self._table.columns, key_col) == val
for key_col, val in zip(self.key_columns, key))))
try:
return self.engine.execute(delete_statement)
except SQLAlchemyError as e:
raise ge_exceptions.StoreBackendError(
f"Unable to delete key: got sqlalchemy error {str(e)}")
def _set(self, key, value, **kwargs):
cols = {k: v for (k, v) in zip(self.key_columns, key)}
cols["value"] = value
ins = self._table.insert().values(**cols)
try:
self.engine.execute(ins)
except IntegrityError as e:
if self._get(key) == value:
logger.info(
f"Key {str(key)} already exists with the same value.")
else:
raise ge_exceptions.StoreBackendError(
"Integrity error {str(e)} while trying to store key")
def __init__(
self,
table_name,
key_columns,
fixed_length_key=True,
credentials=None,
url=None,
connection_string=None,
engine=None,
store_name=None,
suppress_store_backend_id=False,
manually_initialize_store_backend_id: str = "",
**kwargs,
):
super().__init__(
fixed_length_key=fixed_length_key,
suppress_store_backend_id=suppress_store_backend_id,
manually_initialize_store_backend_id=
manually_initialize_store_backend_id,
store_name=store_name,
)
if not sa:
raise ge_exceptions.DataContextError(
"ModuleNotFoundError: No module named 'sqlalchemy'")
if not self.fixed_length_key:
raise ge_exceptions.InvalidConfigError(
"DatabaseStoreBackend requires use of a fixed-length-key")
self._schema_name = None
self._credentials = credentials
self._connection_string = connection_string
self._url = url
if engine is not None:
if credentials is not None:
logger.warning(
"Both credentials and engine were provided during initialization of SqlAlchemyExecutionEngine. "
"Ignoring credentials.")
self.engine = engine
elif credentials is not None:
self.engine = self._build_engine(credentials=credentials, **kwargs)
elif connection_string is not None:
self.engine = sa.create_engine(connection_string, **kwargs)
elif url is not None:
self.drivername = urlparse(url).scheme
self.engine = sa.create_engine(url, **kwargs)
else:
raise ge_exceptions.InvalidConfigError(
"Credentials, url, connection_string, or an engine are required for a DatabaseStoreBackend."
)
meta = MetaData(schema=self._schema_name)
self.key_columns = key_columns
# Dynamically construct a SQLAlchemy table with the name and column names we'll use
cols = []
for column in key_columns:
if column == "value":
raise ge_exceptions.InvalidConfigError(
"'value' cannot be used as a key_element name")
cols.append(Column(column, String, primary_key=True))
cols.append(Column("value", String))
try:
table = Table(table_name,
meta,
autoload=True,
autoload_with=self.engine)
# We do a "light" check: if the columns' names match, we will proceed, otherwise, create the table
if {str(col.name).lower()
for col in table.columns
} != (set(key_columns) | {"value"}):
raise ge_exceptions.StoreBackendError(
f"Unable to use table {table_name}: it exists, but does not have the expected schema."
)
except NoSuchTableError:
table = Table(table_name, meta, *cols)
try:
if self._schema_name:
self.engine.execute(
f"CREATE SCHEMA IF NOT EXISTS {self._schema_name};")
meta.create_all(self.engine)
except SQLAlchemyError as e:
raise ge_exceptions.StoreBackendError(
f"Unable to connect to table {table_name} because of an error. It is possible your table needs to be migrated to a new schema. SqlAlchemyError: {str(e)}"
)
self._table = table
# Initialize with store_backend_id
self._store_backend_id = None
self._store_backend_id = self.store_backend_id
# Gather the call arguments of the present function (include the "module_name" and add the "class_name"), filter
# out the Falsy values, and set the instance "_config" variable equal to the resulting dictionary.
self._config = {
"table_name": table_name,
"key_columns": key_columns,
"fixed_length_key": fixed_length_key,
"credentials": credentials,
"url": url,
"connection_string": connection_string,
"engine": engine,
"store_name": store_name,
"suppress_store_backend_id": suppress_store_backend_id,
"manually_initialize_store_backend_id":
manually_initialize_store_backend_id,
"module_name": self.__class__.__module__,
"class_name": self.__class__.__name__,
}
self._config.update(kwargs)
filter_properties_dict(properties=self._config,
clean_falsy=True,
inplace=True)
