Merge pull request #3 from dorosch/develop

Release 0.4

Author: dorosch

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlalchemy-timescaledb"
-version = "0.3.1"
+version = "0.4"
 authors = [
   { name="Andrei Kliatsko", email="andrey.daraschenka@gmail.com" },
 ]
@@ -15,9 +15,11 @@ test = [
     "pytest==7.2.1",
     "pytest-cov==4.0.0",
     "pytest-factoryboy==2.5.1",
-    "sqlalchemy>=1.3",
+    "sqlalchemy[asyncio]>=1.3",
     "psycopg2-binary==2.9.5",
-    "alembic==1.9.4"
+    "alembic==1.9.4",
+    "asyncpg==0.27.0",
+    "pytest-asyncio==0.20.3"
 ]
 
 [project.urls]

sqlalchemy_timescaledb/dialect.py

@@ -1,4 +1,4 @@
-from sqlalchemy import schema
+from sqlalchemy import schema, event, DDL
 from sqlalchemy.dialects.postgresql.asyncpg import PGDialect_asyncpg
 from sqlalchemy.dialects.postgresql.base import PGDDLCompiler
 from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
@@ -19,15 +19,30 @@ def post_create_table(self, table):
         hypertable = table.kwargs.get('timescaledb_hypertable', {})
 
         if hypertable:
-            hypertable_sql = \
-                ';\n\nSELECT create_hypertable(' \
-                f"'{table.name}', '{hypertable['time_column_name']}'" \
-                ");"
-
-            return super().post_create_table(table) + hypertable_sql
+            event.listen(
+                table,
+                'after_create',
+                self.ddl_hypertable(
+                    table.name, hypertable
+                ).execute_if(
+                    dialect='timescaledb'
+                )
+            )
 
         return super().post_create_table(table)
 
+    @staticmethod
+    def ddl_hypertable(table_name, hypertable):
+        return DDL(
+            f"""
+            SELECT create_hypertable(
+                '{table_name}',
+                '{hypertable['time_column_name']}',
+                if_not_exists => TRUE
+            );
+            """
+        )
+
 
 class TimescaledbDialect:
     name = 'timescaledb'

tests/async/__init__.py

@@ -0,0 +1,26 @@
+import pytest_asyncio
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+
+from tests.models import Base, DATABASE_URL
+
+
+@pytest_asyncio.fixture
+def async_engine():
+    yield create_async_engine(
+        DATABASE_URL.set(drivername='timescaledb+asyncpg')
+    )
+
+
+@pytest_asyncio.fixture
+async def async_session(async_engine):
+    async with AsyncSession(async_engine) as session:
+        yield session
+
+
+@pytest_asyncio.fixture(autouse=True)
+async def setup(async_engine):
+    async with async_engine.begin() as connection:
+        await connection.run_sync(Base.metadata.create_all)
+    yield
+    async with async_engine.begin() as connection:
+        await connection.run_sync(Base.metadata.drop_all)

tests/async/conftest.py

@@ -0,0 +1,29 @@
+import pytest
+from sqlalchemy import text
+
+from tests.models import Metric, User
+
+
+@pytest.mark.asyncio
+class TestHypertable:
+    async def test_is_hypertable(self, async_session):
+        assert (await async_session.execute(
+            text(
+                f"""
+                    SELECT count(*)
+                    FROM _timescaledb_catalog.hypertable
+                    WHERE table_name = '{Metric.__tablename__}'
+                    """
+            )
+        )).scalar_one()
+
+    async def test_is_not_hypertable(self, async_session):
+        assert not (await async_session.execute(
+            text(
+                f"""
+                    SELECT count(*)
+                    FROM _timescaledb_catalog.hypertable
+                    WHERE table_name = '{User.__tablename__}'
+                """
+            )
+        )).scalar_one()

tests/async/test_hypertable.py

@@ -1,26 +1,33 @@
 import pytest
-from sqlalchemy import text
 from pytest_factoryboy import register
+from sqlalchemy import text, create_engine
+from sqlalchemy.orm import Session
 
