Skip to content
This repository was archived by the owner on Sep 23, 2024. It is now read-only.

Commit 9c9db6a

Browse files
gwaxgwax
committed
Add support for transferring postgres range types as json objects
1 parent a97d7d8 commit 9c9db6a

File tree

6 files changed

+108
-6
lines changed

6 files changed

+108
-6
lines changed

‎tap_postgres/db.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,20 @@ def selected_value_to_singer_value_impl(elem, sql_datatype):
159159
cleaned_elem = elem
160160
else:
161161
raise Exception(f"do not know how to marshall a dict if its not an hstore or json: {sql_datatype}")
162+
elif 'range' in sql_datatype:
163+
child_type = {
164+
'int4range': 'int4',
165+
'int8range': 'int8',
166+
'numrange': 'numeric',
167+
'tsrange': 'timestamp without time zone',
168+
'tstzrange': 'timestamp with time zone',
169+
'daterange': 'date',
170+
}[sql_datatype]
171+
cleaned_elem = {
172+
'lower': selected_value_to_singer_value_impl(elem.lower, child_type),
173+
'upper': selected_value_to_singer_value_impl(elem.upper, child_type),
174+
'bounds': elem._bounds,
175+
}
162176
else:
163177
raise Exception(
164178
f"do not know how to marshall value of class( {elem.__class__} ) and sql_datatype ( {sql_datatype} )")

‎tap_postgres/discovery_utils.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -280,6 +280,10 @@ def schema_for_column_datatype(col):
280280
schema['type'] = nullable_column('string', col.is_primary_key)
281281
return schema
282282

283+
if 'range' in data_type:
284+
schema['type'] = nullable_column('object', col.is_primary_key)
285+
return schema
286+
283287
return schema
284288

285289

‎tests/integration/test_discovery.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,72 @@ def test_catalog(self):
348348
'type': 'object'},
349349
stream_dict.get('schema'))
350350

351+
class TestRangeTables(unittest.TestCase):
352+
maxDiff = None
353+
table_name = 'CHICKEN TIMES'
354+
355+
def setUp(self):
356+
table_spec = {"columns": [{"name": 'our_int_range', "type": "int4range"},
357+
{"name": 'our_tstz_range', "type": "tstzrange"}],
358+
"name": TestRangeTables.table_name}
359+
ensure_test_table(table_spec)
360+
361+
def test_catalog(self):
362+
conn_config = get_test_connection_config()
363+
364+
my_stdout = io.StringIO()
365+
with contextlib.redirect_stdout(my_stdout):
366+
streams = tap_postgres.do_discovery(conn_config)
367+
368+
chicken_streams = [s for s in streams if s['tap_stream_id'] == 'public-CHICKEN TIMES']
369+
self.assertEqual(len(chicken_streams), 1)
370+
stream_dict = chicken_streams[0]
371+
372+
stream_dict.get('metadata').sort(key=lambda md: md['breadcrumb'])
373+
374+
self.assertEqual(metadata.to_map(stream_dict.get('metadata')),
375+
{(): {'database-name': 'postgres',
376+
'is-view': False,
377+
'row-count': 0,
378+
'schema-name': 'public',
379+
'table-key-properties': []},
380+
('properties', 'our_int_range'): {'inclusion': 'available',
381+
'selected-by-default': True,
382+
'sql-datatype': 'int4range'},
383+
('properties', 'our_tstz_range'): {'inclusion': 'available',
384+
'selected-by-default': True,
385+
'sql-datatype': 'tstzrange'}})
386+
387+
self.assertEqual(stream_dict.get('schema'),
388+
{'definitions': {'sdc_recursive_boolean_array': {'items': {'$ref': '#/definitions/sdc_recursive_boolean_array'},
389+
'type': ['null',
390+
'boolean',
391+
'array']},
392+
'sdc_recursive_integer_array': {'items': {'$ref': '#/definitions/sdc_recursive_integer_array'},
393+
'type': ['null',
394+
'integer',
395+
'array']},
396+
'sdc_recursive_number_array': {'items': {'$ref': '#/definitions/sdc_recursive_number_array'},
397+
'type': ['null',
398+
'number',
399+
'array']},
400+
'sdc_recursive_object_array': {'items': {'$ref': '#/definitions/sdc_recursive_object_array'},
401+
'type': ['null',
402+
'object',
403+
'array']},
404+
'sdc_recursive_string_array': {'items': {'$ref': '#/definitions/sdc_recursive_string_array'},
405+
'type': ['null',
406+
'string',
407+
'array']},
408+
'sdc_recursive_timestamp_array': {'format': 'date-time',
409+
'items': {'$ref': '#/definitions/sdc_recursive_timestamp_array'},
410+
'type': ['null',
411+
'string',
412+
'array']}},
413+
'properties': {'our_int_range': {'type': ['null', 'object']},
414+
'our_tstz_range': {'type': ['null', 'object']}},
415+
'type': 'object'})
416+
351417

