meerschaum.utils.dtypes.sql

  1#! /usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3# vim:fenc=utf-8
  4
  5"""
  6Utility functions for working with SQL data types.
  7"""
  8
  9from __future__ import annotations
 10from meerschaum.utils.typing import Dict, Union, Tuple, Optional
 11
 12NUMERIC_PRECISION_FLAVORS: Dict[str, Tuple[int, int]] = {
 13    'mariadb': (38, 20),
 14    'mysql': (38, 20),
 15    'mssql': (28, 10),
 16}
 17NUMERIC_AS_TEXT_FLAVORS = {'sqlite', 'duckdb'}
 18TIMEZONE_NAIVE_FLAVORS = {'oracle', 'mysql', 'mariadb'}
 19
 20### MySQL doesn't allow for casting as BIGINT, so this is a workaround.
 21DB_FLAVORS_CAST_DTYPES = {
 22    'mariadb': {
 23        'BIGINT': 'DECIMAL',
 24        'TINYINT': 'SIGNED INT',
 25        'TEXT': 'CHAR(10000) CHARACTER SET utf8',
 26        'BOOL': 'SIGNED INT',
 27        'BOOLEAN': 'SIGNED INT',
 28        'DOUBLE PRECISION': 'DECIMAL',
 29        'DOUBLE': 'DECIMAL',
 30        'FLOAT': 'DECIMAL',
 31    },
 32    'mysql': {
 33        'BIGINT': 'DECIMAL',
 34        'TINYINT': 'SIGNED INT',
 35        'TEXT': 'CHAR(10000) CHARACTER SET utf8',
 36        'INT': 'SIGNED INT',
 37        'INTEGER': 'SIGNED INT',
 38        'BOOL': 'SIGNED INT',
 39        'BOOLEAN': 'SIGNED INT',
 40        'DOUBLE PRECISION': 'DECIMAL',
 41        'DOUBLE': 'DECIMAL',
 42        'FLOAT': 'DECIMAL',
 43    },
 44    'sqlite': {
 45        'BOOLEAN': 'INTEGER',
 46        'REAL': 'FLOAT',
 47    },
 48    'oracle': {
 49        'NVARCHAR(2000)': 'NVARCHAR2(2000)',
 50        'NVARCHAR': 'NVARCHAR2(2000)',
 51        'NVARCHAR2': 'NVARCHAR2(2000)',
 52        'CHAR': 'CHAR(36)',  # UUID columns
 53    },
 54    'mssql': {
 55        'NVARCHAR COLLATE "SQL Latin1 General CP1 CI AS"': 'NVARCHAR(MAX)',
 56        'NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS"': 'NVARCHAR(MAX)',
 57        'VARCHAR COLLATE "SQL Latin1 General CP1 CI AS"': 'NVARCHAR(MAX)',
 58        'VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS"': 'NVARCHAR(MAX)',
 59        'NVARCHAR': 'NVARCHAR(MAX)',
 60    },
 61}
 62for _flavor, (_precision, _scale) in NUMERIC_PRECISION_FLAVORS.items():
 63    if _flavor not in DB_FLAVORS_CAST_DTYPES:
 64        DB_FLAVORS_CAST_DTYPES[_flavor] = {}
 65    DB_FLAVORS_CAST_DTYPES[_flavor].update({
 66        'NUMERIC': f"NUMERIC({_precision}, {_scale})",
 67        'DECIMAL': f"DECIMAL({_precision}, {_scale})",
 68    })
 69
 70DB_TO_PD_DTYPES: Dict[str, Union[str, Dict[str, str]]] = {
 71    'FLOAT': 'float64[pyarrow]',
 72    'REAL': 'float64[pyarrow]',
 73    'DOUBLE_PRECISION': 'float64[pyarrow]',
 74    'DOUBLE': 'float64[pyarrow]',
 75    'DECIMAL': 'numeric',
 76    'BIGINT': 'int64[pyarrow]',
 77    'INT': 'int32[pyarrow]',
 78    'INTEGER': 'int32[pyarrow]',
 79    'NUMBER': 'numeric',
 80    'NUMERIC': 'numeric',
 81    'GEOMETRY': 'geometry',
 82    'GEOMETRY(GEOMETRY)': 'geometry',
 83    'TIMESTAMP': 'datetime64[ns]',
 84    'TIMESTAMP WITHOUT TIMEZONE': 'datetime64[ns]',
 85    'TIMESTAMP WITH TIMEZONE': 'datetime64[ns, UTC]',
 86    'TIMESTAMP WITH TIME ZONE': 'datetime64[ns, UTC]',
 87    'TIMESTAMPTZ': 'datetime64[ns, UTC]',
 88    'DATE': 'datetime64[ns]',
 89    'DATETIME': 'datetime64[ns]',
 90    'DATETIME2': 'datetime64[ns]',
 91    'DATETIMEOFFSET': 'datetime64[ns, UTC]',
 92    'TEXT': 'string[pyarrow]',
 93    'VARCHAR': 'string[pyarrow]',
 94    'CLOB': 'string[pyarrow]',
 95    'BOOL': 'bool[pyarrow]',
 96    'BOOLEAN': 'bool[pyarrow]',
 97    'BOOLEAN()': 'bool[pyarrow]',
 98    'TINYINT': 'bool[pyarrow]',
 99    'TINYINT(1)': 'bool[pyarrow]',
100    'BIT': 'bool[pyarrow]',
101    'BIT(1)': 'bool[pyarrow]',
102    'JSON': 'json',
103    'JSONB': 'json',
104    'UUID': 'uuid',
105    'UNIQUEIDENTIFIER': 'uuid',
106    'BYTEA': 'bytes',
107    'BLOB': 'bytes',
108    'VARBINARY': 'bytes',
109    'VARBINARY(MAX)': 'bytes',
110    'substrings': {
111        'CHAR': 'string[pyarrow]',
112        'TIMESTAMP': 'datetime64[ns]',
113        'TIME': 'datetime64[ns]',
114        'DATE': 'datetime64[ns]',
115        'DOUBLE': 'double[pyarrow]',
116        'DECIMAL': 'numeric',
117        'NUMERIC': 'numeric',
118        'NUMBER': 'numeric',
119        'INT': 'int64[pyarrow]',
120        'BOOL': 'bool[pyarrow]',
121        'JSON': 'json',
122        'BYTE': 'bytes',
123        'LOB': 'bytes',
124        'BINARY': 'bytes',
125        'GEOMETRY': 'geometry',
126        'GEOGRAPHY': 'geography',
127    },
128    'default': 'object',
129}
130### Map pandas dtypes to flavor-specific dtypes.
131PD_TO_DB_DTYPES_FLAVORS: Dict[str, Dict[str, str]] = {
132    'int': {
133        'timescaledb': 'BIGINT',
134        'postgresql': 'BIGINT',
135        'postgis': 'BIGINT',
136        'mariadb': 'BIGINT',
137        'mysql': 'BIGINT',
138        'mssql': 'BIGINT',
139        'oracle': 'INT',
140        'sqlite': 'BIGINT',
141        'duckdb': 'BIGINT',
142        'citus': 'BIGINT',
143        'cockroachdb': 'BIGINT',
144        'default': 'INT',
145    },
146    'uint': {
147        'timescaledb': 'BIGINT',
148        'postgresql': 'BIGINT',
149        'postgis': 'BIGINT',
150        'mariadb': 'BIGINT',
151        'mysql': 'BIGINT',
152        'mssql': 'BIGINT',
153        'oracle': 'INT',
154        'sqlite': 'BIGINT',
155        'duckdb': 'BIGINT',
156        'citus': 'BIGINT',
157        'cockroachdb': 'BIGINT',
158        'default': 'INT',
159    },
160    'int8': {
161        'timescaledb': 'SMALLINT',
162        'postgresql': 'SMALLINT',
163        'postgis': 'SMALLINT',
164        'mariadb': 'SMALLINT',
165        'mysql': 'SMALLINT',
166        'mssql': 'SMALLINT',
167        'oracle': 'INT',
168        'sqlite': 'INT',
169        'duckdb': 'SMALLINT',
170        'citus': 'SMALLINT',
171        'cockroachdb': 'SMALLINT',
172        'default': 'INT',
173    },
174    'uint8': {
175        'timescaledb': 'SMALLINT',
176        'postgresql': 'SMALLINT',
177        'postgis': 'SMALLINT',
178        'mariadb': 'SMALLINT',
179        'mysql': 'SMALLINT',
180        'mssql': 'SMALLINT',
181        'oracle': 'INT',
182        'sqlite': 'INT',
183        'duckdb': 'SMALLINT',
184        'citus': 'SMALLINT',
185        'cockroachdb': 'SMALLINT',
186        'default': 'INT',
187    },
188    'int16': {
189        'timescaledb': 'SMALLINT',
190        'postgresql': 'SMALLINT',
191        'postgis': 'SMALLINT',
192        'mariadb': 'SMALLINT',
193        'mysql': 'SMALLINT',
194        'mssql': 'SMALLINT',
195        'oracle': 'INT',
196        'sqlite': 'INT',
197        'duckdb': 'SMALLINT',
198        'citus': 'SMALLINT',
199        'cockroachdb': 'SMALLINT',
200        'default': 'INT',
201    },
202    'int32': {
203        'timescaledb': 'INT',
204        'postgresql': 'INT',
205        'postgis': 'INT',
206        'mariadb': 'INT',
207        'mysql': 'INT',
208        'mssql': 'INT',
209        'oracle': 'INT',
210        'sqlite': 'INT',
211        'duckdb': 'INT',
212        'citus': 'INT',
213        'cockroachdb': 'INT',
214        'default': 'INT',
215    },
216    'int64': {
217        'timescaledb': 'BIGINT',
218        'postgresql': 'BIGINT',
219        'postgis': 'BIGINT',
220        'mariadb': 'BIGINT',
221        'mysql': 'BIGINT',
222        'mssql': 'BIGINT',
223        'oracle': 'INT',
224        'sqlite': 'BIGINT',
225        'duckdb': 'BIGINT',
226        'citus': 'BIGINT',
227        'cockroachdb': 'BIGINT',
228        'default': 'INT',
229    },
230    'float': {
231        'timescaledb': 'DOUBLE PRECISION',
232        'postgresql': 'DOUBLE PRECISION',
233        'postgis': 'DOUBLE PRECISION',
234        'mariadb': 'DOUBLE PRECISION',
235        'mysql': 'DOUBLE PRECISION',
236        'mssql': 'FLOAT',
237        'oracle': 'FLOAT',
238        'sqlite': 'FLOAT',
239        'duckdb': 'DOUBLE PRECISION',
240        'citus': 'DOUBLE PRECISION',
241        'cockroachdb': 'DOUBLE PRECISION',
242        'default': 'DOUBLE',
243    },
244    'double': {
245        'timescaledb': 'DOUBLE PRECISION',
246        'postgresql': 'DOUBLE PRECISION',
247        'postgis': 'DOUBLE PRECISION',
248        'mariadb': 'DOUBLE PRECISION',
249        'mysql': 'DOUBLE PRECISION',
250        'mssql': 'FLOAT',
251        'oracle': 'FLOAT',
252        'sqlite': 'FLOAT',
253        'duckdb': 'DOUBLE PRECISION',
254        'citus': 'DOUBLE PRECISION',
255        'cockroachdb': 'DOUBLE PRECISION',
256        'default': 'DOUBLE',
257    },
258    'datetime64[ns]': {
259        'timescaledb': 'TIMESTAMP',
260        'postgresql': 'TIMESTAMP',
261        'postgis': 'TIMESTAMP',
262        'mariadb': 'DATETIME',
263        'mysql': 'DATETIME',
264        'mssql': 'DATETIME2',
265        'oracle': 'TIMESTAMP(9)',
266        'sqlite': 'DATETIME',
267        'duckdb': 'TIMESTAMP',
268        'citus': 'TIMESTAMP',
269        'cockroachdb': 'TIMESTAMP',
270        'default': 'DATETIME',
271    },
272    'datetime64[ns, UTC]': {
273        'timescaledb': 'TIMESTAMPTZ',
274        'postgresql': 'TIMESTAMPTZ',
275        'postgis': 'TIMESTAMPTZ',
276        'mariadb': 'DATETIME',
277        'mysql': 'DATETIME',
278        'mssql': 'DATETIMEOFFSET',
279        'oracle': 'TIMESTAMP(9)',
280        'sqlite': 'TIMESTAMP',
281        'duckdb': 'TIMESTAMPTZ',
282        'citus': 'TIMESTAMPTZ',
283        'cockroachdb': 'TIMESTAMPTZ',
284        'default': 'TIMESTAMPTZ',
285    },
286    'datetime': {
287        'timescaledb': 'TIMESTAMPTZ',
288        'postgresql': 'TIMESTAMPTZ',
289        'postgis': 'TIMESTAMPTZ',
290        'mariadb': 'DATETIME',
291        'mysql': 'DATETIME',
292        'mssql': 'DATETIMEOFFSET',
293        'oracle': 'TIMESTAMP(9)',
294        'sqlite': 'TIMESTAMP',
295        'duckdb': 'TIMESTAMPTZ',
296        'citus': 'TIMESTAMPTZ',
297        'cockroachdb': 'TIMESTAMPTZ',
298        'default': 'TIMESTAMPTZ',
299    },
300    'datetimetz': {
301        'timescaledb': 'TIMESTAMPTZ',
302        'postgresql': 'TIMESTAMPTZ',
303        'postgis': 'TIMESTAMPTZ',
304        'mariadb': 'DATETIME',
305        'mysql': 'DATETIME',
306        'mssql': 'DATETIMEOFFSET',
307        'oracle': 'TIMESTAMP(9)',
308        'sqlite': 'TIMESTAMP',
309        'duckdb': 'TIMESTAMPTZ',
310        'citus': 'TIMESTAMPTZ',
311        'cockroachdb': 'TIMESTAMPTZ',
312        'default': 'TIMESTAMPTZ',
313    },
314    'bool': {
315        'timescaledb': 'BOOLEAN',
316        'postgresql': 'BOOLEAN',
317        'postgis': 'BOOLEAN',
318        'mariadb': 'BOOLEAN',
319        'mysql': 'BOOLEAN',
320        'mssql': 'BIT',
321        'oracle': 'INTEGER',
322        'sqlite': 'FLOAT',
323        'duckdb': 'BOOLEAN',
324        'citus': 'BOOLEAN',
325        'cockroachdb': 'BOOLEAN',
326        'default': 'BOOLEAN',
327    },
328    'object': {
329        'timescaledb': 'TEXT',
330        'postgresql': 'TEXT',
331        'postgis': 'TEXT',
332        'mariadb': 'TEXT',
333        'mysql': 'TEXT',
334        'mssql': 'NVARCHAR(MAX)',
335        'oracle': 'NVARCHAR2(2000)',
336        'sqlite': 'TEXT',
337        'duckdb': 'TEXT',
338        'citus': 'TEXT',
339        'cockroachdb': 'TEXT',
340        'default': 'TEXT',
341    },
342    'string': {
343        'timescaledb': 'TEXT',
344        'postgresql': 'TEXT',
345        'postgis': 'TEXT',
346        'mariadb': 'TEXT',
347        'mysql': 'TEXT',
348        'mssql': 'NVARCHAR(MAX)',
349        'oracle': 'NVARCHAR2(2000)',
350        'sqlite': 'TEXT',
351        'duckdb': 'TEXT',
352        'citus': 'TEXT',
353        'cockroachdb': 'TEXT',
354        'default': 'TEXT',
355    },
356    'unicode': {
357        'timescaledb': 'TEXT',
358        'postgresql': 'TEXT',
359        'postgis': 'TEXT',
360        'mariadb': 'TEXT',
361        'mysql': 'TEXT',
362        'mssql': 'NVARCHAR(MAX)',
363        'oracle': 'NVARCHAR2(2000)',
364        'sqlite': 'TEXT',
365        'duckdb': 'TEXT',
366        'citus': 'TEXT',
367        'cockroachdb': 'TEXT',
368        'default': 'TEXT',
369    },
370    'json': {
371        'timescaledb': 'JSONB',
372        'postgresql': 'JSONB',
373        'postgis': 'JSONB',
374        'mariadb': 'TEXT',
375        'mysql': 'TEXT',
376        'mssql': 'NVARCHAR(MAX)',
377        'oracle': 'NVARCHAR2(2000)',
378        'sqlite': 'TEXT',
379        'duckdb': 'TEXT',
380        'citus': 'JSONB',
381        'cockroachdb': 'JSONB',
382        'default': 'TEXT',
383    },
384    'numeric': {
385        'timescaledb': 'NUMERIC',
386        'postgresql': 'NUMERIC',
387        'postgis': 'NUMERIC',
388        'mariadb': f'DECIMAL{NUMERIC_PRECISION_FLAVORS["mariadb"]}',
389        'mysql': f'DECIMAL{NUMERIC_PRECISION_FLAVORS["mysql"]}',
390        'mssql': f'NUMERIC{NUMERIC_PRECISION_FLAVORS["mssql"]}',
391        'oracle': 'NUMBER',
392        'sqlite': 'TEXT',
393        'duckdb': 'TEXT',
394        'citus': 'NUMERIC',
395        'cockroachdb': 'NUMERIC',
396        'default': 'NUMERIC',
397    },
398    'uuid': {
399        'timescaledb': 'UUID',
400        'postgresql': 'UUID',
401        'postgis': 'UUID',
402        'mariadb': 'CHAR(36)',
403        'mysql': 'CHAR(36)',
404        'mssql': 'UNIQUEIDENTIFIER',
405        ### I know this is too much space, but erring on the side of caution.
406        'oracle': 'CHAR(36)',
407        'sqlite': 'TEXT',
408        'duckdb': 'VARCHAR',
409        'citus': 'UUID',
410        'cockroachdb': 'UUID',
411        'default': 'TEXT',
412    },
413    'bytes': {
414        'timescaledb': 'BYTEA',
415        'postgresql': 'BYTEA',
416        'postgis': 'BYTEA',
417        'mariadb': 'BLOB',
418        'mysql': 'BLOB',
419        'mssql': 'VARBINARY(MAX)',
420        'oracle': 'BLOB',
421        'sqlite': 'BLOB',
422        'duckdb': 'BLOB',
423        'citus': 'BYTEA',
424        'cockroachdb': 'BYTEA',
425        'default': 'BLOB',
426    },
427    'geometry': {
428        'timescaledb': 'TEXT',
429        'postgresql': 'TEXT',
430        'postgis': 'GEOMETRY',
431        'mariadb': 'TEXT',
432        'mysql': 'TEXT',
433        'mssql': 'NVARCHAR(MAX)',
434        'oracle': 'NVARCHAR2(2000)',
435        'sqlite': 'TEXT',
436        'duckdb': 'TEXT',
437        'citus': 'TEXT',
438        'cockroachdb': 'TEXT',
439        'default': 'TEXT',
440    },
441    'geography': {
442        'timescaledb': 'TEXT',
443        'postgresql': 'TEXT',
444        'postgis': 'GEOGRAPHY',
445        'mariadb': 'TEXT',
446        'mysql': 'TEXT',
447        'mssql': 'NVARCHAR(MAX)',
448        'oracle': 'NVARCHAR2(2000)',
449        'sqlite': 'TEXT',
450        'duckdb': 'TEXT',
451        'citus': 'TEXT',
452        'cockroachdb': 'TEXT',
453        'default': 'TEXT',
454    },
455}
456PD_TO_SQLALCHEMY_DTYPES_FLAVORS: Dict[str, Dict[str, str]] = {
457    'int': {
458        'timescaledb': 'BigInteger',
459        'postgresql': 'BigInteger',
460        'postgis': 'BigInteger',
461        'mariadb': 'BigInteger',
462        'mysql': 'BigInteger',
463        'mssql': 'BigInteger',
464        'oracle': 'BigInteger',
465        'sqlite': 'BigInteger',
466        'duckdb': 'BigInteger',
467        'citus': 'BigInteger',
468        'cockroachdb': 'BigInteger',
469        'default': 'BigInteger',
470    },
471    'uint': {
472        'timescaledb': 'BigInteger',
473        'postgresql': 'BigInteger',
474        'postgis': 'BigInteger',
475        'mariadb': 'BigInteger',
476        'mysql': 'BigInteger',
477        'mssql': 'BigInteger',
478        'oracle': 'BigInteger',
479        'sqlite': 'BigInteger',
480        'duckdb': 'BigInteger',
481        'citus': 'BigInteger',
482        'cockroachdb': 'BigInteger',
483        'default': 'BigInteger',
484    },
485    'int8': {
486        'timescaledb': 'SmallInteger',
487        'postgresql': 'SmallInteger',
488        'postgis': 'SmallInteger',
489        'mariadb': 'SmallInteger',
490        'mysql': 'SmallInteger',
491        'mssql': 'SmallInteger',
492        'oracle': 'SmallInteger',
493        'sqlite': 'SmallInteger',
494        'duckdb': 'SmallInteger',
495        'citus': 'SmallInteger',
496        'cockroachdb': 'SmallInteger',
497        'default': 'SmallInteger',
498    },
499    'uint8': {
500        'timescaledb': 'SmallInteger',
501        'postgresql': 'SmallInteger',
502        'postgis': 'SmallInteger',
503        'mariadb': 'SmallInteger',
504        'mysql': 'SmallInteger',
505        'mssql': 'SmallInteger',
506        'oracle': 'SmallInteger',
507        'sqlite': 'SmallInteger',
508        'duckdb': 'SmallInteger',
509        'citus': 'SmallInteger',
510        'cockroachdb': 'SmallInteger',
511        'default': 'SmallInteger',
512    },
513    'int16': {
514        'timescaledb': 'SmallInteger',
515        'postgresql': 'SmallInteger',
516        'postgis': 'SmallInteger',
517        'mariadb': 'SmallInteger',
518        'mysql': 'SmallInteger',
519        'mssql': 'SmallInteger',
520        'oracle': 'SmallInteger',
521        'sqlite': 'SmallInteger',
522        'duckdb': 'SmallInteger',
523        'citus': 'SmallInteger',
524        'cockroachdb': 'SmallInteger',
525        'default': 'SmallInteger',
526    },
527    'int32': {
528        'timescaledb': 'Integer',
529        'postgresql': 'Integer',
530        'postgis': 'Integer',
531        'mariadb': 'Integer',
532        'mysql': 'Integer',
533        'mssql': 'Integer',
534        'oracle': 'Integer',
535        'sqlite': 'Integer',
536        'duckdb': 'Integer',
537        'citus': 'Integer',
538        'cockroachdb': 'Integer',
539        'default': 'Integer',
540    },
541    'int64': {
542        'timescaledb': 'BigInteger',
543        'postgresql': 'BigInteger',
544        'postgis': 'BigInteger',
545        'mariadb': 'BigInteger',
546        'mysql': 'BigInteger',
547        'mssql': 'BigInteger',
548        'oracle': 'BigInteger',
549        'sqlite': 'BigInteger',
550        'duckdb': 'BigInteger',
551        'citus': 'BigInteger',
552        'cockroachdb': 'BigInteger',
553        'default': 'BigInteger',
554    },
555    'float': {
556        'timescaledb': 'Float',
557        'postgresql': 'Float',
558        'postgis': 'Float',
559        'mariadb': 'Float',
560        'mysql': 'Float',
561        'mssql': 'Float',
562        'oracle': 'Float',
563        'sqlite': 'Float',
564        'duckdb': 'Float',
565        'citus': 'Float',
566        'cockroachdb': 'Float',
567        'default': 'Float',
568    },
569    'datetime': {
570        'timescaledb': 'DateTime(timezone=True)',
571        'postgresql': 'DateTime(timezone=True)',
572        'postgis': 'DateTime(timezone=True)',
573        'mariadb': 'DateTime(timezone=True)',
574        'mysql': 'DateTime(timezone=True)',
575        'mssql': 'sqlalchemy.dialects.mssql.DATETIMEOFFSET',
576        'oracle': 'sqlalchemy.dialects.oracle.TIMESTAMP(timezone=True)',
577        'sqlite': 'DateTime(timezone=True)',
578        'duckdb': 'DateTime(timezone=True)',
579        'citus': 'DateTime(timezone=True)',
580        'cockroachdb': 'DateTime(timezone=True)',
581        'default': 'DateTime(timezone=True)',
582    },
583    'datetime64[ns]': {
584        'timescaledb': 'DateTime',
585        'postgresql': 'DateTime',
586        'postgis': 'DateTime',
587        'mariadb': 'DateTime',
588        'mysql': 'DateTime',
589        'mssql': 'sqlalchemy.dialects.mssql.DATETIME2',
590        'oracle': 'DateTime',
591        'sqlite': 'DateTime',
592        'duckdb': 'DateTime',
593        'citus': 'DateTime',
594        'cockroachdb': 'DateTime',
595        'default': 'DateTime',
596    },
597    'datetime64[ns, UTC]': {
598        'timescaledb': 'DateTime(timezone=True)',
599        'postgresql': 'DateTime(timezone=True)',
600        'postgis': 'DateTime(timezone=True)',
601        'mariadb': 'DateTime(timezone=True)',
602        'mysql': 'DateTime(timezone=True)',
603        'mssql': 'sqlalchemy.dialects.mssql.DATETIMEOFFSET',
604        'oracle': 'sqlalchemy.dialects.oracle.TIMESTAMP(timezone=True)',
605        'sqlite': 'DateTime(timezone=True)',
606        'duckdb': 'DateTime(timezone=True)',
607        'citus': 'DateTime(timezone=True)',
608        'cockroachdb': 'DateTime(timezone=True)',
609        'default': 'DateTime(timezone=True)',
610    },
611    'bool': {
612        'timescaledb': 'Boolean',
613        'postgresql': 'Boolean',
614        'postgis': 'Boolean',
615        'mariadb': 'Integer',
616        'mysql': 'Integer',
617        'mssql': 'sqlalchemy.dialects.mssql.BIT',
618        'oracle': 'Integer',
619        'sqlite': 'Float',
620        'duckdb': 'Boolean',
621        'citus': 'Boolean',
622        'cockroachdb': 'Boolean',
623        'default': 'Boolean',
624    },
625    'object': {
626        'timescaledb': 'UnicodeText',
627        'postgresql': 'UnicodeText',
628        'postgis': 'UnicodeText',
629        'mariadb': 'UnicodeText',
630        'mysql': 'UnicodeText',
631        'mssql': 'UnicodeText',
632        'oracle': 'UnicodeText',
633        'sqlite': 'UnicodeText',
634        'duckdb': 'UnicodeText',
635        'citus': 'UnicodeText',
636        'cockroachdb': 'UnicodeText',
637        'default': 'UnicodeText',
638    },
639    'string': {
640        'timescaledb': 'UnicodeText',
641        'postgresql': 'UnicodeText',
642        'postgis': 'UnicodeText',
643        'mariadb': 'UnicodeText',
644        'mysql': 'UnicodeText',
645        'mssql': 'UnicodeText',
646        'oracle': 'UnicodeText',
647        'sqlite': 'UnicodeText',
648        'duckdb': 'UnicodeText',
649        'citus': 'UnicodeText',
650        'cockroachdb': 'UnicodeText',
651        'default': 'UnicodeText',
652    },
653    'json': {
654        'timescaledb': 'sqlalchemy.dialects.postgresql.JSONB',
655        'postgresql': 'sqlalchemy.dialects.postgresql.JSONB',
656        'postgis': 'sqlalchemy.dialects.postgresql.JSONB',
657        'mariadb': 'UnicodeText',
658        'mysql': 'UnicodeText',
659        'mssql': 'UnicodeText',
660        'oracle': 'UnicodeText',
661        'sqlite': 'UnicodeText',
662        'duckdb': 'TEXT',
663        'citus': 'sqlalchemy.dialects.postgresql.JSONB',
664        'cockroachdb': 'sqlalchemy.dialects.postgresql.JSONB',
665        'default': 'UnicodeText',
666    },
667    'numeric': {
668        'timescaledb': 'Numeric',
669        'postgresql': 'Numeric',
670        'postgis': 'Numeric',
671        'mariadb': 'Numeric',
672        'mysql': 'Numeric',
673        'mssql': 'Numeric',
674        'oracle': 'Numeric',
675        'sqlite': 'UnicodeText',
676        'duckdb': 'Numeric',
677        'citus': 'Numeric',
678        'cockroachdb': 'Numeric',
679        'default': 'Numeric',
680    },
681    'uuid': {
682        'timescaledb': 'Uuid',
683        'postgresql': 'Uuid',
684        'postgis': 'Uuid',
685        'mariadb': 'sqlalchemy.dialects.mysql.CHAR(36)',
686        'mysql': 'sqlalchemy.dialects.mysql.CHAR(36)',
687        'mssql': 'Uuid',
688        'oracle': 'sqlalchemy.dialects.oracle.CHAR(36)',
689        'sqlite': 'UnicodeText',
690        'duckdb': 'UnicodeText',
691        'citus': 'Uuid',
692        'cockroachdb': 'Uuid',
693        'default': 'Uuid',
694    },
695    'bytes': {
696        'timescaledb': 'LargeBinary',
697        'postgresql': 'LargeBinary',
698        'postgis': 'LargeBinary',
699        'mariadb': 'LargeBinary',
700        'mysql': 'LargeBinary',
701        'mssql': 'LargeBinary',
702        'oracle': 'LargeBinary',
703        'sqlite': 'LargeBinary',
704        'duckdb': 'LargeBinary',
705        'citus': 'LargeBinary',
706        'cockroachdb': 'LargeBinary',
707        'default': 'LargeBinary',
708    },
709    'geometry': {
710        'timescaledb': 'UnicodeText',
711        'postgresql': 'UnicodeText',
712        'postgis': 'geoalchemy2.Geometry',
713        'mariadb': 'UnicodeText',
714        'mysql': 'UnicodeText',
715        'mssql': 'UnicodeText',
716        'oracle': 'UnicodeText',
717        'sqlite': 'UnicodeText',
718        'duckdb': 'UnicodeText',
719        'citus': 'UnicodeText',
720        'cockroachdb': 'UnicodeText',
721        'default': 'UnicodeText',
722    },
723    'geography': {
724        'timescaledb': 'UnicodeText',
725        'postgresql': 'UnicodeText',
726        'postgis': 'geoalchemy2.Geography',
727        'mariadb': 'UnicodeText',
728        'mysql': 'UnicodeText',
729        'mssql': 'UnicodeText',
730        'oracle': 'UnicodeText',
731        'sqlite': 'UnicodeText',
732        'duckdb': 'UnicodeText',
733        'citus': 'UnicodeText',
734        'cockroachdb': 'UnicodeText',
735        'default': 'UnicodeText',
736    },
737}
738
739AUTO_INCREMENT_COLUMN_FLAVORS: Dict[str, str] = {
740    'timescaledb': 'GENERATED BY DEFAULT AS IDENTITY',
741    'postgresql': 'GENERATED BY DEFAULT AS IDENTITY',
742    'postgis': 'GENERATED BY DEFAULT AS IDENTITY',
743    'mariadb': 'AUTO_INCREMENT',
744    'mysql': 'AUTO_INCREMENT',
745    'mssql': 'IDENTITY(1,1)',
746    'oracle': 'GENERATED BY DEFAULT ON NULL AS IDENTITY',
747    'sqlite': 'AUTOINCREMENT',
748    'duckdb': 'GENERATED BY DEFAULT',
749    'citus': 'GENERATED BY DEFAULT',
750    'cockroachdb': 'GENERATED BY DEFAULT AS IDENTITY',
751    'default': 'GENERATED BY DEFAULT AS IDENTITY',
752}
753
754
755def get_pd_type_from_db_type(db_type: str, allow_custom_dtypes: bool = True) -> str:
756    """
757    Parse a database type to a pandas data type.
758
759    Parameters
760    ----------
761    db_type: str
762        The database type, e.g. `DATETIME`, `BIGINT`, etc.
763
764    allow_custom_dtypes: bool, default False
765        If `True`, allow for custom data types like `json` and `str`.
766
767    Returns
768    -------
769    The equivalent datatype for a pandas DataFrame.
770    """
771    from meerschaum.utils.dtypes import are_dtypes_equal, get_geometry_type_srid
772    def parse_custom(_pd_type: str, _db_type: str) -> str:
773        if 'json' in _db_type.lower():
774            return 'json'
775        if are_dtypes_equal(_pd_type, 'numeric') and _pd_type != 'object':
776            precision, scale = get_numeric_precision_scale(None, dtype=_db_type.upper())
777            if precision and scale:
778                return f"numeric[{precision},{scale}]"
779        if are_dtypes_equal(_pd_type, 'geometry') and _pd_type != 'object':
780            geometry_type, srid = get_geometry_type_srid(_db_type.upper())
781            modifiers = [str(modifier) for modifier in (geometry_type, srid) if modifier]
782            typ = "geometry" if 'geography' not in _pd_type.lower() else 'geography'
783            if not modifiers:
784                return typ
785            return f"{typ}[{', '.join(modifiers)}]"
786        return _pd_type
787
788    pd_type = DB_TO_PD_DTYPES.get(db_type.upper().split('(', maxsplit=1)[0].strip(), None)
789    if pd_type is not None:
790        return (
791            parse_custom(pd_type, db_type)
792            if allow_custom_dtypes
793            else pd_type
794        )
795    for db_t, pd_t in DB_TO_PD_DTYPES['substrings'].items():
796        if db_t in db_type.upper():
797            return (
798                parse_custom(pd_t, db_t)
799                if allow_custom_dtypes
800                else pd_t
801            )
802    return DB_TO_PD_DTYPES['default']
803
804
805def get_db_type_from_pd_type(
806    pd_type: str,
807    flavor: str = 'default',
808    as_sqlalchemy: bool = False,
809) -> Union[str, 'sqlalchemy.sql.visitors.TraversibleType']:
810    """
811    Parse a Pandas data type into a flavor's database type.
812
813    Parameters
814    ----------
815    pd_type: str
816        The Pandas datatype. This must be a string, not the actual dtype object.
817
818    flavor: str, default 'default'
819        The flavor of the database to be mapped to.
820
821    as_sqlalchemy: bool, default False
822        If `True`, return a type from `sqlalchemy.types`.
823
824    Returns
825    -------
826    The database data type for the incoming Pandas data type.
827    If nothing can be found, a warning will be thrown and 'TEXT' will be returned.
828    """
829    from meerschaum.utils.warnings import warn
830    from meerschaum.utils.packages import attempt_import
831    from meerschaum.utils.dtypes import are_dtypes_equal, MRSM_ALIAS_DTYPES, get_geometry_type_srid
832    from meerschaum.utils.misc import parse_arguments_str
833    sqlalchemy_types = attempt_import('sqlalchemy.types', lazy=False)
834
835    types_registry = (
836        PD_TO_DB_DTYPES_FLAVORS
837        if not as_sqlalchemy
838        else PD_TO_SQLALCHEMY_DTYPES_FLAVORS
839    )
840
841    precision, scale = None, None
842    geometry_type, geometry_srid = None, None
843    og_pd_type = pd_type
844    if pd_type in MRSM_ALIAS_DTYPES:
845        pd_type = MRSM_ALIAS_DTYPES[pd_type]
846
847    ### Check whether we are able to match this type (e.g. pyarrow support).
848    found_db_type = False
849    if (
850        pd_type not in types_registry
851        and not any(
852            pd_type.startswith(f'{typ}[')
853            for typ in ('numeric', 'geometry', 'geography')
854        )
855    ):
856        for mapped_pd_type in types_registry:
857            if are_dtypes_equal(mapped_pd_type, pd_type):
858                pd_type = mapped_pd_type
859                found_db_type = True
860                break
861    elif (pd_type.startswith('geometry[') or pd_type.startswith('geography[')):
862        og_pd_type = pd_type
863        pd_type = 'geometry' if 'geometry' in pd_type else 'geography'
864        geometry_type, geometry_srid = get_geometry_type_srid(og_pd_type)
865        found_db_type = True
866    elif pd_type.startswith('numeric['):
867        og_pd_type = pd_type
868        pd_type = 'numeric'
869        precision, scale = get_numeric_precision_scale(flavor, og_pd_type)
870        found_db_type = True
871    else:
872        found_db_type = True
873
874    if not found_db_type:
875        warn(f"Unknown Pandas data type '{pd_type}'. Falling back to 'TEXT'.", stacklevel=3)
876        return (
877            'TEXT'
878            if not as_sqlalchemy
879            else sqlalchemy_types.UnicodeText
880        )
881    flavor_types = types_registry.get(
882        pd_type,
883        {
884            'default': (
885                'TEXT'
886                if not as_sqlalchemy
887                else 'UnicodeText'
888            ),
889        },
890    )
891    default_flavor_type = flavor_types.get(
892        'default',
893        (
894            'TEXT'
895            if not as_sqlalchemy
896            else 'UnicodeText'
897        ),
898    )
899    if flavor not in flavor_types:
900        warn(f"Unknown flavor '{flavor}'. Falling back to '{default_flavor_type}' (default).")
901    db_type = flavor_types.get(flavor, default_flavor_type)
902    if not as_sqlalchemy:
903        if precision is not None and scale is not None:
904            db_type_bare = db_type.split('(', maxsplit=1)[0]
905            return f"{db_type_bare}({precision},{scale})"
906        if geometry_type is not None and geometry_srid is not None:
907            if 'geometry' not in db_type.lower() and 'geography' not in db_type.lower():
908                return db_type
909            db_type_bare = db_type.split('(', maxsplit=1)[0]
910            return f"{db_type_bare}({geometry_type.upper()}, {geometry_srid})"
911        return db_type
912
913    if db_type.startswith('sqlalchemy.dialects'):
914        dialect, typ_class_name = db_type.replace('sqlalchemy.dialects.', '').split('.', maxsplit=2)
915        cls_args, cls_kwargs = None, None
916        if '(' in typ_class_name:
917            typ_class_name, args_str = typ_class_name.split('(', maxsplit=1)
918            args_str = args_str.rstrip(')')
919            cls_args, cls_kwargs = parse_arguments_str(args_str)
920        sqlalchemy_dialects_flavor_module = attempt_import(f'sqlalchemy.dialects.{dialect}')
921        cls = getattr(sqlalchemy_dialects_flavor_module, typ_class_name)
922        if cls_args is None:
923            return cls
924        return cls(*cls_args, **cls_kwargs)
925
926    if 'geometry' in db_type.lower() or 'geography' in db_type.lower():
927        geoalchemy2 = attempt_import('geoalchemy2', lazy=False)
928        geometry_class = (
929            geoalchemy2.Geometry
930            if 'geometry' in db_type.lower()
931            else geoalchemy2.Geography
932        )
933        if geometry_type is None or geometry_srid is None:
934            return geometry_class
935        return geometry_class(geometry_type=geometry_type, srid=geometry_srid)
936
937    if 'numeric' in db_type.lower():
938        if precision is None or scale is None:
939            return sqlalchemy_types.Numeric
940        return sqlalchemy_types.Numeric(precision, scale)
941
942    cls_args, cls_kwargs = None, None
943    typ_class_name = db_type
944    if '(' in db_type:
945        typ_class_name, args_str = db_type.split('(', maxsplit=1)
946        args_str = args_str.rstrip(')')
947        cls_args, cls_kwargs = parse_arguments_str(args_str)
948
949    cls = getattr(sqlalchemy_types, typ_class_name)
950    if cls_args is None:
951        return cls
952    return cls(*cls_args, **cls_kwargs)
953
954
955def get_numeric_precision_scale(
956    flavor: str,
957    dtype: Optional[str] = None,
958) -> Union[Tuple[int, int], Tuple[None, None]]:
959    """
960    Return the precision and scale to use for a numeric column for a given database flavor.
961
962    Parameters
963    ----------
964    flavor: str
965        The database flavor for which to return the precision and scale.
966    
967    dtype: Optional[str], default None
968        If provided, return the precision and scale provided in the dtype (if applicable).
969        If all caps, treat this as a DB type.
970
971    Returns
972    -------
973    A tuple of ints or a tuple of Nones.
974    """
975    if not dtype:
976        return None, None
977
978    lbracket = '[' if '[' in dtype else '('
979    rbracket = ']' if lbracket == '[' else ')'
980    if lbracket in dtype and dtype.count(',') == 1 and dtype.endswith(rbracket):
981        try:
982            parts = dtype.split(lbracket, maxsplit=1)[-1].rstrip(rbracket).split(',', maxsplit=1)
983            return int(parts[0].strip()), int(parts[1].strip())
984        except Exception:
985            pass
986
987    return NUMERIC_PRECISION_FLAVORS.get(flavor, (None, None))