Skip to content

Commit 99ea5ef

Browse files
committed
wip: sql source and reader
1 parent fc7f702 commit 99ea5ef

File tree

2 files changed

+61
-0
lines changed

2 files changed

+61
-0
lines changed

src/earthkit/data/readers/sql.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
from . import Reader
2+
3+
4+
class SQLReader(Reader):
5+
r"""Class representing CSV data"""
6+
7+
def __init__(self, source, sql_engine, sql, **kwargs):
8+
super().__init__(source, sql_engine)
9+
self._sql_engine = sql_engine
10+
self._sql = sql
11+
12+
def mutate_source(self):
13+
return self
14+
15+
def to_pandas(self, **kwargs):
16+
import pandas as pd
17+
18+
conn = self._sql_engine.connect()
19+
try:
20+
df = pd.read_sql(self._sql, conn, **kwargs)
21+
finally:
22+
conn.close()
23+
24+
return df
25+
26+
def to_xarray(self, **kwargs):
27+
28+
return self.to_pandas().to_xarray(**kwargs)
29+
30+
31+
def reader(source, path, *, magic=None, deeper_check=False, fwf=False, **kwargs):
32+
return

src/earthkit/data/sources/sql.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from ..readers.sql import SQLReader
2+
from . import Source
3+
4+
5+
def create_sqlalchemy_engine(backend, kwargs):
6+
from sqlalchemy import create_engine
7+
8+
if backend == "sqlite":
9+
return create_engine(f"sqlite:///{kwargs['path']}")
10+
else:
11+
raise NotImplementedError(f"Backend: {backend} is not implemented.")
12+
13+
14+
class SQLSource(Source):
15+
16+
def __init__(self, backend, backend_kwargs, sql, **kwargs):
17+
super().__init__(**kwargs)
18+
sql_engine = create_sqlalchemy_engine(backend, backend_kwargs)
19+
self._reader = SQLReader(self, sql_engine, sql, **kwargs)
20+
21+
def mutate(self):
22+
source = self._reader.mutate_source()
23+
if source not in (None, self):
24+
source._parent = self
25+
return source
26+
return self
27+
28+
29+
source = SQLSource

0 commit comments

Comments
 (0)