def append_data(self, descriptor, table_name=None): table_name = table_name or self.carto_table_name client = CopySQLClient(self.carto_auth_client) query = "COPY {table_name} ({columns}) FROM stdin WITH (FORMAT csv, HEADER true)".format( table_name=table_name, columns=",".join(self.carto_field_names)) client.copyfrom_file_object(query, descriptor)
def __init__(self, credentials): self.credentials = credentials or get_default_credentials() check_credentials(self.credentials) self.auth_client = _create_auth_client(self.credentials) self.sql_client = SQLClient(self.auth_client) self.copy_client = CopySQLClient(self.auth_client) self.batch_sql_client = BatchSQLClient(self.auth_client)
def copy_from(self, data, filepath, to_table): if self._copy_client is None: from carto.sql import CopySQLClient self._copy_client = CopySQLClient(self._auth_client) headers = data.readline().decode('utf-8') data.seek(0) from_query = 'COPY %s (%s) FROM stdin WITH (FORMAT csv, HEADER true)' % ( to_table, headers) return self._copy_client.copyfrom_file_object(from_query, data)
def copy(tablename, rows, delimiter=",", quote='"', headers=None): copy_client = CopySQLClient(_get_auth_client()) rows = iter(rows) if headers is None: headers = delimiter.join(next(rows)) else: headers = delimiter.join(headers) from_query = f"""COPY {tablename} ({headers}) FROM stdin (FORMAT CSV, DELIMITER '{delimiter}', HEADER false)""" try: return copy_client.copyfrom(from_query, rows_generator(rows, delimiter, quote)) except CartoException as e: log.error(f"Error importing \n\n {rows}") log.exception(e)
def initialize(self): if not self.api_url and self.user_name: self.api_url = "https://{}.carto.com/api/".format(self.user_name) elif not self.api_url and not self.user_name: raise Exception( 'Not enough data provided to initialize the client') if self.org_name: self.client = APIKeyAuthClient(self.api_url, self.api_key, self.org_name) else: self.client = APIKeyAuthClient(self.api_url, self.api_key) self.sql_client = SQLClient(self.client) self.batch_client = BatchSQLClient(self.client) self.copy_client = CopySQLClient(self.client)
logger.error('You need to provide valid credentials, run with ' '-h parameter for details') sys.exit(1) # Create and cartodbfy a table sqlClient = SQLClient(auth_client) sqlClient.send(""" CREATE TABLE IF NOT EXISTS copy_example ( the_geom geometry(Geometry,4326), name text, age integer ) """) sqlClient.send("SELECT CDB_CartodbfyTable(current_schema, 'copy_example')") copyClient = CopySQLClient(auth_client) # COPY FROM example logger.info("COPY'ing FROM file...") query = ('COPY copy_example (the_geom, name, age) ' 'FROM stdin WITH (FORMAT csv, HEADER true)') result = copyClient.copyfrom_file_path(query, 'files/copy_from.csv') logger.info('result = %s' % result) # COPY TO example query = 'COPY copy_example TO stdout WITH (FORMAT csv, HEADER true)' output_file = 'files/copy_export.csv' copyClient.copyto_file_path(query, output_file) logger.info('Table copied to %s' % output_file) # Truncate the table to make this example repeatable
class ContextManager: def __init__(self, credentials): self.credentials = credentials or get_default_credentials() check_credentials(self.credentials) self.auth_client = _create_auth_client(self.credentials) self.sql_client = SQLClient(self.auth_client) self.copy_client = CopySQLClient(self.auth_client) self.batch_sql_client = BatchSQLClient(self.auth_client) def execute_query(self, query, parse_json=True, do_post=True, format=None, **request_args): return self.sql_client.send(query.strip(), parse_json, do_post, format, **request_args) def execute_long_running_query(self, query): return self.batch_sql_client.create_and_wait_for_completion( query.strip()) def copy_to(self, source, schema, limit=None, retry_times=DEFAULT_RETRY_TIMES): query = self.compute_query(source, schema) columns = self._get_query_columns_info(query) copy_query = self._get_copy_query(query, columns, limit) return self._copy_to(copy_query, columns, retry_times) def copy_from(self, gdf, table_name, if_exists='fail', cartodbfy=True): schema = self.get_schema() table_name = self.normalize_table_name(table_name) columns = get_dataframe_columns_info(gdf) if if_exists == 'replace' or not self.has_table(table_name, schema): log.debug('Creating table "{}"'.format(table_name)) self._create_table_from_columns(table_name, columns, schema, cartodbfy) elif if_exists == 'fail': raise Exception( 'Table "{schema}.{table_name}" already exists in your CARTO account. ' 'Please choose a different `table_name` or use ' 'if_exists="replace" to overwrite it.'.format( table_name=table_name, schema=schema)) else: # 'append' pass self._copy_from(gdf, table_name, columns) return table_name def create_table_from_query(self, query, table_name, if_exists, cartodbfy=True): schema = self.get_schema() table_name = self.normalize_table_name(table_name) if if_exists == 'replace' or not self.has_table(table_name, schema): log.debug('Creating table "{}"'.format(table_name)) self._create_table_from_query(query, table_name, schema, cartodbfy) elif if_exists == 'fail': raise Exception( 'Table "{schema}.{table_name}" already exists in your CARTO account. ' 'Please choose a different `table_name` or use ' 'if_exists="replace" to overwrite it.'.format( table_name=table_name, schema=schema)) else: # 'append' pass return table_name def has_table(self, table_name, schema=None): query = self.compute_query(table_name, schema) return self._check_exists(query) def delete_table(self, table_name): query = _drop_table_query(table_name) output = self.execute_query(query) return not ('notices' in output and 'does not exist' in output['notices'][0]) def rename_table(self, table_name, new_table_name, if_exists='fail'): new_table_name = self.normalize_table_name(new_table_name) if table_name == new_table_name: raise ValueError( 'Table names are equal. Please choose a different table name.') if not self.has_table(table_name): raise Exception( 'Table "{table_name}" does not exist in your CARTO account.'. format(table_name=table_name)) if self.has_table(new_table_name): if if_exists == 'replace': log.debug('Removing table "{}"'.format(new_table_name)) self.delete_table(new_table_name) elif if_exists == 'fail': raise Exception( 'Table "{new_table_name}" already exists in your CARTO account. ' 'Please choose a different `new_table_name` or use ' 'if_exists="replace" to overwrite it.'.format( new_table_name=new_table_name)) self._rename_table(table_name, new_table_name) return new_table_name def update_privacy_table(self, table_name, privacy=None): DatasetInfo(self.auth_client, table_name).update_privacy(privacy) def get_privacy(self, table_name): return DatasetInfo(self.auth_client, table_name).privacy def get_schema(self): """Get user schema from current credentials""" query = 'SELECT current_schema()' result = self.execute_query(query, do_post=False) return result['rows'][0]['current_schema'] def get_geom_type(self, query): """Fetch geom type of a remote table or query""" distict_query = ''' SELECT distinct ST_GeometryType(the_geom) AS geom_type FROM ({}) q LIMIT 5 '''.format(query) response = self.execute_query(distict_query, do_post=False) if response and response.get('rows') and len(response.get('rows')) > 0: st_geom_type = response.get('rows')[0].get('geom_type') if st_geom_type: return map_geom_type(st_geom_type[3:]) return None def get_num_rows(self, query): """Get the number of rows in the query""" result = self.execute_query( "SELECT COUNT(*) FROM ({query}) _query".format(query=query)) return result.get('rows')[0].get('count') def get_bounds(self, query): extent_query = ''' SELECT ARRAY[ ARRAY[st_xmin(geom_env), st_ymin(geom_env)], ARRAY[st_xmax(geom_env), st_ymax(geom_env)] ] bounds FROM ( SELECT ST_Extent(the_geom) geom_env FROM ({}) q ) q; '''.format(query) response = self.execute_query(extent_query, do_post=False) if response and response.get('rows') and len(response.get('rows')) > 0: return response.get('rows')[0].get('bounds') return None def get_column_names(self, source, schema=None, exclude=None): query = self.compute_query(source, schema) columns = [c.name for c in self._get_query_columns_info(query)] if exclude and isinstance(exclude, list): columns = list(set(columns) - set(exclude)) return columns def is_public(self, query): # Used to detect public tables in queries in the publication, # because privacy only works for tables. public_auth_client = _create_auth_client(self.credentials, public=True) public_sql_client = SQLClient(public_auth_client) exists_query = 'EXPLAIN {}'.format(query) try: public_sql_client.send(exists_query, do_post=False) return True except CartoException: return False def get_table_names(self, query): # Used to detect tables in queries in the publication. query = 'SELECT CDB_QueryTablesText(\'{}\') as tables'.format(query) result = self.execute_query(query) tables = [] if result['total_rows'] > 0 and result['rows'][0]['tables']: # Dataset_info only works with tables without schema tables = [ table.split('.')[1] if '.' in table else table for table in result['rows'][0]['tables'] ] return tables def _create_table_from_query(self, query, table_name, schema, cartodbfy=True): query = 'BEGIN; {drop}; {create}; {cartodbfy}; COMMIT;'.format( drop=_drop_table_query(table_name), create=_create_table_from_query_query(table_name, query), cartodbfy=_cartodbfy_query(table_name, schema) if cartodbfy else '') self.execute_long_running_query(query) def _create_table_from_columns(self, table_name, columns, schema, cartodbfy=True): query = 'BEGIN; {drop}; {create}; {cartodbfy}; COMMIT;'.format( drop=_drop_table_query(table_name), create=_create_table_from_columns_query(table_name, columns), cartodbfy=_cartodbfy_query(table_name, schema) if cartodbfy else '') self.execute_long_running_query(query) def compute_query(self, source, schema=None): if is_sql_query(source): return source schema = schema or self.get_schema() return self._compute_query_from_table(source, schema) def _compute_query_from_table(self, table_name, schema): return 'SELECT * FROM "{schema}"."{table_name}"'.format( schema=schema or 'public', table_name=table_name) def _check_exists(self, query): exists_query = 'EXPLAIN {}'.format(query) try: self.execute_query(exists_query, do_post=False) return True except CartoException: return False def _get_query_columns_info(self, query): query = 'SELECT * FROM ({}) _q LIMIT 0'.format(query) table_info = self.execute_query(query) return Column.from_sql_api_fields(table_info['fields']) def _get_copy_query(self, query, columns, limit): query_columns = [ column.name for column in columns if (column.name != 'the_geom_webmercator') ] query = 'SELECT {columns} FROM ({query}) _q'.format( query=query, columns=','.join(query_columns)) if limit is not None: if isinstance(limit, int) and (limit >= 0): query += ' LIMIT {limit}'.format(limit=limit) else: raise ValueError("`limit` parameter must an integer >= 0") return query def _copy_to(self, query, columns, retry_times): copy_query = 'COPY ({0}) TO stdout WITH (FORMAT csv, HEADER true, NULL \'{1}\')'.format( query, PG_NULL) try: raw_result = self.copy_client.copyto_stream(copy_query) except CartoRateLimitException as err: if retry_times > 0: retry_times -= 1 warn('Read call rate limited. Waiting {s} seconds'.format( s=err.retry_after)) time.sleep(err.retry_after) warn('Retrying...') return self._copy_to(query, columns, retry_times) else: warn(( 'Read call was rate-limited. ' 'This usually happens when there are multiple queries being read at the same time.' )) raise err converters = obtain_converters(columns) parse_dates = date_columns_names(columns) df = read_csv(raw_result, converters=converters, parse_dates=parse_dates) return df def _copy_from(self, dataframe, table_name, columns): query = """ COPY {table_name}({columns}) FROM stdin WITH (FORMAT csv, DELIMITER '|', NULL '{null}'); """.format(table_name=table_name, null=PG_NULL, columns=','.join(column.dbname for column in columns)).strip() data = _compute_copy_data(dataframe, columns) self.copy_client.copyfrom(query, data) def _rename_table(self, table_name, new_table_name): query = _rename_table_query(table_name, new_table_name) self.execute_query(query) def normalize_table_name(self, table_name): norm_table_name = normalize_name(table_name) if norm_table_name != table_name: log.debug('Table name normalized: "{}"'.format(norm_table_name)) return norm_table_name
class CartoDataSource(DataSource): SUBDOMAIN_URL_PATTERN = "https://%s.carto.com" ON_PREMISES_URL_PATTERN = "https://%s/user/%s" DEFAULT_API_VERSION = 'v2' def __init__(self, user, api_key, options={}): super().__init__(options) self.do_post = options.get('do_post', False) self.parse_json = options.get('parse_json', True) self.format = options.get('format', 'json') self.base_url_option = options.get('base_url', '') self.api_version = options.get('api_version', self.DEFAULT_API_VERSION) self.batch = options.get('batch', False) self.user = user self.api_key = api_key self.base_url = self._generate_base_url(user, self.base_url_option) # Carto Context for DataFrame handling self._carto_context = None # Carto client for COPYs self._copy_client = None self._auth_client = APIKeyAuthClient(api_key=api_key, base_url=self.base_url) self._sql_client = SQLClient(self._auth_client, api_version=self.api_version) self._batch_client = None if self.batch: self._batch_client = BatchSQLClient(self._auth_client) @property def cc(self): """ Creates and returns a CartoContext object to work with Panda Dataframes :return: """ # TODO: The CartoContext documentaton says that SSL must be disabled sometimes if an on # premise host is used. # We are not taking this into account. It would need to create a requests.Session() # object, set its SSL to false and pass it to the CartoContext init. if self._carto_context is None: self._carto_context = cartoframes.CartoContext( base_url=self.base_url, api_key=self.api_key) return self._carto_context def _generate_base_url(self, user, base_url_option): if base_url_option: base_url = self.ON_PREMISES_URL_PATTERN % (base_url_option, user) else: base_url = self.SUBDOMAIN_URL_PATTERN % user return base_url def execute_query(self, query_template, params, query_config, **opts): # TODO: Here we are parsing the parameters and taking responsability for it. We do not make # any safe parsing as this will be used in a backend-to-backend context and we build our # own queries. # --- # This is also problematic as quoting is not done and relies in the query template # --- # Can we use the .mogrify method in psycopg2 to render a query as it is going to be # executed ? -> NO # -> .mogrify is a cursor method but in CARTO connections we lack a cursor. # --- # There is an open issue in CARTO about having separated parameters and binding them in # the server: # https://github.com/CartoDB/Geographica-Product-Coordination/issues/57 params = {k: "'" + v + "'" for k, v in params.items()} formatted_query = query_template % params try: return self._sql_client.send(formatted_query, parse_json=self.parse_json, do_post=self.do_post, format=self.format) except CartoException as e: raise LongitudeQueryCannotBeExecutedException(str(e)) def parse_response(self, response): return LongitudeQueryResponse(rows=response['rows'], fields=response['fields'], meta={ 'response_time': response.get('time'), 'total_rows': response.get('total_rows') }) def copy_from(self, data, filepath, to_table): if self._copy_client is None: from carto.sql import CopySQLClient self._copy_client = CopySQLClient(self._auth_client) headers = data.readline().decode('utf-8') data.seek(0) from_query = 'COPY %s (%s) FROM stdin WITH (FORMAT csv, HEADER true)' % ( to_table, headers) return self._copy_client.copyfrom_file_object(from_query, data) def read_dataframe(self, table_name='', *args, **kwargs): return self.cc.read(table_name=table_name, *args, **kwargs) def query_dataframe(self, query='', *args, **kwargs): return self.cc.query(query=query, *args, **kwargs) def write_dataframe(self, df, table_name='', *args, **kwargs): return self.cc.write(df=df, table_name=table_name, *args, **kwargs)
' https://username.carto.com/' ' (defaults to env variable CARTO_API_URL)')) parser.add_argument('--api_key', dest='CARTO_API_KEY', default=os.environ.get('CARTO_API_KEY', ''), help=('Api key of the account' ' (defaults to env variable CARTO_API_KEY)')) args = parser.parse_args() if not args.CARTO_BASE_URL or not args.CARTO_API_KEY: sys.exit(parser.print_usage()) auth_client = APIKeyAuthClient(args.CARTO_BASE_URL, args.CARTO_API_KEY) sql_client = SQLClient(auth_client) copy_client = CopySQLClient(auth_client) # Create a table suitable to receive the data logger.info('Creating table nexrad_copy_example...') sql_client.send("""CREATE TABLE IF NOT EXISTS nexrad_copy_example ( the_geom geometry(Geometry,4326), reflectivity numeric )""") sql_client.send( "SELECT CDB_CartodbfyTable(current_schema, 'nexrad_copy_example')") logger.info('Done') logger.info('Trying to connect to the THREDDS radar query service') rs = RadarServer( 'http://thredds.ucar.edu/thredds/radarServer/nexrad/level2/IDD/')
def copy_client(api_key_auth_client_usr): return CopySQLClient(api_key_auth_client_usr)
logger.error('You need to provide valid credentials, run with ' '-h parameter for details') sys.exit(1) # Create and cartodbfy a table sqlClient = SQLClient(auth_client) sqlClient.send(""" CREATE TABLE IF NOT EXISTS copy_example ( the_geom geometry(Geometry,4326), name text, age integer ) """) sqlClient.send("SELECT CDB_CartodbfyTable(current_schema, 'copy_example')") copyClient = CopySQLClient(auth_client) # COPY FROM example logger.info("COPY'ing FROM file...") query = ('COPY copy_example (the_geom, name, age) ' 'FROM stdin WITH (FORMAT csv, HEADER true)') result = copyClient.copyfrom_file_path(query, 'files/copy_from.csv') logger.info('result = %s' % result) # COPY TO example with pandas DataFrame logger.info("COPY'ing TO pandas DataFrame...") query = 'COPY copy_example TO stdout WITH (FORMAT csv, HEADER true)' result = copyClient.copyto_stream(query) df = pd.read_csv(result) logger.info(df.head())
# This is a bit of a trick: we omit the sequences to avoid # dependencies on other objects Normally this just affects the # cartodb_id and can optionally be fixed by cartodbfy'ing create_table_no_seqs = re.sub(r'DEFAULT nextval\([^\)]+\)', ' ', create_table) logger.info(create_table_no_seqs) # Create the table in the destination account logger.info('Creating the table in the destination account...') res = sql_dst_client.send(create_table_no_seqs) logger.info('Response: {}'.format(res)) # Cartodbfy the table (this is optional) logger.info("Cartodbfy'ing the destination table...") res = sql_dst_client.send( "SELECT CDB_CartodbfyTable(current_schema, '%s')" % TABLE_NAME ) logger.info('Response: {}'.format(res)) # Create COPY clients copy_src_client = CopySQLClient(auth_src_client) copy_dst_client = CopySQLClient(auth_dst_client) # COPY (streaming) the data from the source to the dest table. We use # here all the COPY defaults. Note that we take the `response` from # the `copyto`, which can be iterated, and we pipe it directly into # the `copyfrom`. logger.info("Streaming the data from source to destination...") response = copy_src_client.copyto('COPY %s TO STDOUT' % TABLE_NAME) result = copy_dst_client.copyfrom('COPY %s FROM STDIN' % TABLE_NAME, response) logger.info('Result: {}'.format(result))
class CARTOUser(object): def __init__(self, user_name=None, org_name=None, api_url=None, api_key=None, check_ssl=True): self.user_name = user_name self.org_name = org_name self.api_url = api_url self.api_key = api_key if not check_ssl: old_request = requests.Session.request requests.Session.request = partialmethod(old_request, verify=False) warnings.filterwarnings('ignore', 'Unverified HTTPS request') def initialize(self): if not self.api_url and self.user_name: self.api_url = "https://{}.carto.com/api/".format(self.user_name) elif not self.api_url and not self.user_name: raise Exception( 'Not enough data provided to initialize the client') if self.org_name: self.client = APIKeyAuthClient(self.api_url, self.api_key, self.org_name) else: self.client = APIKeyAuthClient(self.api_url, self.api_key) self.sql_client = SQLClient(self.client) self.batch_client = BatchSQLClient(self.client) self.copy_client = CopySQLClient(self.client) def execute_sql(self, query, parse_json=True, format=None, do_post=False): try: try: self.client except AttributeError: self.initialize() return self.sql_client.send(query, parse_json=parse_json, format=format, do_post=do_post) except CartoException as e: raise Exception(e.args[0].args[0][0]) def batch_check(self, job_id): try: self.batch_client except AttributeError: self.initialize() return self.batch_client.read(job_id) def batch_create(self, query): try: self.batch_client except AttributeError: self.initialize() return self.batch_client.create(query) def batch_cancel(self, job_id): try: self.batch_client except AttributeError: self.initialize() return self.batch_client.cancel(job_id) def get_dataset_manager(self): try: self.sql_client except AttributeError: self.initialize() return DatasetManager(self.client) def get_sync_manager(self): try: self.sql_client except AttributeError: self.initialize() return SyncTableJobManager(self.client) def upload(self, uri, sync_time=None): try: self.sql_client except AttributeError: self.initialize() dataset_manager = DatasetManager(self.client) if sync_time: return dataset_manager.create(uri, sync_time) else: return dataset_manager.create(uri) def copy_from(self, path, query, tablename=None, delimiter=','): try: self.copy_client except AttributeError: self.initialize() if tablename is None: tablename = Path(path).stem if query is None: with open(path, 'rb') as myfile: headers = next(myfile).strip().decode('utf8') query = f"""COPY {tablename} ({headers}) FROM stdin (FORMAT CSV, DELIMITER '{delimiter}', HEADER false, QUOTE '"')""" return self.copy_client.copyfrom_file_object(query, myfile) return self.copy_client.copyfrom_file_path(query, path) def copy_to(self, query, output, delimiter=','): try: self.copy_client except AttributeError: self.initialize() copy_query = f"""COPY ({query}) TO stdout WITH (FORMAT CSV, DELIMITER '{delimiter}', HEADER true, QUOTE '"')""" return self.copy_client.copyto_file_path(copy_query, output)
# This is a bit of a trick: we omit the sequences to avoid # dependencies on other objects Normally this just affects the # cartodb_id and can optionally be fixed by cartodbfy'ing create_table_no_seqs = re.sub(r'DEFAULT nextval\([^\)]+\)', ' ', create_table) logger.info(create_table_no_seqs) # Create the table in the destination account logger.info('Creating the table in the destination account...') res = sql_dst_client.send(create_table_no_seqs) logger.info('Response: {}'.format(res)) # Cartodbfy the table (this is optional) logger.info("Cartodbfy'ing the destination table...") res = sql_dst_client.send("SELECT CDB_CartodbfyTable(current_schema, '%s')" % TABLE_NAME) logger.info('Response: {}'.format(res)) # Create COPY clients copy_src_client = CopySQLClient(auth_src_client) copy_dst_client = CopySQLClient(auth_dst_client) # COPY (streaming) the data from the source to the dest table. We use # here all the COPY defaults. Note that we take the `response` from # the `copyto`, which can be iterated, and we pipe it directly into # the `copyfrom`. logger.info("Streaming the data from source to destination...") response = copy_src_client.copyto('COPY %s TO STDOUT' % TABLE_NAME) result = copy_dst_client.copyfrom('COPY %s FROM STDIN' % TABLE_NAME, response) logger.info('Result: {}'.format(result))
logger.error('You need to provide valid credentials, run with ' '-h parameter for details') sys.exit(1) # Create and cartodbfy a table sqlClient = SQLClient(auth_client) sqlClient.send(""" CREATE TABLE IF NOT EXISTS copy_example ( the_geom geometry(Geometry,4326), name text, age integer ) """) sqlClient.send("SELECT CDB_CartodbfyTable(current_schema, 'copy_example')") copyClient = CopySQLClient(auth_client) # COPY FROM example logger.info("COPY'ing FROM file...") query = ('COPY copy_example (the_geom, name, age) ' 'FROM stdin WITH (FORMAT csv, HEADER true)') result = copyClient.copyfrom_file_path(query, 'files/copy_from.csv') logger.info('result = %s' % result) # COPY TO example query = 'COPY copy_example TO stdout WITH (FORMAT csv, HEADER true)' output_file = 'files/copy_export.csv' copyClient.copyto_file_path(query, output_file) logger.info('Table copied to %s' % output_file) # Truncate the table to make this example repeatable
logger.error('You need to provide valid credentials, run with ' '-h parameter for details') sys.exit(1) # Create and cartodbfy a table sqlClient = SQLClient(auth_client) sqlClient.send(""" CREATE TABLE IF NOT EXISTS copy_example ( the_geom geometry(Geometry,4326), name text, age integer ) """) sqlClient.send("SELECT CDB_CartodbfyTable(current_schema, 'copy_example')") copyClient = CopySQLClient(auth_client) # COPY FROM example logger.info("COPY'ing FROM file...") query = ('COPY copy_example (the_geom, name, age) ' 'FROM stdin WITH (FORMAT csv, HEADER true)') result = copyClient.copyfrom_file_path(query, 'files/copy_from.csv') logger.info('result = %s' % result) # COPY TO example with pandas DataFrame logger.info("COPY'ing TO pandas DataFrame...") query = 'COPY copy_example TO stdout WITH (FORMAT csv, HEADER true)' result = copyClient.copyto_stream(query) df = pd.read_csv(result) logger.info(df.head())
class ContextManager: def __init__(self, credentials): self.credentials = credentials or get_default_credentials() check_credentials(self.credentials) self.auth_client = _create_auth_client(self.credentials) self.sql_client = SQLClient(self.auth_client) self.copy_client = CopySQLClient(self.auth_client) self.batch_sql_client = BatchSQLClient(self.auth_client) @not_found def execute_query(self, query, parse_json=True, do_post=True, format=None, **request_args): return self.sql_client.send(query.strip(), parse_json, do_post, format, **request_args) @not_found def execute_long_running_query(self, query): return self.batch_sql_client.create_and_wait_for_completion( query.strip()) def copy_to(self, source, schema=None, limit=None, retry_times=DEFAULT_RETRY_TIMES): query = self.compute_query(source, schema) columns = self._get_query_columns_info(query) copy_query = self._get_copy_query(query, columns, limit) return self._copy_to(copy_query, columns, retry_times) def copy_from(self, gdf, table_name, if_exists='fail', cartodbfy=True, retry_times=DEFAULT_RETRY_TIMES): schema = self.get_schema() table_name = self.normalize_table_name(table_name) df_columns = get_dataframe_columns_info(gdf) if self.has_table(table_name, schema): if if_exists == 'replace': table_query = self._compute_query_from_table( table_name, schema) table_columns = self._get_query_columns_info(table_query) if self._compare_columns(df_columns, table_columns): # Equal columns: truncate table self._truncate_table(table_name, schema, cartodbfy) else: # Diff columns: truncate table and drop + add columns self._truncate_and_drop_add_columns( table_name, schema, df_columns, table_columns, cartodbfy) elif if_exists == 'fail': raise Exception( 'Table "{schema}.{table_name}" already exists in your CARTO account. ' 'Please choose a different `table_name` or use ' 'if_exists="replace" to overwrite it.'.format( table_name=table_name, schema=schema)) else: # 'append' pass else: self._create_table_from_columns(table_name, schema, df_columns, cartodbfy) self._copy_from(gdf, table_name, df_columns, retry_times) return table_name def create_table_from_query(self, query, table_name, if_exists, cartodbfy=True): schema = self.get_schema() table_name = self.normalize_table_name(table_name) if self.has_table(table_name, schema): if if_exists == 'replace': # TODO: review logic copy_from self._drop_create_table_from_query(table_name, schema, query, cartodbfy) elif if_exists == 'fail': raise Exception( 'Table "{schema}.{table_name}" already exists in your CARTO account. ' 'Please choose a different `table_name` or use ' 'if_exists="replace" to overwrite it.'.format( table_name=table_name, schema=schema)) else: # 'append' pass else: self._drop_create_table_from_query(table_name, schema, query, cartodbfy) return table_name def list_tables(self, schema=None): datasets = DatasetManager(self.auth_client).filter( show_table_size_and_row_count='false', show_table='false', show_stats='false', show_likes='false', show_liked='false', show_permission='false', show_uses_builder_features='false', show_synchronization='false', load_totals='false') datasets.sort(key=lambda x: x.updated_at, reverse=True) return pd.DataFrame([dataset.name for dataset in datasets], columns=['tables']) def has_table(self, table_name, schema=None): query = self.compute_query(table_name, schema) return self._check_exists(query) def delete_table(self, table_name): query = _drop_table_query(table_name) output = self.execute_query(query) return not ('notices' in output and 'does not exist' in output['notices'][0]) def rename_table(self, table_name, new_table_name, if_exists='fail'): new_table_name = self.normalize_table_name(new_table_name) if table_name == new_table_name: raise ValueError( 'Table names are equal. Please choose a different table name.') if not self.has_table(table_name): raise Exception( 'Table "{table_name}" does not exist in your CARTO account.'. format(table_name=table_name)) if self.has_table(new_table_name): if if_exists == 'replace': log.debug('Removing table "{}"'.format(new_table_name)) self.delete_table(new_table_name) elif if_exists == 'fail': raise Exception( 'Table "{new_table_name}" already exists in your CARTO account. ' 'Please choose a different `new_table_name` or use ' 'if_exists="replace" to overwrite it.'.format( new_table_name=new_table_name)) self._rename_table(table_name, new_table_name) return new_table_name def update_privacy_table(self, table_name, privacy=None): DatasetInfo(self.auth_client, table_name).update_privacy(privacy) def get_privacy(self, table_name): return DatasetInfo(self.auth_client, table_name).privacy def get_schema(self): """Get user schema from current credentials""" query = 'SELECT current_schema()' result = self.execute_query(query, do_post=False) schema = result['rows'][0]['current_schema'] log.debug('schema: {}'.format(schema)) return schema def get_geom_type(self, query): """Fetch geom type of a remote table or query""" distict_query = ''' SELECT distinct ST_GeometryType(the_geom) AS geom_type FROM ({}) q LIMIT 5 '''.format(query) response = self.execute_query(distict_query, do_post=False) if response and response.get('rows') and len(response.get('rows')) > 0: st_geom_type = response.get('rows')[0].get('geom_type') if st_geom_type: return map_geom_type(st_geom_type[3:]) return None def get_num_rows(self, query): """Get the number of rows in the query""" result = self.execute_query( 'SELECT COUNT(*) FROM ({query}) _query'.format(query=query)) return result.get('rows')[0].get('count') def get_bounds(self, query): extent_query = ''' SELECT ARRAY[ ARRAY[st_xmin(geom_env), st_ymin(geom_env)], ARRAY[st_xmax(geom_env), st_ymax(geom_env)] ] bounds FROM ( SELECT ST_Extent(the_geom) geom_env FROM ({}) q ) q; '''.format(query) response = self.execute_query(extent_query, do_post=False) if response and response.get('rows') and len(response.get('rows')) > 0: return response.get('rows')[0].get('bounds') return None def get_column_names(self, source, schema=None, exclude=None): query = self.compute_query(source, schema) columns = [c.name for c in self._get_query_columns_info(query)] if exclude and isinstance(exclude, list): columns = list(set(columns) - set(exclude)) return columns def is_public(self, query): # Used to detect public tables in queries in the publication, # because privacy only works for tables. public_auth_client = _create_auth_client(self.credentials, public=True) public_sql_client = SQLClient(public_auth_client) exists_query = 'EXPLAIN {}'.format(query) try: public_sql_client.send(exists_query, do_post=False) return True except CartoException: return False def get_table_names(self, query): # Used to detect tables in queries in the publication. query = 'SELECT CDB_QueryTablesText($q${}$q$) as tables'.format(query) result = self.execute_query(query) tables = [] if result['total_rows'] > 0 and result['rows'][0]['tables']: # Dataset_info only works with tables without schema tables = [ table.split('.')[1] if '.' in table else table for table in result['rows'][0]['tables'] ] return tables def _compare_columns(self, a, b): a_copy = [i for i in a if _not_reserved(i.name)] b_copy = [i for i in b if _not_reserved(i.name)] a_copy.sort() b_copy.sort() return a_copy == b_copy def _drop_create_table_from_query(self, table_name, schema, query, cartodbfy): log.debug('DROP + CREATE table "{}"'.format(table_name)) query = 'BEGIN; {drop}; {create}; {cartodbfy}; COMMIT;'.format( drop=_drop_table_query(table_name), create=_create_table_from_query_query(table_name, query), cartodbfy=_cartodbfy_query(table_name, schema) if cartodbfy else '') self.execute_long_running_query(query) def _create_table_from_columns(self, table_name, schema, columns, cartodbfy): log.debug('CREATE table "{}"'.format(table_name)) query = 'BEGIN; {create}; {cartodbfy}; COMMIT;'.format( create=_create_table_from_columns_query(table_name, columns), cartodbfy=_cartodbfy_query(table_name, schema) if cartodbfy else '') self.execute_long_running_query(query) def _truncate_table(self, table_name, schema, cartodbfy): log.debug('TRUNCATE table "{}"'.format(table_name)) query = 'BEGIN; {truncate}; {cartodbfy}; COMMIT;'.format( truncate=_truncate_table_query(table_name), cartodbfy=_cartodbfy_query(table_name, schema) if cartodbfy else '') self.execute_long_running_query(query) def _truncate_and_drop_add_columns(self, table_name, schema, df_columns, table_columns, cartodbfy): log.debug( 'TRUNCATE AND DROP + ADD columns table "{}"'.format(table_name)) query = '{regenerate}; BEGIN; {truncate}; {drop_columns}; {add_columns}; {cartodbfy}; COMMIT;'.format( regenerate=_regenerate_table_query(table_name, schema) if self._check_regenerate_table_exists() else '', truncate=_truncate_table_query(table_name), drop_columns=_drop_columns_query(table_name, table_columns), add_columns=_add_columns_query(table_name, df_columns), cartodbfy=_cartodbfy_query(table_name, schema) if cartodbfy else '') self.execute_long_running_query(query) def compute_query(self, source, schema=None): if is_sql_query(source): return source schema = schema or self.get_schema() return self._compute_query_from_table(source, schema) def _compute_query_from_table(self, table_name, schema): return 'SELECT * FROM "{schema}"."{table_name}"'.format( schema=schema or 'public', table_name=table_name) def _check_exists(self, query): exists_query = 'EXPLAIN {}'.format(query) try: self.execute_query(exists_query, do_post=False) return True except CartoException: return False def _check_regenerate_table_exists(self): query = ''' SELECT 1 FROM pg_catalog.pg_proc p LEFT JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace WHERE p.proname = 'cdb_regeneratetable' AND n.nspname = 'cartodb'; ''' result = self.execute_query(query) return len(result['rows']) > 0 def _get_query_columns_info(self, query): query = 'SELECT * FROM ({}) _q LIMIT 0'.format(query) table_info = self.execute_query(query) return get_query_columns_info(table_info['fields']) def _get_copy_query(self, query, columns, limit): query_columns = [ double_quote(column.name) for column in columns if (column.name != 'the_geom_webmercator') ] query = 'SELECT {columns} FROM ({query}) _q'.format( query=query, columns=','.join(query_columns)) if limit is not None: if isinstance(limit, int) and (limit >= 0): query += ' LIMIT {limit}'.format(limit=limit) else: raise ValueError("`limit` parameter must an integer >= 0") return query @retry_copy def _copy_to(self, query, columns, retry_times=DEFAULT_RETRY_TIMES): log.debug('COPY TO') copy_query = "COPY ({0}) TO stdout WITH (FORMAT csv, HEADER true, NULL '{1}')".format( query, PG_NULL) raw_result = self.copy_client.copyto_stream(copy_query) converters = obtain_converters(columns) parse_dates = date_columns_names(columns) df = pd.read_csv(raw_result, converters=converters, parse_dates=parse_dates) return df @retry_copy def _copy_from(self, dataframe, table_name, columns, retry_times=DEFAULT_RETRY_TIMES): log.debug('COPY FROM') query = """ COPY {table_name}({columns}) FROM stdin WITH (FORMAT csv, DELIMITER '|', NULL '{null}'); """.format(table_name=table_name, null=PG_NULL, columns=','.join( double_quote(column.dbname) for column in columns)).strip() data = _compute_copy_data(dataframe, columns) self.copy_client.copyfrom(query, data) def _rename_table(self, table_name, new_table_name): query = _rename_table_query(table_name, new_table_name) self.execute_query(query) def normalize_table_name(self, table_name): norm_table_name = normalize_name(table_name) if norm_table_name != table_name: log.debug('Table name normalized: "{}"'.format(norm_table_name)) return norm_table_name