Skip to content

Commit

Permalink
Support customizable SYMBOL column
Browse files Browse the repository at this point in the history
QuestDB allows SYMBOL columns to specify:
1. Capacity
2. Cache: ON/OFF

This changeset is adding this to SQLAlchemy. It allows to define columns as this:
`instrument = Column(Symbol(capacity=10000, cache=False))`
  • Loading branch information
jerrinot committed Dec 13, 2024
1 parent da17d32 commit 919355f
Show file tree
Hide file tree
Showing 3 changed files with 139 additions and 1 deletion.
2 changes: 1 addition & 1 deletion src/examples/hello_world.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ class Signal(Base):
__table_args__ = (
qdbc.QDBTableEngine("signal", "ts", qdbc.PartitionBy.HOUR, is_wal=True),
)
source = Column(qdbc.Symbol)
source = Column(qdbc.Symbol(capacity=1024, cache=False))
value = Column(qdbc.Double)
ts = Column(qdbc.Timestamp, primary_key=True)

Expand Down
33 changes: 33 additions & 0 deletions src/questdb_connect/types.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from typing import Optional

import sqlalchemy

from .common import quote_identifier
Expand Down Expand Up @@ -121,9 +123,40 @@ class String(QDBTypeMixin):


class Symbol(QDBTypeMixin):
"""
QuestDB SYMBOL type implementation with support for capacity and cache parameters.
Example usage:
source = Column(Symbol(capacity=128, cache=True))
"""
__visit_name__ = "SYMBOL"
type_code = 12

def __init__(
self,
capacity: Optional[int] = None,
cache: Optional[bool] = None,
*args, **kwargs
):
super().__init__(*args, **kwargs)
self.capacity = capacity
self.cache = cache

def compile(self, dialect=None):
params = []

if self.capacity is not None:
params.append(f"CAPACITY {self.capacity}")
if self.cache is not None:
params.append("CACHE" if self.cache else "NOCACHE")

if params:
return f"{self.__visit_name__} {' '.join(params)}"
return self.__visit_name__

def column_spec(self, column_name):
return f"{quote_identifier(column_name)} {self.compile()}"


class Long256(QDBTypeMixin):
__visit_name__ = "LONG256"
Expand Down
105 changes: 105 additions & 0 deletions tests/test_types.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import re

import questdb_connect as qdbc
from questdb_connect.common import quote_identifier


def test_resolve_type_from_name():
Expand All @@ -16,6 +17,110 @@ def test_resolve_type_from_name():
assert isinstance(g_class(), qdbc.geohash_class(n))


def test_symbol_type():
# Test basic Symbol without parameters
symbol = qdbc.Symbol()
assert symbol.__visit_name__ == "SYMBOL"
assert symbol.compile() == "SYMBOL"
assert symbol.column_spec("test_col") == "\"test_col\" SYMBOL"

# Test Symbol with capacity
symbol_cap = qdbc.Symbol(capacity=128)
assert symbol_cap.compile() == "SYMBOL CAPACITY 128"
assert symbol_cap.column_spec("test_col") == "\"test_col\" SYMBOL CAPACITY 128"

# Test Symbol with cache true
symbol_cache = qdbc.Symbol(cache=True)
assert symbol_cache.compile() == "SYMBOL CACHE"
assert symbol_cache.column_spec("test_col") == "\"test_col\" SYMBOL CACHE"

# Test Symbol with cache false
symbol_nocache = qdbc.Symbol(cache=False)
assert symbol_nocache.compile() == "SYMBOL NOCACHE"
assert symbol_nocache.column_spec("test_col") == "\"test_col\" SYMBOL NOCACHE"

# Test Symbol with both parameters
symbol_full = qdbc.Symbol(capacity=256, cache=True)
assert symbol_full.compile() == "SYMBOL CAPACITY 256 CACHE"
assert symbol_full.column_spec("test_col") == "\"test_col\" SYMBOL CAPACITY 256 CACHE"

# Test inheritance and type resolution
assert isinstance(symbol, qdbc.QDBTypeMixin)
resolved_class = qdbc.resolve_type_from_name("SYMBOL")
assert resolved_class.__visit_name__ == "SYMBOL"
assert isinstance(symbol, resolved_class)

# Test that parameters don't affect type resolution
symbol_with_params = qdbc.Symbol(capacity=128, cache=True)
assert isinstance(symbol_with_params, resolved_class)
assert isinstance(resolved_class(), type(symbol_with_params))


def test_symbol_backward_compatibility():
"""Verify that the parametrized Symbol type maintains backward compatibility with older code."""
# Test all the ways Symbol type could be previously instantiated
symbol1 = qdbc.Symbol
symbol2 = qdbc.Symbol()

# Check that both work in column definitions
from sqlalchemy import Column, MetaData, Table

metadata = MetaData()
test_table = Table(
'test_table',
metadata,
Column('old_style1', symbol1), # Old style: direct class reference
Column('old_style2', symbol2), # Old style: basic instantiation
)

# Verify type resolution still works
for column in test_table.columns:
# Check inheritance
assert isinstance(column.type, qdbc.QDBTypeMixin)

# Check type resolution
resolved_class = qdbc.resolve_type_from_name("SYMBOL")
assert isinstance(column.type, resolved_class)

# Check SQL generation matches old behavior
assert column.type.compile() == "SYMBOL"
assert column.type.column_spec(column.name) == f"{quote_identifier(column.name)} SYMBOL"

def test_symbol_type_in_column():
# Test Symbol type in Column definition
from sqlalchemy import Column, MetaData, Table

metadata = MetaData()

# Create a test table with different Symbol column variations
test_table = Table(
'test_table',
metadata,
Column('basic_symbol', qdbc.Symbol()),
Column('symbol_with_capacity', qdbc.Symbol(capacity=128)),
Column('symbol_with_cache', qdbc.Symbol(cache=True)),
Column('symbol_with_nocache', qdbc.Symbol(cache=False)),
Column('symbol_full', qdbc.Symbol(capacity=256, cache=True))
)

# Get the create table SQL (implementation-dependent)
# This part might need adjustment based on your actual SQL compilation logic
for column in test_table.columns:
assert isinstance(column.type, qdbc.Symbol)
assert isinstance(column.type, qdbc.QDBTypeMixin)

if column.name == 'basic_symbol':
assert column.type.compile() == "SYMBOL"
elif column.name == 'symbol_with_capacity':
assert column.type.compile() == "SYMBOL CAPACITY 128"
elif column.name == 'symbol_with_cache':
assert column.type.compile() == "SYMBOL CACHE"
elif column.name == 'symbol_with_nocache':
assert column.type.compile() == "SYMBOL NOCACHE"
elif column.name == 'symbol_full':
assert column.type.compile() == "SYMBOL CAPACITY 256 CACHE"


def test_superset_default_mappings():
default_column_type_mappings = (
(re.compile("^BOOLEAN$", re.IGNORECASE), qdbc.Boolean),
Expand Down

0 comments on commit 919355f

Please sign in to comment.