def main():
for table in ("empresa", "socio", "cnae_secundaria"):
pgimport(
str(DATA_DIRECTORY / f"{table}.csv.gz"),
POSTGRES_URI,
table,
schema=load_schema(str(SCHEMA_DIRECTORY / f"{table}.csv")),
)
create_index(table)
with cnaes_csv() as (source, schema):
pgimport(source, POSTGRES_URI, "cnae", schema=load_schema(schema))
create_index("cnae", "codigo")
def from_file(cls, filename):
obj = cls()
obj.filename = filename
obj.fields = load_schema(
str(filename)
) # TODO: load_schema must support Path objects
return obj
def read_cases(input_filename, order_by=None):
cases = rows.import_from_csv(input_filename,
force_types=load_schema(
str(SCHEMA_PATH / "caso.csv")))
if order_by:
cases.order_by(order_by)
return cases
def get_cities():
table = rows.import_from_csv(
POPULATION_DATA_PATH, force_types=load_schema(str(POPULATION_SCHEMA_PATH)),
)
cities = defaultdict(dict)
for row in table:
cities[row.state][row.city] = row
return cities
def main():
if wait_for_postgres() and tables_exist():
print(("There are existing tables in the database. "
"Please, start with a clean database."))
return
tables = ("empresa", "socio", "cnae_secundaria")
files = ("empresa", "socio", "cnae-secundaria")
for table, filename in zip(tables, files):
pgimport(
str(DATA_DIRECTORY / f"{filename}.csv.gz"),
POSTGRES_URI,
table,
schema=load_schema(str(SCHEMA_DIRECTORY / f"{filename}.csv")),
)
create_index(table)
with cnaes_csv() as (source, schema):
pgimport(source, POSTGRES_URI, "cnae", schema=load_schema(schema))
create_index("cnae", "codigo")
def _get_schemas_for_inputs(schemas, inputs):
if schemas is None:
schemas = [None for _ in inputs]
else:
schemas = [schema.strip() or None for schema in schemas.split(",")]
if len(schemas) > len(inputs):
click.echo("ERROR: number of schemas is greater than sources", err=True)
sys.exit(9)
elif len(schemas) < len(inputs):
diff = len(inputs) - len(schemas)
for _ in range(diff):
schemas.append(None)
return [load_schema(schema) if schema else None for schema in schemas]
def read_schema(filename):
"""Read schema and with custom fields"""
fields = load_schema(
str(filename),
context={
"cpf": CPFField,
"text": rows.fields.TextField,
"decimal": CustomDecimalField,
"date": rows.fields.DateField,
"integer": rows.fields.IntegerField,
},
)
# Add "optional" parameter
for row in rows.import_from_csv(filename):
fields[row.field_name] = fields[row.field_name]()
fields[row.field_name].optional = "optional" in (row.options or "")
return fields
def schema(self):
return load_schema(str(self.schema_filename))
def read_population():
return rows.import_from_csv(
DATA_PATH / "populacao-estimada-2019.csv",
force_types=load_schema(
str(SCHEMA_PATH / "populacao-estimada-2019.csv")),
)