-from tests.models import engine, Session, Base
-from tests.factories import MetricFactory
+from tests.factories import MetricFactory, FactorySession
+from tests.models import Base, DATABASE_URL
 
 register(MetricFactory)
 
 
+@pytest.fixture
+def engine():
+    yield create_engine(DATABASE_URL)
+
+
+@pytest.fixture
+def session(engine):
+    with Session(engine) as session:
+        yield session
+
+
 @pytest.fixture(autouse=True)
-def setup():
+def setup(engine):
+    FactorySession.configure(bind=engine)
     Base.metadata.create_all(bind=engine)
     yield
     Base.metadata.drop_all(bind=engine)
 
 
-@pytest.fixture
-def session():
-    yield Session
-    Session.close()
-
-
 @pytest.fixture
 def is_hypertable(session):
     def check_hypertable(table):

tests/conftest.py

@@ -1,13 +1,17 @@
 import factory
 from factory.fuzzy import FuzzyText
+from sqlalchemy import orm
 
-from tests.models import Metric, Session
+from tests.models import Metric
+
+FactorySession = orm.scoped_session(orm.sessionmaker())
 
 
 class MetricFactory(factory.alchemy.SQLAlchemyModelFactory):
     class Meta:
         model = Metric
-        sqlalchemy_session = Session
+        sqlalchemy_session = FactorySession
+        sqlalchemy_session_persistence = 'commit'
 
     name = FuzzyText()
     value = 0

tests/factories.py

@@ -13,12 +13,12 @@
 if config.config_file_name is not None:
     fileConfig(config.config_file_name)
 
-from models import Base #, database_url, Metric, User
+from models import Base, DATABASE_URL
+
 config.set_section_option(
     config.config_ini_section,
     "sqlalchemy.url",
-    # database_url.render_as_string(hide_password=False)
-    'timescaledb://user:password@0.0.0.0:5432/database'
+    DATABASE_URL.render_as_string(hide_password=False)
 )
 
 target_metadata = Base.metadata

tests/migrations/env.py

@@ -2,21 +2,18 @@
 import os
 
 from sqlalchemy import Column, String, DateTime, Float, Integer
-from sqlalchemy import create_engine
 from sqlalchemy.engine import URL
 from sqlalchemy.orm import declarative_base
-from sqlalchemy.orm import sessionmaker, scoped_session
 
-database_url = URL.create(
+DATABASE_URL = URL.create(
     host=os.environ.get('POSTGRES_HOST', '0.0.0.0'),
     port=os.environ.get('POSTGRES_PORT', 5432),
     username=os.environ.get('POSTGRES_USER', 'user'),
     password=os.environ.get('POSTGRES_PASSWORD', 'password'),
     database=os.environ.get('POSTGRES_DB', 'database'),
     drivername=os.environ.get('DRIVERNAME', 'timescaledb')
 )
-engine = create_engine(database_url)
-Session = scoped_session(sessionmaker(bind=engine))
+
 Base = declarative_base()
 
 

tests/models.py

@@ -1,20 +1,33 @@
 import os
+from pathlib import Path
 
 from alembic import command
 from alembic.config import Config
 
-from tests.models import Base, engine
+from tests.models import Base
 
 
 class TestAlembic:
     def setup_class(self):
+        # TODO: Disable output for alembic
         self.config = Config(
             os.path.join(os.path.dirname(__file__), 'alembic.ini')
         )
+        self.migration_versions_path = os.path.join(
+            os.path.dirname(__file__), 'migrations', 'versions'
+        )
 
-    def test_create_revision(self):
+    def test_create_revision(self, engine):
         Base.metadata.drop_all(bind=engine)
-        script = command.revision(self.config, autogenerate=True)
-        Base.metadata.create_all(bind=engine)
+        script = command.revision(
+            self.config, message='initial', autogenerate=True
+        )
+        migration_file = os.path.join(
+            self.migration_versions_path, f'{script.revision}_initial.py'
+        )
 
         assert script.down_revision is None
+        assert Path(migration_file).is_file()
+
+        Path(migration_file).unlink()
+        Base.metadata.create_all(bind=engine)