352418
class TestUUIDTables(unittest.TestCase):
353419
maxDiff = None

‎tests/integration/test_logical_replication.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ def setUpClass(cls) -> None:
2424
{"name": 'colour', "type": "character varying"},
2525
{"name": 'timestamp_ntz', "type": "timestamp without time zone"},
2626
{"name": 'timestamp_tz', "type": "timestamp with time zone"},
27+
{"name": 'int_range', "type": "int4range"},
2728
],
2829
"name": cls.table_name}
2930

@@ -53,19 +54,22 @@ def test_logical_replication(self):
5354
'name': 'betty',
5455
'colour': 'blue',
5556
'timestamp_ntz': '2020-09-01 10:40:59',
56-
'timestamp_tz': '2020-09-01 00:50:59+02'
57+
'timestamp_tz': '2020-09-01 00:50:59+02',
58+
'int_range': '[1,2)',
5759
},
5860
{
5961
'name': 'smelly',
6062
'colour': 'brown',
6163
'timestamp_ntz': '2020-09-01 10:40:59 BC',
62-
'timestamp_tz': '2020-09-01 00:50:59+02 BC'
64+
'timestamp_tz': '2020-09-01 00:50:59+02 BC',
65+
'int_range': '[2,5)',
6366
},
6467
{
6568
'name': 'pooper',
6669
'colour': 'green',
6770
'timestamp_ntz': '30000-09-01 10:40:59',
68-
'timestamp_tz': '10000-09-01 00:50:59+02'
71+
'timestamp_tz': '10000-09-01 00:50:59+02',
72+
'int_range': '[100,)',
6973
}
7074
]
7175

@@ -101,6 +105,7 @@ def test_logical_replication(self):
101105
'name': 'betty',
102106
'timestamp_ntz': '2020-09-01T10:40:59+00:00',
103107
'timestamp_tz': '2020-08-31T22:50:59+00:00',
108+
'int_range': {'lower': 1, 'upper': 2, 'bounds': '[)'},
104109
},
105110
'time_extracted': unittest.mock.ANY,
106111
'version': unittest.mock.ANY
@@ -114,6 +119,7 @@ def test_logical_replication(self):
114119
'name': 'smelly',
115120
'timestamp_ntz': '9999-12-31T23:59:59.999000+00:00',
116121
'timestamp_tz': '9999-12-31T23:59:59.999000+00:00',
122+
'int_range': {'lower': 2, 'upper': 5, 'bounds': '[)'},
117123
},
118124
'time_extracted': unittest.mock.ANY,
119125
'version': unittest.mock.ANY
@@ -127,6 +133,7 @@ def test_logical_replication(self):
127133
'name': 'pooper',
128134
'timestamp_ntz': '9999-12-31T23:59:59.999000+00:00',
129135
'timestamp_tz': '9999-12-31T23:59:59.999000+00:00',
136+
'int_range': {'lower': 100, 'upper': None, 'bounds': '[)'},
130137
},
131138
'time_extracted': unittest.mock.ANY,
132139
'version': unittest.mock.ANY
@@ -175,6 +182,7 @@ def test_logical_replication(self):
175182
'name': 'betty',
176183
'timestamp_ntz': '2020-09-01T10:40:59+00:00',
177184
'timestamp_tz': '2020-08-31T22:50:59+00:00',
185+
'int_range': '[1,2)',
178186
},
179187
'time_extracted': unittest.mock.ANY,
180188
'version': unittest.mock.ANY,
@@ -190,6 +198,7 @@ def test_logical_replication(self):
190198
'nice_flag': False,
191199
'timestamp_ntz': '2022-09-01T10:40:59+00:00',
192200
'timestamp_tz': '9999-12-31T23:59:59.999+00:00',
201+
'int_range': None,
193202
},
194203
'time_extracted': unittest.mock.ANY,
195204
'version': unittest.mock.ANY,

