def save_file(self, name, file_path, file_name=None, company_id=None):
""" Save the file to our store
Args:
name (str): with that name file will be available in sql api
file_name (str): file name
file_path (str): path to the file
company_id (int): company id
Returns:
int: id of 'file' record in db
"""
if file_name is None:
file_name = Path(file_path).name
try:
ds_meta_dir = Path(self.dir).joinpath(f'{company_id}@@@@@{name}')
ds_meta_dir.mkdir()
source = ds_meta_dir.joinpath(file_name)
shutil.move(file_path, str(source))
ds = FileDS(str(source))
ds_meta = self._get_ds_meta(ds)
column_names = ds_meta['column_names']
if ds_meta['column_names'] is not None:
column_names = json.dumps(
[dict(name=x) for x in ds_meta['column_names']])
file_record = File(name=name,
company_id=company_id,
source_file_path=file_name,
file_path=str(source),
row_count=ds_meta['row_count'],
columns=column_names)
session.add(file_record)
session.commit()
self.fs_store.put(f'{company_id}@@@@@{name}',
f'file_{company_id}_{file_record.id}', self.dir)
except Exception as e:
log.error(e)
shutil.rmtree(ds_meta_dir)
raise
return file_record.id
def create_datasource(self,
source_type,
source,
file_path=None,
company_id=None,
ds_meta_dir=None):
datasource_controller = DatasourceController()
if source_type == 'file':
source = os.path.join(ds_meta_dir, source)
shutil.move(file_path, source)
ds = FileDS(source)
creation_info = {'class': 'FileDS', 'args': [source], 'kwargs': {}}
elif datasource_controller.get_db_integration(source_type,
company_id) is not None:
integration = datasource_controller.get_db_integration(
source_type, company_id)
ds_class_map = {
'clickhouse': ClickhouseDS,
'mariadb': MariaDS,
'mysql': MySqlDS,
'singlestore': MySqlDS,
'postgres': PostgresDS,
'cockroachdb': PostgresDS,
'mssql': MSSQLDS,
'mongodb': MongoDS,
'snowflake': SnowflakeDS,
'athena': AthenaDS,
'cassandra': CassandraDS,
'scylladb': ScyllaDS,
'trinodb': TrinoDS
}
try:
dsClass = ds_class_map[integration['type']]
except KeyError:
raise KeyError(
f"Unknown DS type: {source_type}, type is {integration['type']}"
)
if dsClass is None:
raise Exception(
f"Unsupported datasource: {source_type}, type is {integration['type']}, please install required dependencies!"
)
if integration['type'] in ['clickhouse']:
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] in [
'mssql', 'postgres', 'cockroachdb', 'mariadb', 'mysql',
'singlestore', 'cassandra', 'scylladb'
]:
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port']
}
}
kwargs = creation_info['kwargs']
integration_folder_name = f'integration_files_{company_id}_{integration["id"]}'
if integration['type'] in ('mysql', 'mariadb'):
kwargs['ssl'] = integration.get('ssl')
kwargs['ssl_ca'] = integration.get('ssl_ca')
kwargs['ssl_cert'] = integration.get('ssl_cert')
kwargs['ssl_key'] = integration.get('ssl_key')
for key in ['ssl_ca', 'ssl_cert', 'ssl_key']:
if isinstance(kwargs[key],
str) and len(kwargs[key]) > 0:
kwargs[key] = os.path.join(
self.integrations_dir, integration_folder_name,
kwargs[key])
elif integration['type'] in ('cassandra', 'scylla'):
kwargs['secure_connect_bundle'] = integration.get(
'secure_connect_bundle')
if (isinstance(kwargs['secure_connect_bundle'], str)
and len(kwargs['secure_connect_bundle']) > 0):
kwargs['secure_connect_bundle'] = os.path.join(
self.integrations_dir, integration_folder_name,
kwargs['secure_connect_bundle'])
if 'database' in integration:
kwargs['database'] = integration['database']
if 'database' in source:
kwargs['database'] = source['database']
ds = dsClass(**kwargs)
elif integration['type'] == 'snowflake':
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query':
source['query'].replace('"', "'"),
'schema':
source.get('schema', integration['schema']),
'warehouse':
source.get('warehouse', integration['warehouse']),
'database':
source.get('database', integration['database']),
'host':
integration['host'],
'password':
integration['password'],
'user':
integration['user'],
'account':
integration['account']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] == 'mongodb':
if isinstance(source['find'], str):
source['find'] = json.loads(source['find'])
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'database': source['database'],
'collection': source['collection'],
'query': source['find'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] == 'athena':
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'staging_dir': source['staging_dir'],
'database': source['database'],
'access_key': source['access_key'],
'secret_key': source['secret_key'],
'region_name': source['region_name']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] == 'trinodb':
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port'],
'schema': integration['schema'],
'catalog': integration['catalog']
}
}
ds = dsClass(**creation_info['kwargs'])
else:
# This probably only happens for urls
ds = FileDS(source)
creation_info = {'class': 'FileDS', 'args': [source], 'kwargs': {}}
return ds, creation_info
def save_datasource(self, name, source_type, source, file_path=None):
if source_type == 'file' and (file_path is None):
raise Exception(
'`file_path` argument required when source_type == "file"')
datasource_record = session.query(Datasource).filter_by(
company_id=self.company_id, name=name).first()
while datasource_record is not None:
raise Exception(f'Datasource with name {name} already exists')
try:
datasource_record = Datasource(
company_id=self.company_id,
name=name,
datasources_version=mindsdb_datasources.__version__,
mindsdb_version=mindsdb_version)
session.add(datasource_record)
session.commit()
datasource_record = session.query(Datasource).filter_by(
company_id=self.company_id, name=name).first()
ds_meta_dir = os.path.join(self.dir, name)
os.mkdir(ds_meta_dir)
if source_type == 'file':
source = os.path.join(ds_meta_dir, source)
shutil.move(file_path, source)
ds = FileDS(source)
creation_info = {
'class': 'FileDS',
'args': [source],
'kwargs': {}
}
elif source_type in self.config['integrations']:
integration = self.config['integrations'][source_type]
ds_class_map = {
'clickhouse': ClickhouseDS,
'mariadb': MariaDS,
'mysql': MySqlDS,
'postgres': PostgresDS,
'mssql': MSSQLDS,
'mongodb': MongoDS,
'snowflake': SnowflakeDS,
'athena': AthenaDS
}
try:
dsClass = ds_class_map[integration['type']]
except KeyError:
raise KeyError(
f"Unknown DS type: {source_type}, type is {integration['type']}"
)
if dsClass is None:
raise Exception(
f'Unsupported datasource: {source_type}, please install required dependencies!'
)
if integration['type'] in ['clickhouse']:
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] in [
'mssql', 'postgres', 'mariadb', 'mysql'
]:
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port']
}
}
if 'database' in integration:
creation_info['kwargs']['database'] = integration[
'database']
if 'database' in source:
creation_info['kwargs']['database'] = source[
'database']
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] == 'snowflake':
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'schema': source['schema'],
'warehouse': source['warehouse'],
'database': source['database'],
'host': integration['host'],
'password': integration['password'],
'user': integration['user'],
'account': integration['account']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] == 'mongodb':
if isinstance(source['find'], str):
source['find'] = json.loads(source['find'])
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'database': source['database'],
'collection': source['collection'],
'query': source['find'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] == 'athena':
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'staging_dir': source['staging_dir'],
'database': source['database'],
'access_key': source['access_key'],
'secret_key': source['secret_key'],
'region_name': source['region_name']
}
}
ds = dsClass(**creation_info['kwargs'])
else:
# This probably only happens for urls
ds = FileDS(source)
creation_info = {
'class': 'FileDS',
'args': [source],
'kwargs': {}
}
df = ds.df
if '' in df.columns or len(df.columns) != len(set(df.columns)):
shutil.rmtree(ds_meta_dir)
raise Exception(
'Each column in datasource must have unique non-empty name'
)
datasource_record.creation_info = json.dumps(creation_info)
datasource_record.data = json.dumps({
'source_type':
source_type,
'source':
source,
'row_count':
len(df),
'columns': [dict(name=x) for x in list(df.keys())]
})
self.fs_store.put(
name, f'datasource_{self.company_id}_{datasource_record.id}',
self.dir)
session.commit()
except Exception as e:
log.error(f'{e}')
try:
self.delete_datasource(name)
except Exception:
pass
raise e
return self.get_datasource_obj(name, raw=True), name
def create_datasource(self,
source_type,
source,
file_path=None,
company_id=None):
integration_controller = IntegrationController()
if source_type == 'view_query':
dsClass = QueryDS
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'source': source['source'], # view
'source_type': source_type,
'company_id': company_id
}
}
ds = dsClass(**creation_info['kwargs'])
elif source_type == 'file':
file_name = source.get('mindsdb_file_name')
file_record = session.query(File).filter_by(
company_id=company_id, name=file_name).first()
if file_record is None:
raise Exception(f"Cant find file '{file_name}'")
self.fs_store.get(f'{company_id}@@@@@{file_name}',
f'file_{company_id}_{file_record.id}', self.dir)
kwargs = {}
query = source.get('query')
if query is not None:
kwargs['query'] = query
path = Path(
self.dir).joinpath(f'{company_id}@@@@@{file_name}').joinpath(
file_record.source_file_path)
creation_info = {
'class': 'FileDS',
'args': [str(path)],
'kwargs': kwargs
}
ds = FileDS(str(path), **kwargs)
elif integration_controller.get(source_type, company_id) is not None:
integration = integration_controller.get(source_type, company_id)
ds_class_map = {
'clickhouse': ClickhouseDS,
'mariadb': MariaDS,
'mysql': MySqlDS,
'singlestore': MySqlDS,
'postgres': PostgresDS,
'cockroachdb': PostgresDS,
'mssql': MSSQLDS,
'mongodb': MongoDS,
'snowflake': SnowflakeDS,
'athena': AthenaDS,
'cassandra': CassandraDS,
'scylladb': ScyllaDS,
'trinodb': TrinoDS
}
try:
dsClass = ds_class_map[integration['type']]
except KeyError:
raise KeyError(
f"Unknown DS type: {source_type}, type is {integration['type']}"
)
if dsClass is None:
raise Exception(
f"Unsupported dataset: {source_type}, type is {integration['type']}, please install required dependencies!"
)
if integration['type'] in ['clickhouse']:
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] in [
'mssql', 'postgres', 'cockroachdb', 'mariadb', 'mysql',
'singlestore', 'cassandra', 'scylladb'
]:
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port']
}
}
kwargs = creation_info['kwargs']
integration_folder_name = f'integration_files_{company_id}_{integration["id"]}'
if integration['type'] in ('mysql', 'mariadb'):
kwargs['ssl'] = integration.get('ssl')
kwargs['ssl_ca'] = integration.get('ssl_ca')
kwargs['ssl_cert'] = integration.get('ssl_cert')
kwargs['ssl_key'] = integration.get('ssl_key')
for key in ['ssl_ca', 'ssl_cert', 'ssl_key']:
if isinstance(kwargs[key],
str) and len(kwargs[key]) > 0:
kwargs[key] = os.path.join(
self.integrations_dir, integration_folder_name,
kwargs[key])
elif integration['type'] in ('cassandra', 'scylla'):
kwargs['secure_connect_bundle'] = integration.get(
'secure_connect_bundle')
if (isinstance(kwargs['secure_connect_bundle'], str)
and len(kwargs['secure_connect_bundle']) > 0):
kwargs['secure_connect_bundle'] = os.path.join(
self.integrations_dir, integration_folder_name,
kwargs['secure_connect_bundle'])
if 'database' in integration:
kwargs['database'] = integration['database']
if 'database' in source:
kwargs['database'] = source['database']
ds = dsClass(**kwargs)
elif integration['type'] == 'snowflake':
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query':
source['query'],
'schema':
source.get('schema', integration['schema']),
'warehouse':
source.get('warehouse', integration['warehouse']),
'database':
source.get('database', integration['database']),
'host':
integration['host'],
'password':
integration['password'],
'user':
integration['user'],
'account':
integration['account']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] == 'mongodb':
if isinstance(source['find'], str):
source['find'] = json.loads(source['find'])
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'database': source['database'],
'collection': source['collection'],
'query': source['find'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] == 'athena':
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'staging_dir': source['staging_dir'],
'database': source['database'],
'access_key': source['access_key'],
'secret_key': source['secret_key'],
'region_name': source['region_name']
}
}
ds = dsClass(**creation_info['kwargs'])
elif integration['type'] == 'trinodb':
creation_info = {
'class': dsClass.__name__,
'args': [],
'kwargs': {
'query': source['query'],
'user': integration['user'],
'password': integration['password'],
'host': integration['host'],
'port': integration['port'],
'schema': integration['schema'],
'catalog': integration['catalog']
}
}
ds = dsClass(**creation_info['kwargs'])
else:
# This probably only happens for urls
ds = FileDS(source)
creation_info = {'class': 'FileDS', 'args': [source], 'kwargs': {}}
return ds, creation_info
def test_file_ds_with_issued_csv(self):
ds = FileDS(self.issued_csv)
try:
ds._handle_source()
except Exception as e:
assert False, "issue csv were not handled properly: {}".format(e)