SQLAlchemy の IdentityMap の登録・利用条件

## はじめに

SQLAlchemy を用いて DB にデータ投入した後に、再び select をして返している時がよくあります。これは、本来避けるべきでしょうが、構造が複雑になってきて、このような構造になってしまうことがどうしてもあります。[SDGs が叫ばれる昨今、Python という地球環境によろしくない言語は抹殺されるべき](https://twitter.com/imos/status/1465185411651694598)という ~~言説~~ ネタが飛び交う中、このような構造になってしまった時に少しでも DB の負担を和らげることはできないだろうと思い調べたのが、 Identity Map でした。



## Identity Map とは

[こちらのダイアグラム](https://www.martinfowler.com/eaaCatalog/identityMap.html)を見ていただければわかると思いますが、データをプライマリーキーをキーとしてキャッシュとして保持しておき、取得するデータのプライマリーキーのオブジェクトがキャッシュが存在した場合は、キャッシュの値を用いる仕組みです。これにより、DB の負担を軽くできるメリットがあり、ほとんどの ORM に実装されています。もちろん、SQLAlchemy も例外ではなくIdentity Map があり、データをキャッシュしています。

ただし、SQLAlchemy では確実にキャッシュが使えるようになっているわけではなりません。[Identity Map](https://github.com/sqlalchemy/sqlalchemy/blob/400bf1716f/lib/sqlalchemy/orm/identity.py#L89) は弱参照で ORMとその状態管理をするオブジェクトを保持しており、その Identity Map も [Session](https://github.com/sqlalchemy/sqlalchemy/blob/400bf1716f/lib/sqlalchemy/orm/session.py#L1129) に対して保持されています。そのために GC のタイミングによっては Identity Map のデータがなくなっていることがあります。

## Identity Map の登録・利用条件について

### Identity Map に登録される条件

こちらは ORM で状態管理を行なっているもののみです。具体的には、ORM オブジェクトを作成して、DB に登録を行った場合とQuery 経由で ORM オブジェクトを取得した場合です。まずは、Session.flush を行いますが、 `UOWTransaction.finalize_flush_changes` を挟んで [`Session._register_persistent`](https://github.com/sqlalchemy/sqlalchemy/blob/146a349d81/lib/sqlalchemy/orm/session.py#L2366) を実行しています。そこで Identity Map に存在するオブジェクトをしている処理に移ります。[ここでの処理](https://github.com/sqlalchemy/sqlalchemy/blob/146a349d81/lib/sqlalchemy/orm/identity.py#L118-L131) は Identity Map に存在しない場合は、新規に登録し、存在する場合は、新たにデータを投入するという処理になっています。

次に、Query 経由で ORM オブジェクトを取得する処理においては、クエリ実行時には、SQL Statement に ORM オブジェクトを翻訳してからその SQL Statement を実行し、その戻り値に対して [ORM にマッピングする箇所](https://github.com/sqlalchemy/sqlalchemy/blob/146a349d81023805264f81643db50a5281da90da/lib/sqlalchemy/orm/context.py#L340) に移りうます。 ここからの流れは非常に複雑になり、割愛しますが 最終的には [identity map にデータを投入](https://github.com/sqlalchemy/sqlalchemy/blob/028f9b5cac/lib/sqlalchemy/orm/loading.py#L938)しています。

このように、ORM 経由で DB とのデータやりとりをおこなった際に登録する際に IdentityMap にデータが登録されていることがわかります。



### Identity Map を用いる条件

こちらは [`Session.get`](https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.Session.get) 経由で ORM オブジェクトを取得する時だけです。SQLAlchemy の Identity Map の [key](https://github.com/sqlalchemy/sqlalchemy/blob/0b95f0055be252b13e99b0a944466f60b5e367ff/lib/sqlalchemy/orm/util.py#L278) は ORM を定義したクラスとそのプライマリーキーをセットにしています。`Session.get` を用いるときには、ご存じの通り、プライマリーキーを指定して実行しています。なお、SQLAlchemy で `Query.get` が deprecated になりましたが、おそらく Identity Map を用いているのを明示的に表現するためかもしれません。

Query 経由でデータを取得する際には、プライマリーキーを指定したとしても Identity Map を用いることはなく、DB からデータを fetch されてしまいます。ただし、Query 経由で取得したオブジェクトのメモリアドレスは何度叩いても同一の結果が返ってきています。この点により Query 経由で ORM を取得した場合は Identity Map が用いられると思われる方もいるかもしれませんが、先述の通り、Query 経由で ORM を取得した際ではその対象となるメモリアドレスにオブジェクトを上書きしてしまうためです。

## 検証

これらのことを実際のコードで検証してみましょう。
まずは基本的な設定周りから書いてきます。

```python
from sqlalchemy import create_engine

from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy import Column, Integer, String


engine = create_engine('sqlite:///:memory:', echo=True)
DBSession = sessionmaker(engine)
db = DBSession()


Base = declarative_base()


class ATable(Base):
    __tablename__ = 'a_tables'

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String, nullable=True)


Base.metadata.create_all(bind=engine)
```

まず、ORM 経由で session.flush を行い、identity_map に登録してみましょう。

```python
print("register a ORM")
a1 = ATable(name='hogehoge')
db.add(a1)
db.flush()
print("printout identity_map")
print(db.identity_map._dict)
```

結果

```
register a ORM
2022-01-07 06:47:04,630 INFO sqlalchemy.engine.Engine BEGIN (implicit)
2022-01-07 06:47:04,631 INFO sqlalchemy.engine.Engine INSERT INTO a_tables (name) VALUES (?)
2022-01-07 06:47:04,631 INFO sqlalchemy.engine.Engine [generated in 0.00011s] ('hogehoge',)
printout identity_map
{(<class '__main__.ATable'>, (1,), None): <sqlalchemy.orm.state.InstanceState object at 0x10e1edc10>}
```
これで、Identity Map に登録されたのがわかると思います。
次に、`Session.execute` でデータ投入を行い、Identity Map に登録されないことを見ていきます。

```python
print("insert entity via executing sql statement")
db.execute("insert into a_tables(name) values ('hogehogehoge')")
print(db.identity_map._dict)
```

結果

```
insert entity via executing sql statement
2022-01-07 06:47:04,631 INFO sqlalchemy.engine.Engine insert into a_tables(name) values ('hogehogehoge')
2022-01-07 06:47:04,632 INFO sqlalchemy.engine.Engine [generated in 0.00016s] ()
{(<class '__main__.ATable'>, (1,), None): <sqlalchemy.orm.state.InstanceState object at 0x10e1edc10>}
```

SQL Statement 経由で `Session.execute` を実行した後に再び IdentityMap を参照していますが、Identity Map に存在するオブジェクトはひとつ前のものと同じものが格納されているとわかります。

次に、先ほど作成したデータを参照してみましょう。

```python
print("select inserted")
a2 = db.query(ATable).filter_by(id=2).first()
print(db.identity_map._dict)
```

結果

```
select inserted
2022-01-07 06:47:04,634 INFO sqlalchemy.engine.Engine SELECT a_tables.id AS a_tables_id, a_tables.name AS a_tables_name
FROM a_tables
WHERE a_tables.id = ?
 LIMIT ? OFFSET ?
2022-01-07 06:47:04,634 INFO sqlalchemy.engine.Engine [generated in 0.00012s] (2, 1, 0)
{(<class '__main__.ATable'>, (1,), None): <sqlalchemy.orm.state.InstanceState object at 0x10e1edc10>, (<class '__main__.ATable'>, (2,), None): <sqlalchemy.orm.state.InstanceState object at 0x10e2cb760>}
```

今回は先ほどとは違い、Identity Map に id = 2 のオブジェクトが登録されています。このように、Query 経由で ORM を取得すると Identity Map にデータが登録されます。

最後に、既存のデータを `Session.get` で取得してみます。

```python
print("Session.get to existing object")
a1_ = db.get(ATable, a1.id)
print(db.identity_map._dict)
```

結果

```
Session.get to existing object
{(<class '__main__.ATable'>, (1,), None): <sqlalchemy.orm.state.InstanceState object at 0x10e1edc10>, (<class '__main__.ATable'>, (2,), None): <sqlalchemy.orm.state.InstanceState object at 0x10e2cb760>}
```

こちらの場合では、SQLAlchemy が実行する Statement 文がログとして出力されていません。そのために、ここでは DB に格納されたデータを用いずに Identity Map に存在しているデータを取得してきていることがわかると思います。


## まとめ

* SQLAlchemy にも Identity Map が存在して DB の負担を軽減してくれている。
* しかしながら、そのデータを用いる時には `Session.get` を用いる必要がある。


## 所感

- Python を使っていて、弱参照を目にするとは思わなかった。

## References

- [利用したコード](https://gist.github.com/kou2kkkt1984/98768fa1764b62b2ddc8aaed2fbb49b7)