'''Memory store using SQL databases supported by sqlalchemy. (eg. Postgresql)'''

def __init__(self, configs: Dict = MEMORYDB_CONFIG):

'''Initialize memory storage'''

self.engine = create_engine(url=configs['connect_str'])

self.meta = MetaData()

def add_history(self, project: str, session_id: str, messages: List[dict]):

project_table = self._create_table_if_not_exists(project)

data = [

{'session_id': session_id, 'message': json.dumps(

self._message_to_dict(message))}

for message in messages

]

with self.engine.connect() as conn:

query = project_table.insert()

conn.execute(query, data)

conn.commit()

def get_history(self, project: str, session_id: str):

if self.check(project):

project_table = Table(project, self.meta,

autoload_with=self.engine,

extend_existing=True

)

query = project_table.select().where(project_table.c.session_id == session_id)

messages = []

with self.engine.connect() as conn:

for row in conn.execute(query):

message = row.message

if isinstance(message, str):

message = json.loads(message)

message = self._message_from_dict(message)

messages.append(message)

else:

messages = []

return messages

def _create_table_if_not_exists(self, project):

if not self.check(project):

project_table = Table(project, self.meta,

Column('id', Integer, primary_key=True,

autoincrement=True),

Column('session_id', String, nullable=False),

Column('message', JSON, nullable=False),

extend_existing=True

)

self.meta.create_all(self.engine)

project_table = Table(project, self.meta, autoload_with=self.engine, extend_existing=True)

return project_table

def drop(self, project, session_id=None):

if self.check(project):

project_table = Table(project, self.meta,

autoload_with=self.engine, extend_existing=True)

if session_id and len(session_id) > 0:

query = project_table.delete().where(project_table.c.session_id == session_id)

with self.engine.connect() as conn:

conn.execute(query)

conn.commit()

else:

query = project_table.drop(self.engine)

if not session_id or len(session_id) == 0:

existence = self.check(project)

assert not existence, f'Failed to drop table {project}.'

def check(self, project):

return inspect(self.engine).has_table(project)

@staticmethod

def _message_from_dict(message: dict):

return tuple(message['data'])

@staticmethod

def _message_to_dict(message: tuple) -> dict: