From b8256849f545a115c6c8020da3ef4d98d6b1327d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mois=C3=A9s?= <7888669+moisses89@users.noreply.github.com>
Date: Wed, 11 Dec 2024 13:19:20 +0100
Subject: [PATCH] WIP

---
 app/datasources/db/models.py             | 34 +++++++++-
 app/services/abi.py                      |  0
 app/services/chain.py                    | 36 +++++++++++
 app/tests/db/factories.py                | 11 ++++
 app/tests/db/test_model.py               |  6 +-
 app/tests/routers/test_contracts.py      |  9 ++-
 migrations/versions/c52699d26409_init.py | 80 ++++++++++++++++++++++++
 migrations/versions/d0c5d72aa50b_init.py | 37 -----------
 8 files changed, 168 insertions(+), 45 deletions(-)
 create mode 100644 app/services/abi.py
 create mode 100644 app/services/chain.py
 create mode 100644 app/tests/db/factories.py
 create mode 100644 migrations/versions/c52699d26409_init.py
 delete mode 100644 migrations/versions/d0c5d72aa50b_init.py

diff --git a/app/datasources/db/models.py b/app/datasources/db/models.py
index 8223cd5..410369a 100644
--- a/app/datasources/db/models.py
+++ b/app/datasources/db/models.py
@@ -1,9 +1,37 @@
-from typing import Optional
+from sqlmodel import JSON, Column, Field, SQLModel, UniqueConstraint
 
-from sqlmodel import Field, SQLModel
+
+class AbiSource(SQLModel, table=True):
+    id: int | None = Field(default=None, primary_key=True)
+    name: str = Field(nullable=False)
+    url: str = Field(nullable=False)
+
+
+class Abi(SQLModel, table=True):
+    abi_hash: bytes = Field(nullable=False, primary_key=True)
+    relevance: int = Field(nullable=False, default=0)
+    abi_json: dict = Field(default_factory=dict, sa_column=Column(JSON))
+    source_id: int = Field(default=None, foreign_key="abisource.id")
+
+
+class Chain(SQLModel, table=True):
+    id: int = Field(primary_key=True)  # Chain ID
+    name: str = Field(nullable=False)
 
 
 class Contract(SQLModel, table=True):
+    __table_args__ = (
+        UniqueConstraint("address", "chain_id", name="address_chain_unique"),
+    )
+
     address: bytes = Field(nullable=False, primary_key=True)
     name: str = Field(nullable=False)
-    description: Optional[str] = None
+    display_name: str | None = None
+    description: str | None = None
+    trusted_for_delegate: bool = Field(nullable=False, default=False)
+    proxy: bool = Field(nullable=False, default=False)
+    fetch_retries: int = Field(nullable=False, default=0)
+    abi_id: bytes | None = Field(
+        nullable=True, default=None, foreign_key="abi.abi_hash"
+    )
+    chain_id: int = Field(default=None, foreign_key="chain.id")
diff --git a/app/services/abi.py b/app/services/abi.py
new file mode 100644
index 0000000..e69de29
diff --git a/app/services/chain.py b/app/services/chain.py
new file mode 100644
index 0000000..38b485e
--- /dev/null
+++ b/app/services/chain.py
@@ -0,0 +1,36 @@
+from typing import Sequence
+
+from sqlmodel import select
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+from app.datasources.db.database import get_database_session
+from app.datasources.db.models import Chain
+
+
+class ChainService:
+
+    @staticmethod
+    @get_database_session
+    async def get_all(session: AsyncSession) -> Sequence[Chain]:
+        """
+        Get all chains
+
+        :param session: passed by the decorator
+        :return:
+        """
+        result = await session.exec(select(Chain))
+        return result.all()
+
+    @staticmethod
+    @get_database_session
+    async def create(chain: Chain, session: AsyncSession) -> Chain:
+        """
+        Create a new chain
+
+        :param chain:
+        :param session:
+        :return:
+        """
+        session.add(chain)
+        await session.commit()
+        return chain
diff --git a/app/tests/db/factories.py b/app/tests/db/factories.py
new file mode 100644
index 0000000..8f65db7
--- /dev/null
+++ b/app/tests/db/factories.py
@@ -0,0 +1,11 @@
+from factory.alchemy import SQLAlchemyModelFactory
+from sqlmodel import Session
+
+from app.datasources.db.models import Chain
+
+
+class ChainFactory(SQLAlchemyModelFactory):
+    class Meta:
+        model = Chain
+        sqlalchemy_session_factory = Session
+        sqlalchemy_session_persistence = "commit"
diff --git a/app/tests/db/test_model.py b/app/tests/db/test_model.py
index 67191d2..3a2f653 100644
--- a/app/tests/db/test_model.py
+++ b/app/tests/db/test_model.py
@@ -2,14 +2,16 @@
 from sqlmodel.ext.asyncio.session import AsyncSession
 
 from app.datasources.db.database import database_session
-from app.datasources.db.models import Contract
+from app.datasources.db.models import Chain, Contract
 from app.tests.db.db_async_conn import DbAsyncConn
 
 
 class TestModel(DbAsyncConn):
     @database_session
     async def test_contract(self, session: AsyncSession):
-        contract = Contract(address=b"a", name="A Test Contracts")
+        chain = Chain(id=1, name="mainnet")
+        session.add(chain)
+        contract = Contract(address=b"a", name="A Test Contracts", chain_id=chain.id)
         session.add(contract)
         await session.commit()
         statement = select(Contract).where(Contract.address == b"a")
diff --git a/app/tests/routers/test_contracts.py b/app/tests/routers/test_contracts.py
index 4895869..5bb2961 100644
--- a/app/tests/routers/test_contracts.py
+++ b/app/tests/routers/test_contracts.py
@@ -1,10 +1,10 @@
 from fastapi.testclient import TestClient
 
 from sqlmodel.ext.asyncio.session import AsyncSession
-
 from ...datasources.db.database import database_session
-from ...datasources.db.models import Contract
+from ...datasources.db.models import Chain, Contract
 from ...main import app
+from ...services.chain import ChainService
 from ...services.contract import ContractService
 from ..db.db_async_conn import DbAsyncConn
 
@@ -16,9 +16,12 @@ class TestRouterContract(DbAsyncConn):
     def setUpClass(cls):
         cls.client = TestClient(app)
 
+
     @database_session
     async def test_view_contracts(self, session: AsyncSession):
-        contract = Contract(address=b"a", name="A Test Contracts")
+        chain = Chain(id=1, name="mainnet")
+        await ChainService.create(chain)
+        contract = Contract(address=b"a", name="A Test Contracts", chain_id=1)
         expected_response = {
             "name": "A Test Contracts",
             "description": None,
diff --git a/migrations/versions/c52699d26409_init.py b/migrations/versions/c52699d26409_init.py
new file mode 100644
index 0000000..1f956f3
--- /dev/null
+++ b/migrations/versions/c52699d26409_init.py
@@ -0,0 +1,80 @@
+"""init
+
+Revision ID: c52699d26409
+Revises:
+Create Date: 2024-12-11 10:58:44.921603
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "c52699d26409"
+down_revision: Union[str, None] = None
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "abisource",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("url", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_table(
+        "chain",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_table(
+        "abi",
+        sa.Column("abi_hash", sa.LargeBinary(), nullable=False),
+        sa.Column("relevance", sa.Integer(), nullable=False),
+        sa.Column("abi_json", sa.JSON(), nullable=True),
+        sa.Column("source_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["source_id"],
+            ["abisource.id"],
+        ),
+        sa.PrimaryKeyConstraint("abi_hash"),
+    )
+    op.create_table(
+        "contract",
+        sa.Column("address", sa.LargeBinary(), nullable=False),
+        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("display_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column("description", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column("trusted_for_delegate", sa.Boolean(), nullable=False),
+        sa.Column("proxy", sa.Boolean(), nullable=False),
+        sa.Column("fetch_retries", sa.Integer(), nullable=False),
+        sa.Column("abi_id", sa.LargeBinary(), nullable=True),
+        sa.Column("chain_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["abi_id"],
+            ["abi.abi_hash"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["chain_id"],
+            ["chain.id"],
+        ),
+        sa.PrimaryKeyConstraint("address"),
+        sa.UniqueConstraint("address", "chain_id", name="address_chain_unique"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("contract")
+    op.drop_table("abi")
+    op.drop_table("chain")
+    op.drop_table("abisource")
+    # ### end Alembic commands ###
diff --git a/migrations/versions/d0c5d72aa50b_init.py b/migrations/versions/d0c5d72aa50b_init.py
deleted file mode 100644
index d7fd35f..0000000
--- a/migrations/versions/d0c5d72aa50b_init.py
+++ /dev/null
@@ -1,37 +0,0 @@
-"""init
-
-Revision ID: d0c5d72aa50b
-Revises:
-Create Date: 2024-12-03 14:16:45.445610
-
-"""
-
-from typing import Sequence, Union
-
-import sqlalchemy as sa
-import sqlmodel
-from alembic import op
-
-# revision identifiers, used by Alembic.
-revision: str = "d0c5d72aa50b"
-down_revision: Union[str, None] = None
-branch_labels: Union[str, Sequence[str], None] = None
-depends_on: Union[str, Sequence[str], None] = None
-
-
-def upgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
-    op.create_table(
-        "contract",
-        sa.Column("address", sa.LargeBinary(), nullable=False),
-        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
-        sa.Column("description", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
-        sa.PrimaryKeyConstraint("address"),
-    )
-    # ### end Alembic commands ###
-
-
-def downgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_table("contract")
-    # ### end Alembic commands ###