‎tests/integration/test_unsupported_pk.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@ def setUp(self):
2828
{"name": "circle_col", "type": "circle"},
2929
{"name": "xml_col", "type": "xml"},
3030
{"name": "composite_col", "type": "person_composite"},
31-
{"name": "int_range_col", "type": "int4range"},
3231
],
3332
"name": Unsupported.table_name}
3433
with get_test_connection() as conn:
@@ -53,7 +52,6 @@ def test_catalog(self):
5352
('properties', 'bit_string_col'): {'sql-datatype': 'bit(5)', 'selected-by-default': False, 'inclusion': 'unsupported'},
5453
('properties', 'line_col'): {'sql-datatype': 'line', 'selected-by-default': False, 'inclusion': 'unsupported'},
5554
('properties', 'xml_col'): {'sql-datatype': 'xml', 'selected-by-default': False, 'inclusion': 'unsupported'},
56-
('properties', 'int_range_col'): {'sql-datatype': 'int4range', 'selected-by-default': False, 'inclusion': 'unsupported'},
5755
('properties', 'circle_col'): {'sql-datatype': 'circle', 'selected-by-default': False, 'inclusion': 'unsupported'},
5856
('properties', 'polygon_col'): {'sql-datatype': 'polygon', 'selected-by-default': False, 'inclusion': 'unsupported'},
5957
('properties', 'box_col'): {'sql-datatype': 'box', 'selected-by-default': False, 'inclusion': 'unsupported'},

‎tests/unit/test_db.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import unittest
33

44
import datetime
5-
5+
import psycopg2.extras
66
from tap_postgres import db
77

88

@@ -171,6 +171,17 @@ def test_selected_value_to_singer_value_impl_with_non_empty_jsonb_returns_equiva
171171
'key2': [{'kk': 'yo'}, {}]
172172
}, output)
173173

174+
def test_selected_value_to_singer_value_impl_with_intrange(self):
175+
output = db.selected_value_to_singer_value_impl(psycopg2.extras.NumericRange(1,4,'[)'), 'int4range')
176+
177+
self.assertEqual({'lower': 1, 'upper': 4, 'bounds': '[)'}, output)
178+
179+
def test_selected_value_to_singer_value_impl_with_tstzrange(self):
180+
output = db.selected_value_to_singer_value_impl(psycopg2.extras.DateTimeTZRange(datetime.datetime(2020,1,2,3,4,5,tzinfo=datetime.timezone.utc),None,'[)'), 'tstzrange')
181+
182+
self.assertEqual({'lower': '2020-01-02T03:04:05+00:00', 'upper': None, 'bounds': '[)'}, output)
183+
184+
174185
def test_fully_qualified_column_name(self):
175186
schema = 'foo_schema'
176187
table = 'foo_table'

0 commit comments

Comments
 (0)